NeMo 并行训练完全指南从 DDP 到 AutomodelParallelStrategy 的 Speech AI 多卡扩展方案【免费下载链接】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 语音 AI 框架中多 GPU/多节点训练依赖一套层次分明的并行策略默认的 DDP 覆盖绝大多数 ASR、TTS 与 Audio 工作负载而 SpeechLM2 集合则通过AutomodelParallelStrategy与ModelParallelStrategy提供 FSDP2、张量并行TP、序列并行SP、上下文并行CP、专家并行EP与混合分片数据并行HSDP等大规模扩展手段。本文以 docs/source/features/parallelisms.rst 为核心结合 parallel.py 源码、salm_automodel.yaml 等真实配置与 test_parallel.py 测试用例系统讲解各并行维度的原理、适用场景与可复制的最小配置读完即可根据模型规模与 GPU 拓扑选择合适的策略并落地训练。一、DDP所有集合的默认并行策略Distributed Data ParallelismDDP是 NeMo 全部集合ASR、TTS、Audio、SpeechLM2的默认策略。它的工作方式非常直接在每个 GPU 上复制完整模型每个 GPU 消费不同的数据分片每次反向传播后通过 all-reduce 同步参数梯度。适用场景当完整模型可以放入单张 GPU 显存时DDP 是最简单高效的选择——这覆盖了绝大多数 ASR、TTS 与 Audio 训练负载。它的优势在于无额外分片通信开销、实现简单、与 PyTorch 生态完全一致。NeMo 默认即启用 DDP也可以在配置中显式声明YAML Hydra 方式trainer: strategy: _target_: lightning.pytorch.strategies.DDPStrategy gradient_as_bucket_view: true find_unused_parameters: true或者通过 Python 直接构造from lightning.pytorch.strategies import DDPStrategy trainer pl.Trainer( strategyDDPStrategy(gradient_as_bucket_viewTrue, find_unused_parametersTrue), devices8, acceleratorgpu, )关键参数说明gradient_as_bucket_view: true让梯度直接以桶视图bucket view形式存在避免反向传播时为梯度分配额外内存显著降低峰值显存占用find_unused_parameters: true允许模型存在未被使用的参数例如冻结部分的 ASR 编码器防止 DDP 在梯度同步时因找不到梯度而报错。SpeechLM2 的冻结策略如freeze_params正则与此配合良好。在真实配置中SpeechLM2 的非 Automodel 训练脚本默认也走 DDP例如 salm.yaml 中 trainer 段默认使用lightning.pytorch.strategies.DDPStrategy并在注释中提示“Replace DDPStrategy with ModelParallelStrategy to enable model parallelism”即把 DDP 替换为模型并行策略即可开启 FSDP2/TP。二、AutomodelParallelStrategySpeechLM2 大模型扩展的核心当 SpeechLM2 模型例如SALMAutomodel的骨干 LLM 无法放入单张 GPU 时NeMo 提供了nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy。这是一个 Lightning 策略其核心设计是把设备网格device mesh的创建委托给 NeMo Automodel并原生支持FSDP2完全分片数据并行Tensor ParallelismTP张量并行Sequence ParallelismSP序列并行Context ParallelismCP上下文并行Expert ParallelismEPMoE 专家并行Hybrid Sharded Data ParallelismHSDP混合分片数据并行适用场景训练或微调骨干 LLM 无法放入单卡显存的 SpeechLM2 模型或者希望比 DDP 更高效地把训练扩展到大量 GPU。AutomodelParallelStrategy专门用于SALMAutomodel以及 MoE LLM 骨干如 NVIDIA Nemotron Nano V3。硬性要求每个模型必须实现configure_model()方法定义其层如何分片与并行化。SALMAutomodel已实现该方法并在configure_model()中接收 Automodel 解析好的distributed_setup。换句话说不能不加实现地把任意模型从 DDP 直接切换到AutomodelParallelStrategy——缺少configure_model()的并行化定义会导致策略无法对模型进行分片。2.1 六个并行维度逐一拆解FSDP2Fully Sharded Data Parallelism在数据并行维度上对模型参数、梯度与优化器状态进行全量分片大幅降低单卡显存占用——这是相对 DDP 最本质的收益使原本放不下的模型可以训练。由dp_size控制当dp_size为null时NeMo Automodel 会根据 world size 与其它并行维度自动推断数据并行大小。Tensor ParallelismTP把单个权重矩阵切分到多张 GPU例如一个大线性层的权重按列或按行切分每张 GPU 只持有其中一片。由tp_size控制。模型必须定义 TP 分片方案哪些层被切分、如何切分Automodel 支撑的 SpeechLM2 模型对骨干 LLM 使用 Automodel 内置的 TP 方案。Sequence ParallelismSP在 TP 组内沿着序列维度分散激活内存通常与 TP 一起启用进一步压低激活显存。通过distributed_config.sequence_parallel: true开启。需要注意的是SP 依赖 TP 维度tp_size 1时才有意义源码与配置注释中均有说明。Context ParallelismCP在上下文并行组内切分长序列处理。由cp_size控制。对 SpeechLM2 模型而言CP 面向打包序列packed-sequence训练——每条语音被当作独立的注意力片段处理。源码 parallel.py 中的validate_parallelism_compatibility明确校验packed_sequencesfalseBSHD 路径与cp_size 1不兼容会直接抛错必须改用model.packed_sequences: true走 THD 路径。Expert ParallelismEP面向 MoE LLM 骨干把 MoE 专家路由到多张 GPU。由ep_size控制。EP复用 FSDP2 的数据并行轴稠密层通过 FSDP2 分片MoE 专家层则在同样的 ranks 上使用 all-to-all 专家路由而不是额外增加一个数据并行维度。这与配置注释中“ep_sizecontrols MoE expert routing on the>trainer: devices: 8 num_nodes: 1 accelerator: gpu precision: bf16-true strategy: _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy dp_size: null # inferred from world_size / other dimensions dp_replicate_size: 1 # HSDP replication group size tp_size: 1 # tensor parallel size cp_size: 1 # context parallel size ep_size: 8 # expert parallel size for MoE models distributed_config: sequence_parallel: false activation_checkpointing_llm: false activation_checkpointing_perception: falsePython 写法from nemo.collections.speechlm2.parts.parallel import AutomodelParallelStrategy trainer pl.Trainer( strategyAutomodelParallelStrategy( dp_sizeNone, dp_replicate_size1, tp_size1, cp_size1, ep_size8, ), devices8, acceleratorgpu, precisionbf16-true, use_distributed_samplerFalse, )重要提醒使用AutomodelParallelStrategy时必须在 trainer 中设置use_distributed_samplerFalse——NeMo 的数据模块内部自行处理分布式采样。2.3 源码级实现原理从 parallel.py 可以确认该策略的完整工作链路AutomodelParallelStrategy继承自 PyTorch Lightning 的ModelParallelStrategy测试用例test_is_subclass_of_model_parallel_strategy也验证了这一点见 test_parallel.py构造时把所有并行维度存入内部字段默认值均为 1dp_sizeNone表示待推断并支持timeout_minutes参数YAML 友好地绕过 timedelta 的 allowlist 限制适用于 rank-0 从 Lustre 加载模型超时的场景create_device_mesh()调用nemo_automodel.components.distributed.DistributedSetup.build(...)把ParallelismSizes(dp_size, dp_replicate_size, tp_size, pp_size, cp_size, ep_size)与FSDP2Config默认、MoEParallelizerConfig一起解析得到维度为(pp, dp_replicate, dp_shard, cp, tp)的 device meshsetup_environment()在初始化进程组后自动调用create_device_mesh()并把device_mesh、moe_mesh、distributed_setup挂到 LightningModule 上供configure_model()读取使用distributed_sampler_kwargs使用 Automodel 扁平化的dp子网格覆盖dp_replicate * dp_shard返回num_replicas与rank这正是必须关闭use_distributed_sampler的底层原因。此外parallel.py中还有一个独立的纯函数setup_distributed(...)它初始化torch.distributed、设置 CUDA 设备并创建 device mesh返回一个distributed_setup已就绪的策略对象——这是为推理脚本在无 Lightning Trainer 场景下做模型并行加载准备的便捷入口。2.4 激活检查点两个独立开关AutomodelParallelStrategy暴露两个可独立启用的激活检查点旋钮activation_checkpointing_llm对 LLM transformer 块做检查点。单个开关同时覆盖标准 FSDP2 路径与 EP/MoE 并行器路径源码注释确认前者通过强制FSDP2Config.activation_checkpointingTrue后者由 EP 并行器作为独立运行时参数读取。因此无论ep_size是 1 还是更大MoE LLM 骨干都应使用这个开关activation_checkpointing_perception在 FSDP2 分片之前对语音感知编码器perception encoder的 transformer 层应用checkpoint_wrapper。两者默认均为false开启后以反向传播时额外重计算为代价换取激活显存下降trainer: strategy: _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy activation_checkpointing_llm: true activation_checkpointing_perception: true测试 test_parallel.py 专门验证了两个开关相互独立可以只开其一且均能通过 YAML 中的 Hydra 实例化正确传递test_automodel_strategy_with_ac_flags_from_yaml。2.5 三种典型配置示例示例一仅 FSDP2最简配置适用于各层都能放入单卡显存的情况8 张 GPU 全部用作数据并行分片trainer: devices: 8 strategy: _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy dp_size: 8 tp_size: 1 ep_size: 1示例二MoE 专家并行如 NVIDIA Nemotron Nano V3稠密层走 FSDP2MoE 层做 8 路专家路由dp_size: null让 Automodel 自动推断数据并行大小trainer: devices: 8 strategy: _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy dp_size: null tp_size: 1 ep_size: 8示例三TP FSDP2更大规模稠密骨干2 路张量并行把每层切到 2 张 GPU剩余 ranks 由 NeMo Automodel 推断为 FSDP2 数据并行大小trainer: devices: 8 strategy: _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy dp_size: null tp_size: 2 ep_size: 12.6 真实配置中的完整策略块仓库中的 salm_automodel.yaml以及 salm_automodel_pee.yaml提供了生产级参考devices: 8、num_nodes: 1、precision: bf16-flash、use_distributed_sampler: Falsestrategy 块中dp_size: null、dp_replicate_size: 1、tp_size: 1、pp_size: 1SALMAutomodel 当前要求pp_size1、cp_size: 1、ep_size: 8。配置中还给出了可选的扩展项distributed_config纯 dict自动解析为FSDP2Configsequence_parallel需tp_size 1时开启可配置offload_policy如torch.distributed.fsdp.CPUOffloadPolicy实现 CPU 卸载moe_config纯 dict自动解析为MoEParallelizerConfigignore_router_for_ac默认 trueAC 重计算时保存 router 输出、reshard_after_forward前向后重新分片省显存但增加通信、lm_head_precision如float32提高稳定性、wrap_outer_model等automodel_backend.dispatcher: torch仅当 GPU 具备 NVLINK/NVSHMEM 时才可设为deepep。这些 dict 到类型化 config 的转换由nemo.utils.trainer_utils._resolve_automodel_configs完成test_parallel.py 的TestResolveAutomodelConfigs用例验证了纯 dict、空 dict、嵌套_target_如mp_policy用 Hydra 实例化MixedPrecisionPolicy等各类写法。三、ModelParallelStrategySALM 与 Duplex 的并行路径原始的 SpeechLM2SALM与 Duplex 模型配置直接使用 PyTorch Lightning 的ModelParallelStrategy。这条路径与SALMAutomodel相互独立基于 PyTorch 原生 DTensor 支持 FSDP2、TP 与 SP。适用场景非 Automodel 的 SpeechLM2 模型SALM、Duplex。Automodel 支撑的模型如SALMAutomodel则必须用AutomodelParallelStrategy。硬性要求与AutomodelParallelStrategy相同模型必须实现configure_model()定义层的分片与并行化方式SpeechLM2 的 SALM 与 Duplex 模型已经实现。3.1 配置方法约束条件data_parallel_size与tensor_parallel_size的乘积必须等于总 GPU 数devices * num_nodes。例如 8 卡下4 × 2 8trainer: devices: 8 num_nodes: 1 accelerator: gpu precision: bf16-true strategy: _target_: lightning.pytorch.strategies.ModelParallelStrategy data_parallel_size: 4 # FSDP2: shard across 4 GPUs tensor_parallel_size: 2 # TP: split layers across 2 GPUsPython 写法from lightning.pytorch.strategies import ModelParallelStrategy trainer pl.Trainer( strategyModelParallelStrategy( data_parallel_size4, tensor_parallel_size2, ), devices8, acceleratorgpu, precisionbf16-true, use_distributed_samplerFalse, )重要提醒与 Automodel 路径一致使用ModelParallelStrategy时同样必须设置use_distributed_samplerFalse。在 salm.yaml 这类配置中切换方式就是把注释掉的ModelParallelStrategy块启用并填上tensor_parallel_size/data_parallel_size注释明确标注data_parallel_size: 8 # This is FSDP2。其余非 Automodel 训练入口如 s2s_duplex.yaml、salm_asr_decoder.yaml 等也都保留了同样的切换注释。四、策略选型速查与注意事项场景推荐策略关键维度模型可放入单卡ASR/TTS/Audio 大多数情况DDP默认无需额外配置Automodel SpeechLM2SALMAutomodel层可入单卡AutomodelParallelStrategy FSDP2dp_size: 8MoE 骨干如 Nemotron Nano V3AutomodelParallelStrategy EPep_size: 8,dp_size: null稠密骨干超大模型AutomodelParallelStrategy TP FSDP2tp_size: 2,dp_size: null非 Automodel 的 SALM / DuplexModelParallelStrategydata_parallel_size × tensor_parallel_size 总卡数需要留意的实践要点两个策略都要求use_distributed_samplerFalseNeMo 数据模块内部完成分布式采样源码中distributed_sampler_kwargs基于 Automodel 扁平化 dp 子网格计算见 parallel.py配置的并行维度必须与总 GPU 数兼容否则 Automodel 解析或模型加载阶段会报错configure_model()是并行化的前提策略本身不定义如何分片分片方案完全由模型的configure_model()提供CP 只面向打包序列训练packed_sequencesfalse时使用cp_size 1会触发 parallel.py 中的硬性兼容性校验并抛错需改用 THD 路径THD 路径还要求automodel_backend.attnte且在 Blackwell sm_120 上需设置NVTE_FUSED_ATTN0强制 FlashAttention 分发否则 TE 2.14 的 cuDNN fused-attention 反向核可能把梯度放大 8x–960x 导致 NaNvalidate_parallelism_compatibility会根据设备能力选择硬报错或告警相关校验测试见 test_salm_parallelism_validation.py激活检查点按需开启先开activation_checkpointing_llm同时覆盖 FSDP2 与 EP/MoE 路径感知编码器显存吃紧再开activation_checkpointing_perception完整的训练配置数据、优化器、调度器可参考examples/speechlm2/conf/下的 yaml以及 training_and_scaling.rst 中的 SLURM/torchrun 启动方式FSDP2/TP 训练产出的分布式 checkpointDCP说明见 checkpoints/intro.rst。五、结语从默认的 DDP 到AutomodelParallelStrategy支持的六维并行NeMo 为 Speech AI 训练提供了一条平滑的扩展路径小模型直接 DDPAutomodel 大模型按“FSDP2 → EP → TP”逐步升级非 Automodel 的 SALM/Duplex 则走ModelParallelStrategy的 DTensor 路线。理解每个维度的显存与通信权衡、牢记use_distributed_samplerFalse与configure_model()两个前提再配合仓库中开箱即用的 YAML 配置与源码级兼容性校验即可在多卡集群上稳定地把 SpeechLM2 训练规模推上去。【免费下载链接】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），仅供参考 NeMo 并行训练完全指南：从 DDP 到 AutomodelParallelStrategy 的 Speech AI 多卡扩展方案
NeMo 并行训练完全指南从 DDP 到 AutomodelParallelStrategy 的 Speech AI 多卡扩展方案【免费下载链接】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 语音 AI 框架中多 GPU/多节点训练依赖一套层次分明的并行策略默认的 DDP 覆盖绝大多数 ASR、TTS 与 Audio 工作负载而 SpeechLM2 集合则通过AutomodelParallelStrategy与ModelParallelStrategy提供 FSDP2、张量并行TP、序列并行SP、上下文并行CP、专家并行EP与混合分片数据并行HSDP等大规模扩展手段。本文以 docs/source/features/parallelisms.rst 为核心结合 parallel.py 源码、salm_automodel.yaml 等真实配置与 test_parallel.py 测试用例系统讲解各并行维度的原理、适用场景与可复制的最小配置读完即可根据模型规模与 GPU 拓扑选择合适的策略并落地训练。一、DDP所有集合的默认并行策略Distributed Data ParallelismDDP是 NeMo 全部集合ASR、TTS、Audio、SpeechLM2的默认策略。它的工作方式非常直接在每个 GPU 上复制完整模型每个 GPU 消费不同的数据分片每次反向传播后通过 all-reduce 同步参数梯度。适用场景当完整模型可以放入单张 GPU 显存时DDP 是最简单高效的选择——这覆盖了绝大多数 ASR、TTS 与 Audio 训练负载。它的优势在于无额外分片通信开销、实现简单、与 PyTorch 生态完全一致。NeMo 默认即启用 DDP也可以在配置中显式声明YAML Hydra 方式trainer: strategy: _target_: lightning.pytorch.strategies.DDPStrategy gradient_as_bucket_view: true find_unused_parameters: true或者通过 Python 直接构造from lightning.pytorch.strategies import DDPStrategy trainer pl.Trainer( strategyDDPStrategy(gradient_as_bucket_viewTrue, find_unused_parametersTrue), devices8, acceleratorgpu, )关键参数说明gradient_as_bucket_view: true让梯度直接以桶视图bucket view形式存在避免反向传播时为梯度分配额外内存显著降低峰值显存占用find_unused_parameters: true允许模型存在未被使用的参数例如冻结部分的 ASR 编码器防止 DDP 在梯度同步时因找不到梯度而报错。SpeechLM2 的冻结策略如freeze_params正则与此配合良好。在真实配置中SpeechLM2 的非 Automodel 训练脚本默认也走 DDP例如 salm.yaml 中 trainer 段默认使用lightning.pytorch.strategies.DDPStrategy并在注释中提示“Replace DDPStrategy with ModelParallelStrategy to enable model parallelism”即把 DDP 替换为模型并行策略即可开启 FSDP2/TP。二、AutomodelParallelStrategySpeechLM2 大模型扩展的核心当 SpeechLM2 模型例如SALMAutomodel的骨干 LLM 无法放入单张 GPU 时NeMo 提供了nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy。这是一个 Lightning 策略其核心设计是把设备网格device mesh的创建委托给 NeMo Automodel并原生支持FSDP2完全分片数据并行Tensor ParallelismTP张量并行Sequence ParallelismSP序列并行Context ParallelismCP上下文并行Expert ParallelismEPMoE 专家并行Hybrid Sharded Data ParallelismHSDP混合分片数据并行适用场景训练或微调骨干 LLM 无法放入单卡显存的 SpeechLM2 模型或者希望比 DDP 更高效地把训练扩展到大量 GPU。AutomodelParallelStrategy专门用于SALMAutomodel以及 MoE LLM 骨干如 NVIDIA Nemotron Nano V3。硬性要求每个模型必须实现configure_model()方法定义其层如何分片与并行化。SALMAutomodel已实现该方法并在configure_model()中接收 Automodel 解析好的distributed_setup。换句话说不能不加实现地把任意模型从 DDP 直接切换到AutomodelParallelStrategy——缺少configure_model()的并行化定义会导致策略无法对模型进行分片。2.1 六个并行维度逐一拆解FSDP2Fully Sharded Data Parallelism在数据并行维度上对模型参数、梯度与优化器状态进行全量分片大幅降低单卡显存占用——这是相对 DDP 最本质的收益使原本放不下的模型可以训练。由dp_size控制当dp_size为null时NeMo Automodel 会根据 world size 与其它并行维度自动推断数据并行大小。Tensor ParallelismTP把单个权重矩阵切分到多张 GPU例如一个大线性层的权重按列或按行切分每张 GPU 只持有其中一片。由tp_size控制。模型必须定义 TP 分片方案哪些层被切分、如何切分Automodel 支撑的 SpeechLM2 模型对骨干 LLM 使用 Automodel 内置的 TP 方案。Sequence ParallelismSP在 TP 组内沿着序列维度分散激活内存通常与 TP 一起启用进一步压低激活显存。通过distributed_config.sequence_parallel: true开启。需要注意的是SP 依赖 TP 维度tp_size 1时才有意义源码与配置注释中均有说明。Context ParallelismCP在上下文并行组内切分长序列处理。由cp_size控制。对 SpeechLM2 模型而言CP 面向打包序列packed-sequence训练——每条语音被当作独立的注意力片段处理。源码 parallel.py 中的validate_parallelism_compatibility明确校验packed_sequencesfalseBSHD 路径与cp_size 1不兼容会直接抛错必须改用model.packed_sequences: true走 THD 路径。Expert ParallelismEP面向 MoE LLM 骨干把 MoE 专家路由到多张 GPU。由ep_size控制。EP复用 FSDP2 的数据并行轴稠密层通过 FSDP2 分片MoE 专家层则在同样的 ranks 上使用 all-to-all 专家路由而不是额外增加一个数据并行维度。这与配置注释中“ep_sizecontrols MoE expert routing on the>trainer: devices: 8 num_nodes: 1 accelerator: gpu precision: bf16-true strategy: _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy dp_size: null # inferred from world_size / other dimensions dp_replicate_size: 1 # HSDP replication group size tp_size: 1 # tensor parallel size cp_size: 1 # context parallel size ep_size: 8 # expert parallel size for MoE models distributed_config: sequence_parallel: false activation_checkpointing_llm: false activation_checkpointing_perception: falsePython 写法from nemo.collections.speechlm2.parts.parallel import AutomodelParallelStrategy trainer pl.Trainer( strategyAutomodelParallelStrategy( dp_sizeNone, dp_replicate_size1, tp_size1, cp_size1, ep_size8, ), devices8, acceleratorgpu, precisionbf16-true, use_distributed_samplerFalse, )重要提醒使用AutomodelParallelStrategy时必须在 trainer 中设置use_distributed_samplerFalse——NeMo 的数据模块内部自行处理分布式采样。2.3 源码级实现原理从 parallel.py 可以确认该策略的完整工作链路AutomodelParallelStrategy继承自 PyTorch Lightning 的ModelParallelStrategy测试用例test_is_subclass_of_model_parallel_strategy也验证了这一点见 test_parallel.py构造时把所有并行维度存入内部字段默认值均为 1dp_sizeNone表示待推断并支持timeout_minutes参数YAML 友好地绕过 timedelta 的 allowlist 限制适用于 rank-0 从 Lustre 加载模型超时的场景create_device_mesh()调用nemo_automodel.components.distributed.DistributedSetup.build(...)把ParallelismSizes(dp_size, dp_replicate_size, tp_size, pp_size, cp_size, ep_size)与FSDP2Config默认、MoEParallelizerConfig一起解析得到维度为(pp, dp_replicate, dp_shard, cp, tp)的 device meshsetup_environment()在初始化进程组后自动调用create_device_mesh()并把device_mesh、moe_mesh、distributed_setup挂到 LightningModule 上供configure_model()读取使用distributed_sampler_kwargs使用 Automodel 扁平化的dp子网格覆盖dp_replicate * dp_shard返回num_replicas与rank这正是必须关闭use_distributed_sampler的底层原因。此外parallel.py中还有一个独立的纯函数setup_distributed(...)它初始化torch.distributed、设置 CUDA 设备并创建 device mesh返回一个distributed_setup已就绪的策略对象——这是为推理脚本在无 Lightning Trainer 场景下做模型并行加载准备的便捷入口。2.4 激活检查点两个独立开关AutomodelParallelStrategy暴露两个可独立启用的激活检查点旋钮activation_checkpointing_llm对 LLM transformer 块做检查点。单个开关同时覆盖标准 FSDP2 路径与 EP/MoE 并行器路径源码注释确认前者通过强制FSDP2Config.activation_checkpointingTrue后者由 EP 并行器作为独立运行时参数读取。因此无论ep_size是 1 还是更大MoE LLM 骨干都应使用这个开关activation_checkpointing_perception在 FSDP2 分片之前对语音感知编码器perception encoder的 transformer 层应用checkpoint_wrapper。两者默认均为false开启后以反向传播时额外重计算为代价换取激活显存下降trainer: strategy: _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy activation_checkpointing_llm: true activation_checkpointing_perception: true测试 test_parallel.py 专门验证了两个开关相互独立可以只开其一且均能通过 YAML 中的 Hydra 实例化正确传递test_automodel_strategy_with_ac_flags_from_yaml。2.5 三种典型配置示例示例一仅 FSDP2最简配置适用于各层都能放入单卡显存的情况8 张 GPU 全部用作数据并行分片trainer: devices: 8 strategy: _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy dp_size: 8 tp_size: 1 ep_size: 1示例二MoE 专家并行如 NVIDIA Nemotron Nano V3稠密层走 FSDP2MoE 层做 8 路专家路由dp_size: null让 Automodel 自动推断数据并行大小trainer: devices: 8 strategy: _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy dp_size: null tp_size: 1 ep_size: 8示例三TP FSDP2更大规模稠密骨干2 路张量并行把每层切到 2 张 GPU剩余 ranks 由 NeMo Automodel 推断为 FSDP2 数据并行大小trainer: devices: 8 strategy: _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy dp_size: null tp_size: 2 ep_size: 12.6 真实配置中的完整策略块仓库中的 salm_automodel.yaml以及 salm_automodel_pee.yaml提供了生产级参考devices: 8、num_nodes: 1、precision: bf16-flash、use_distributed_sampler: Falsestrategy 块中dp_size: null、dp_replicate_size: 1、tp_size: 1、pp_size: 1SALMAutomodel 当前要求pp_size1、cp_size: 1、ep_size: 8。配置中还给出了可选的扩展项distributed_config纯 dict自动解析为FSDP2Configsequence_parallel需tp_size 1时开启可配置offload_policy如torch.distributed.fsdp.CPUOffloadPolicy实现 CPU 卸载moe_config纯 dict自动解析为MoEParallelizerConfigignore_router_for_ac默认 trueAC 重计算时保存 router 输出、reshard_after_forward前向后重新分片省显存但增加通信、lm_head_precision如float32提高稳定性、wrap_outer_model等automodel_backend.dispatcher: torch仅当 GPU 具备 NVLINK/NVSHMEM 时才可设为deepep。这些 dict 到类型化 config 的转换由nemo.utils.trainer_utils._resolve_automodel_configs完成test_parallel.py 的TestResolveAutomodelConfigs用例验证了纯 dict、空 dict、嵌套_target_如mp_policy用 Hydra 实例化MixedPrecisionPolicy等各类写法。三、ModelParallelStrategySALM 与 Duplex 的并行路径原始的 SpeechLM2SALM与 Duplex 模型配置直接使用 PyTorch Lightning 的ModelParallelStrategy。这条路径与SALMAutomodel相互独立基于 PyTorch 原生 DTensor 支持 FSDP2、TP 与 SP。适用场景非 Automodel 的 SpeechLM2 模型SALM、Duplex。Automodel 支撑的模型如SALMAutomodel则必须用AutomodelParallelStrategy。硬性要求与AutomodelParallelStrategy相同模型必须实现configure_model()定义层的分片与并行化方式SpeechLM2 的 SALM 与 Duplex 模型已经实现。3.1 配置方法约束条件data_parallel_size与tensor_parallel_size的乘积必须等于总 GPU 数devices * num_nodes。例如 8 卡下4 × 2 8trainer: devices: 8 num_nodes: 1 accelerator: gpu precision: bf16-true strategy: _target_: lightning.pytorch.strategies.ModelParallelStrategy data_parallel_size: 4 # FSDP2: shard across 4 GPUs tensor_parallel_size: 2 # TP: split layers across 2 GPUsPython 写法from lightning.pytorch.strategies import ModelParallelStrategy trainer pl.Trainer( strategyModelParallelStrategy( data_parallel_size4, tensor_parallel_size2, ), devices8, acceleratorgpu, precisionbf16-true, use_distributed_samplerFalse, )重要提醒与 Automodel 路径一致使用ModelParallelStrategy时同样必须设置use_distributed_samplerFalse。在 salm.yaml 这类配置中切换方式就是把注释掉的ModelParallelStrategy块启用并填上tensor_parallel_size/data_parallel_size注释明确标注data_parallel_size: 8 # This is FSDP2。其余非 Automodel 训练入口如 s2s_duplex.yaml、salm_asr_decoder.yaml 等也都保留了同样的切换注释。四、策略选型速查与注意事项场景推荐策略关键维度模型可放入单卡ASR/TTS/Audio 大多数情况DDP默认无需额外配置Automodel SpeechLM2SALMAutomodel层可入单卡AutomodelParallelStrategy FSDP2dp_size: 8MoE 骨干如 Nemotron Nano V3AutomodelParallelStrategy EPep_size: 8,dp_size: null稠密骨干超大模型AutomodelParallelStrategy TP FSDP2tp_size: 2,dp_size: null非 Automodel 的 SALM / DuplexModelParallelStrategydata_parallel_size × tensor_parallel_size 总卡数需要留意的实践要点两个策略都要求use_distributed_samplerFalseNeMo 数据模块内部完成分布式采样源码中distributed_sampler_kwargs基于 Automodel 扁平化 dp 子网格计算见 parallel.py配置的并行维度必须与总 GPU 数兼容否则 Automodel 解析或模型加载阶段会报错configure_model()是并行化的前提策略本身不定义如何分片分片方案完全由模型的configure_model()提供CP 只面向打包序列训练packed_sequencesfalse时使用cp_size 1会触发 parallel.py 中的硬性兼容性校验并抛错需改用 THD 路径THD 路径还要求automodel_backend.attnte且在 Blackwell sm_120 上需设置NVTE_FUSED_ATTN0强制 FlashAttention 分发否则 TE 2.14 的 cuDNN fused-attention 反向核可能把梯度放大 8x–960x 导致 NaNvalidate_parallelism_compatibility会根据设备能力选择硬报错或告警相关校验测试见 test_salm_parallelism_validation.py激活检查点按需开启先开activation_checkpointing_llm同时覆盖 FSDP2 与 EP/MoE 路径感知编码器显存吃紧再开activation_checkpointing_perception完整的训练配置数据、优化器、调度器可参考examples/speechlm2/conf/下的 yaml以及 training_and_scaling.rst 中的 SLURM/torchrun 启动方式FSDP2/TP 训练产出的分布式 checkpointDCP说明见 checkpoints/intro.rst。五、结语从默认的 DDP 到AutomodelParallelStrategy支持的六维并行NeMo 为 Speech AI 训练提供了一条平滑的扩展路径小模型直接 DDPAutomodel 大模型按“FSDP2 → EP → TP”逐步升级非 Automodel 的 SALM/Duplex 则走ModelParallelStrategy的 DTensor 路线。理解每个维度的显存与通信权衡、牢记use_distributed_samplerFalse与configure_model()两个前提再配合仓库中开箱即用的 YAML 配置与源码级兼容性校验即可在多卡集群上稳定地把 SpeechLM2 训练规模推上去。【免费下载链接】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），仅供参考