NeMo SpeechLM2 训练与扩展实战从 SLURM 集群调度到 FSDP2/EP/CP 多维度并行【免费下载链接】SpeechA scalable generative AI framework built for researchers and developers working on Large Language Models, Multimodal, and Speech AI (Automatic Speech Recognition and Text-to-Speech)项目地址: https://gitcode.com/GitHub_Trending/nem/Speech本文是 NeMo当前仓库 nem/Speech中 SpeechLM2 集合的训练与规模扩展指南围绕 training_and_scaling.rst 展开覆盖 SLURM 集群上的作业编排、本地 torchrun 调试以及 FSDP2、张量并行TP、序列并行SP、上下文并行CP、专家并行EP等并行策略的完整配置与实现原理。读完本文你将能够独立提交 SpeechLM2 训练作业、诊断分布式训练问题并为不同规模的模型和 GPU 集群选择合适的并行方案。概述SpeechLM2 训练扩展的三个关键环节SpeechLM2 集合代码位于 nemo/collections/speechlm2/面向语音大语言模型Speech LLM的训练其规模扩展能力建立在三个环节之上作业编排通过 SLURM 提交脚本与auto_launcher_with_seed.sh批量启动多次实验并统一管理随机种子并行策略以 Lightning 的ModelParallelStrategy为基座新增了面向 NeMo Automodel 的AutomodelParallelStrategy支持 FSDP2、TP、PP、CP、EPMoE与 HSDP 等并行维度数据与内核兼容为长音频训练提供 packed-sequenceTHD路径并内置了对已知不兼容组合如 BSHD CP的运行时校验避免静默出错。下面按“先跑通、再扩展、后深挖原理”的顺序展开。运行实验SLURM 集群上的作业提交SpeechLM2 提供了面向 SLURM 集群的作业启动脚本。文档给出的标准工作流是# 连续提交 8 个带随机种子的作业 scripts/speechlm2/auto_launcher_with_seed.sh -n8 s2s_tinyllama_repro.subauto_launcher_with_seed.sh脚本完成三件事为每个提交的作业生成独立的随机种子借助 Lhotse 数据加载器中的shard_seedrandomized设置保证每个数据并行DPrank 的随机分片互不相同保证每个张量并行TPrank 使用相同的种子从而确保并行切分的一致性。说明该脚本在文档中作为训练编排入口出现具体实现可依据你自己的集群环境在scripts/speechlm2/目录下准备当前仓库该目录包含 oomptimizer.py、distributed_oomptimizer.py、estimate_token_bins.py 等配套工具。训练与推理的实际入口脚本位于 examples/speechlm2/例如 s2s_duplex_train.py 与 salm_train.py。SLURM 提交脚本示例下面是一个标准的 SLURM 提交脚本s2s_tinyllama_repro.sub来自文档它申请 4 个节点、每节点 8 个 GPU#!/bin/bash #SBATCH --job-names2s_training #SBATCH --nodes4 #SBATCH --ntasks-per-node8 #SBATCH --gresgpu:8 #SBATCH --time24:00:00 #SBATCH --exclusive #SBATCH --outputs2s_tinyllama_repro_%j.out # Check that the global random seed base is provided if [ -z $1 ]; then echo Usage: $0 global_random_seed_base exit 1 fi SEED${1} EXP_NAMEs2s_training RESULTS_DIRresults/${EXP_NAME} srun --ntasks${SLURM_NTASKS} --ntasks-per-node${SLURM_NTASKS_PER_NODE} \ python -u examples/speechlm2/s2s_duplex_train.py \ --config-path/path/to/config/dir \ --config-names2s_training.yaml \ exp_manager.name${EXP_NAME} \ exp_manager.wandb_logger_kwargs.name${EXP_NAME} \ trainer.num_nodes$SLURM_JOB_NUM_NODES \ exp_manager.explicit_log_dir${RESULTS_DIR} \ data.train_ds.seed$SEED \ data.validation_ds.seed$SEED几个值得注意的细节脚本接收一个global_random_seed_base位置参数即$1用于设置data.train_ds.seed与data.validation_ds.seedsrun直接继承 SLURM 环境变量SLURM_NTASKS、SLURM_NTASKS_PER_NODE、SLURM_JOB_NUM_NODES无需重复指定节点数训练参数通过 Hydra 命令行覆盖注入exp_manager.name决定实验名exp_manager.explicit_log_dir指定日志目录WB logger 名称与实验名保持一致。配置文件的使用原则所有模型、训练和数据参数集中放在主配置文件如s2s_training.yaml中。文档明确建议复制并修改配置文件而不是在 SLURM 脚本里逐个覆盖参数这样有利于版本管理与配置清晰。配置文件的结构说明可参考 configs.rst。以 salm_automodel.yaml 为例其顶层结构分为model模型与骨干网络、trainer训练器与并行策略、data训练/验证数据集、exp_manager实验管理四大部分。调试本地 torchrun 快速验证在提交大规模集群作业之前建议先用torchrun在本地做小规模调试与性能分析# Run with 4 GPUs locally torchrun --nproc_per_node4 examples/speechlm2/s2s_duplex_train.py \ --config-path/path/to/config/dir \ --config-names2s_training.yamltorchrun负责为每个进程设置LOCAL_RANK环境变量并启动进程组。salm_train.py 中对应逻辑是训练入口在torch.cuda.is_available()时按os.environ[LOCAL_RANK]设置当前设备并调用torch.distributed.init_process_group(backendnccl)初始化进程组然后通过seed_everything(cfg.data.train_ds.seed)固定全局种子、以Trainer(**resolve_trainer_cfg(cfg.trainer))构造 Lightning Trainer。若在配置文件里设置use_nemo_automodel: true脚本会从nemo.collections.speechlm2导入SALMAutomodel作为模型类走 NeMo Automodel 后端。扩展策略SpeechLM2 支持的并行维度SpeechLM2 集合原生支持多种模型并行策略Fully Sharded Data ParallelFSDP2将模型参数分片到多张 GPU前向/反向时先 all-gather、计算后立即释放适合在有限显存下训练更大的模型Tensor ParallelismTP把单个张量如注意力头、嵌入维度拆分到多张 GPU每个 rank 只持有并计算一部分Sequence ParallelismSP把序列维度拆分到多张 GPU 上处理2D ParallelismFSDP2 与 TP/SP 的组合同时获得参数分片与张量/序列切分的能力。从源码看这些策略的统一入口是模型类如SALMAutomodel、SALM的configure_model()方法它读取trainer.strategy配置并完成设备网格创建与模型并行化。AutomodelParallelStrategySALMAutomodel 专用对于SALMAutomodel集合提供了专门的AutomodelParallelStrategy实现在 nemo/collections/speechlm2/parts/parallel.py。它把拓扑解析与并行化委托给 NeMo Automodel 的公开DistributedSetup构建器是ModelParallelStrategy的直接替代品。该策略可以构建 FSDP2、TP、PP、CP、EPMoE与 HSDP 的设备网格但实际支持情况取决于所选 LLM 架构、后端以及并行轴的组合。以下配置片段来自 salm_automodel.yamltrainer: strategy: _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy dp_size: null # 从 world_size / 其他维度自动推断 dp_replicate_size: 1 # HSDP 复制组大小 tp_size: 1 pp_size: 1 cp_size: 1 ep_size: 8 # MoE 模型的专家并行度 # 激活检查点——两个相互独立的开关 activation_checkpointing_llm: false # LLM transformer 块 activation_checkpointing_perception: false # 语音编码器层各维度含义与 parallel.py 中的AutomodelParallelStrategy.__init__一一对应参数默认值含义dp_sizeNone数据并行大小None时由world_size与其他并行维度推断dp_replicate_size1HSDP 复制组大小大于 1 时启用混合分片Hybrid Shardingtp_size1张量并行大小pp_size1流水线并行大小cp_size1上下文并行大小ep_size1专家并行大小MoE 模型activation_checkpointing_llmfalseLLM 侧激活检查点开关activation_checkpointing_perceptionfalse感知编码器侧激活检查点开关策略解析流程见AutomodelParallelStrategy.create_device_mesh是用FSDP2Config()可被distributed_config覆盖作为分布式配置结合ParallelismSizes与world_size调用DistributedSetup.build(...)最终得到形如(pp, dp_replicate, dp_shard, cp, tp)的设备网格及其展平子网格dp、dp_shard_cp、dp_cp。解析结果会通过self._distributed_setup暴露给模型在setup_environment()中_device_mesh、_moe_mesh、_distributed_setup会被写入 LightningModule供configure_model()读取。模型侧的configure_model()拿到解析后的distributed_setup后将其传给 Automodel 的from_pretrained实现显存友好的分片加载——每张 GPU 只加载属于自己的那部分权重分片。此外SALMAutomodel目前拒绝pp_size 1只有在模型集成补上流水线执行支持之后才能启用流水线并行。关于并行度的两个重要限制语音编码器 / 感知模块目前只支持 FSDP2通过dp_size控制MoE 的专家并行复用 FSDP2 的数据并行轴稠密层用 FSDP2 分片MoE 专家层用 EP 做 all-to-all 专家路由两者运行在同一组 GPU 上。ep_size只决定有多少 GPU 参与专家路由并不会新增一个独立的并行维度。激活检查点两个独立旋钮AutomodelParallelStrategy暴露了两个相互独立的激活检查点开关activation_checkpointing_llm一个开关同时覆盖两条路径——非 EP 的 FSDP2 路径强制FSDP2Config.activation_checkpointingTrue与 EP/MoE 并行器路径作为独立的运行时参数透传。无论ep_size是否为 1MoE LLM 都应使用此开关activation_checkpointing_perception在 FSDP2 分片之前用checkpoint_wrapper包裹perception.encoder.layers中的每一层以及 Conformer 的pre_encode前端——当它不是裸nn.Linear时。该逻辑实现在 nemo/collections/speechlm2/modules/perception.py 的AudioPerceptionModule.set_activation_checkpointing中且必须在 FSDP2 分片之前调用。两者默认均为false可以独立开启以便在模型两端分别做“算力换显存”的权衡。注意它们是 SALMAutomodel 专用开关HF Transformers 的 SALM 路径使用的是 HuggingFace 自身的 gradient-checkpointing API。使用 MoE LLM 骨干训练SALMAutomodel支持使用 Mixture-of-ExpertsMoE骨干的高效训练例如 NVIDIA Nemotron Nano V3约 30B 总参、3B 激活参数。NeMo Automodel 提供两项关键 MoE 优化Grouped GEMM把单个 MoE 层内所有专家的计算融合成一次批量矩阵乘法最大化 GPU 利用率与吞吐DeepEPDeep Expert Parallelism面向跨 GPU 专家路由的高效 all-to-all 通信原语显著降低专家并行的通信开销。原生 Nemotron-V3 骨干的已验证并行组合为 FSDP2、EP 与 CP具体要求如下Nemotron-V3 不提供 TP 方案因此必须设置tp_size: 1该路径不支持 HSDP因此dp_replicate_size: 1PP 尚未针对该 SpeechLM2 配置验证因此pp_size: 1CP 需要 packed sequences 与 Transformer EngineTE注意力详见下文“上下文并行”一节。在 8 卡上以 EP8 训练 SALMAutomodelNemotron Nano V3的示例torchrun --nproc_per_node8 examples/speechlm2/salm_train.py \ --config-namesalm_automodel \ model.pretrained_llmnvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 \ trainer.strategy.ep_size8分布式推理同样用torchrun启动salm_eval.py 会按tp_size/ep_size/pp_size/cp_size检测分布式需求并调用 parallel.py 中的setup_distributed便捷函数内部先init_process_group、设置 CUDA 设备再构造AutomodelParallelStrategy并创建设备网格然后以distributed_setup分片加载 checkpointtorchrun --nproc_per_node8 examples/speechlm2/salm_eval.py \ pretrained_namepath/to/checkpoint \ inputspath/to/manifest \ ep_size2Packed SequencesTHD消除变长语音的填充开销SALMAutomodel支持可选的 packed-sequenceTHD训练/验证路径把每条话语的文本 音频 embedding 拼接成一条扁平的[T_total, H]序列并用cu_seqlens索引标出各话语边界取代标准[B, T_max, H]BSHD布局的右填充。TE 的 varlen FlashAttention 按段计算、绝不跨话语相互 attentionMamba 的seq_idx也由同一cu_seqlens推导从而保证 SSM 状态在文档边界处正确重置。对于长度差异大的语音 batch填充开销非常可观BSHD 布局每个 minibatch 要支付B * (T_max - T_avg)的无效计算而 THD 只需要把每条话语舍入到2*cp_size的倍数TE 的 CP DualChunkSwap 模式需要。吞吐收益随话语长度方差增大而放大。按 batch 启用逐条注释见 salm_automodel.yamlmodel: packed_sequences: true # 可选项默认 falseBSHD automodel_backend: attn: te # THD 路径要求 TE 注意力当packed_sequences未设置时保持原有 BSHD 路径不变生成 / 推理始终使用 BSHD不走prepare_inputsTHD 打包在 nemo/collections/speechlm2/models/salm_automodel.py 中通过prepare_packed_llm_inputs完成训练与验证共用该路径。上下文并行CP长音频训练的关键SALMAutomodel支持在混合 Mamba/注意力 LLM如 Nemotron-V3上做长音频训练的上下文并行。CP 把序列维度分片到多张 GPU使每个 rank 的激活与 KV-cache 显存开销从T降到T / cp_size注意力层走 TE 的 DualChunkSwap 模式Mamba 混合器走 NeMo Automodel 中MambaContextParallel的 hidden-parallel all-to-all。在策略中启用trainer: strategy: _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy cp_size: 2 # 上下文并行大小必须整除每个 Mamba 块的 num_heads ep_size: 2 # 可与 CP 共享同一组 rankCP 下只支持 THD packed-sequence 路径每条话语是独立的 attention 段按话语级别的序列舍入天然满足 CP 的2*cp_size对齐要求。兼容性校验与三个已知坑仓库在 parallel.py 中实现了validate_parallelism_compatibility纯函数对三类已知不兼容组合做硬报错或告警相应测试见 test_salm_parallelism_validation.pyBSHD CP 不兼容硬报错TE 的 fused-attention CP 路径只支持causal、不支持padding_causal因此必须丢弃右填充掩码掩码一旦丢弃填充的 K/V 会经因果掩码泄漏进真实 token 的 attention首个优化步之后梯度变为NaN。文档记录的经验数据是2 卡上 BSHD CP2 EP2step 1 的loss4.62从 step 2 起lossnan。该问题与下文 TE/cuDNN 反向问题相互独立设置NVTE_FUSED_ATTN0也无法修复。解决办法是设置model.packed_sequences: true走 THD 路径。THD 要求 TE 注意力硬报错THD 打包通过 TE 的thd_get_partitioned_indices产生 2D[T_total, H]布局并送入 TE varlen FlashAttentionAutomodel 分支中 SDPA 的 3D-THD 路径存在缺陷transpose 假定 4D BSHD 输入因此packed_sequences: true时必须设置automodel_backend.attn: te。THD TE 在部分 GPU 上的梯度爆炸sm_120 硬报错其余架构告警在某些 GPU 架构尤其是 Blackwellsm_120上TransformerEngine 2.14 为qkv_formatthdattn_mask_typepadding_causal挑选的 cuDNN 后端前向结果正确但每层反向梯度被放大 8×–960×。经 LLM 注意力栈逐层累积后step 0 的梯度达到1e22量级按范数裁剪计算出1.0 / inf 0Adam 的矩最终变为 NaN。解决办法是在启动环境中设置NVTE_FUSED_ATTN0强制 TE 走 FlashAttention 派发要求已为你的 GPU 架构安装flash-attnFlashAttention 的 THD /padding_causal反向在相同形状下梯度是正确的。CP 下的数据加载是自动安全的SpeechLM2 的 datamodule 用 nemo/collections/common/data/lhotse/broadcasting.py 中的BroadcastingDataLoader包装 Lhotse 加载器——在 CP/TP 下每个 batch 只在 DP 源 rankcp_rank 0且tp_rank 0上构造一次再广播给子网格中的其他 rank。这消除了按 rank 独立的 Lhotse 非确定性concurrent_bucketing、worker 调度抖动等作为 CP 下 NCCL 死锁来源的问题。独立 API 说明见 dataloaders.rst。通用 ModelParallelStrategy 配置非 Automodel 路径除了AutomodelParallelStrategytrainer.strategy也支持通用的 ModelParallelStrategy对应 HF Transformers 后端的 SALM 路径trainer: strategy: _target_: nemo.core.ModelParallelStrategy find_unused_parameters: False data_parallel: 1 # 数据并行世界大小FSDP2 tensor_parallel: 8 # 张量并行世界大小 devices: 8 num_nodes: 1 accelerator: gpu precision: bf16-true模型的configure_model方法会根据这份配置自动完成并行化设置。FSDP2 配置HF Automodel 路径对 Fully Sharded Data Parallel 训练把data_parallel设为用于数据并行的 GPU 数量把tensor_parallel设为 1禁用。FSDP2 把模型参数分片到多张 GPU前向/反向时全部聚合、计算完成后立即释放从而在有限显存内训练更大的模型。这正是 salm_automodel.yaml 中use_distributed_sampler: False搭配AutomodelParallelStrategy的用法背景——分布式采样器参数由策略的distributed_sampler_kwargs基于 Automodel 展平的dp子网格提供。张量并行配置HF Automodel 路径对张量并行把tensor_parallel设为用于张量并行的 GPU 数量把data_parallel设为 1需要 2D 并行时可调大。parallelize_module函数把并行化计划应用到具体模型组件例如把注意力头或 embedding 维度拆分到多张 GPU。实现细节configure_model 中的三件事模型并行的核心实现集中在模型类的configure_model()方法中Automodel 路径见 salm_automodel.py关键点包括模块分片Module Sharding对模块调用fully_shard把参数分布到各数据并行 rank并行化计划Parallelization Plans创建并应用描述各层如何并行的计划例如通过parallelize_module拆分注意力头与 embedding 维度模型特定适配Model-Specific Adaptations处理不同 LLM 之间的架构差异如 Mamba 与 Transformer 混合架构、MoE 层、num_heads对 CP 的对齐要求等。Automodel 路径下的顺序尤其重要感知编码器的激活检查点包装必须先于 FSDP2 分片checkpoint_wrapper要看到未被分片的原始层对象fully_shard再索引最终结构LoRA 适配器在 Automodel 的from_pretrained内分片之前应用若未传入设备网格则在configure_model中另行安装。高级用法脚本定制注意事项自定义训练脚本时文档建议留意以下几点路径覆盖Path Overrides按需在 YAML 配置文件中覆盖数据、checkpoint、日志等路径——例如salm_automodel.yaml中的data.train_ds.input_cfg[].cuts_path、exp_manager.explicit_log_dir均为待填项需要替换为你自己的实际路径WB 密钥WB Keys如需 WB 日志在配置文件中更新 Weights Biases API 密钥并启用exp_manager.create_wandb_loggerBatch Size 调优根据 GPU 显存与模型大小调整data.train_ds.batch_size对长度差异大的语音数据可改用batch_tokens上限配合 bucketinguse_bucketing: true来稳定吞吐——相关字段在 salm_automodel.yaml 的data段中已给出注释说明batch_tokens、max_tokens、bucket_duration_bins等均以“音频帧 输出文本 token”之和计音频帧数按token_equivalent_duration换算。小结SpeechLM2 的训练扩展是一个层层递进的体系先用 SLURM 脚本与随机种子编排把实验跑起来用torchrun在本地快速调试再依据模型规模与集群拓扑选择并行维度——小模型用纯 FSDP2MoE 骨干如 Nemotron-V3组合 FSDP2 EP CP超长音频启用 THD packed-sequence 并配合 CP 摊薄序列维度。关键的经验教训是并行维度的组合并非任意可行BSHD 与 CP 互斥、THD 强依赖 TE 注意力、sm_120 上需NVTE_FUSED_ATTN0规避 cuDNN 反向梯度放大——这些约束都已内置到validate_parallelism_compatibility的运行时校验中会在配置不合规时第一时间报错而不是等到训练发散才暴露。【免费下载链接】SpeechA scalable generative AI framework built for researchers and developers working on Large Language Models, Multimodal, and Speech AI (Automatic Speech Recognition and Text-to-Speech)项目地址: https://gitcode.com/GitHub_Trending/nem/Speech创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考