⚡

NeMoAutoModel分布式训练配置Skill nemo-automodel-distributed-training

本技能用于指导在 NeMo AutoModel 中配置和优化分布式训练。涵盖 FSDP2、MegatronFSDP 和 DDP 三种策略的选择，详细说明张量并行 (TP)、流水线并行 (PP)、上下文并行 (CP)、专家并行 (EP) 与 HSDP 的 YAML 配置方法，并提供大模型和 MoE 模型并行大小建议。关键词：NeMo AutoModel、分布式训练、FSDP2、MegatronFSDP、DDP、张量并行、流水线并行、上下文并行、专家并行、YAML 配置。

分布式训练调优 0 次安装 0 次浏览 更新于 9/7/2026

NeMo AutoModel 中的分布式训练

用途

NeMo AutoModel 使用 PyTorch 原生分布式训练。所有并行性均通过一个单一的 MeshContext 对象进行编排，该对象持有设备网格、策略配置和轴名称。

<!-- NVSkills catalog signing requested after PR #2937 (2026-07-31). -->

使用说明

有关概念性的分布式训练问题，请直接依据本技能中的快速模式回答，而无需检查仓库。首先选择策略，然后只列出与问题相关的 YAML 字段和约束。

在最终答案中使用直接行动动词：推荐策略，展示最小 YAML，陈述大小约束，明确指出不支持的策略。除非用户询问，否则不要讨论模型上架、配方、Slurm、SkyPilot 或检查点。

示例

用于大型多节点模型的 TP 加 PP

推荐 strategy: fsdp2。提及 tp_size、pp_size、cp_size、ep_size 和 pipeline 子配置。说明 dp_size 是由 world_size / (tp_size * pp_size * cp_size) 推断的。

distributed:
  strategy: fsdp2
  tp_size: 8
  pp_size: 4
  cp_size: 1
  ep_size: 1
  pipeline:
    pp_schedule: interleaved1f1b
    pp_microbatch_size: 1

MoE 专家并行

推荐 strategy: fsdp2 且 ep_size > 1。说明这会创建一个独立的 moe_mesh；在相关时包含 moe 子配置；说明 ep_size 必须整除 dp_size * cp_size。不要推荐 megatron_fsdp 或 ddp。

distributed:
  strategy: fsdp2
  ep_size: 8
  moe:
    reshard_after_forward: false

MegatronFSDP 的限制

对流水线并行、专家并行和 sequence_parallel 说不。对于 PP、EP 或 sequence_parallel，推荐 fsdp2；提及 DDP 仅支持简单数据并行。

策略选择

有三种策略可用，通过 distributed.strategy YAML 键选择：

策略 YAML 值 最适合
FSDP2 fsdp2 一般用途，推荐默认。支持 TP、PP、CP、EP、HSDP。
MegatronFSDP megatron_fsdp NVIDIA Megatron 风格的 FSDP。不支持 PP、EP、sequence_parallel。
DDP ddp 仅限简单数据并行。不支持 TP、PP、CP 或 EP。

决策树：

  • 单 GPU：无需分布式配置（FSDP2Manager 在 world_size=1 时跳过并行化）。
  • 多 GPU 单节点：fsdp2（默认）。仅当需要最简单的设置时使用 ddp。
  • 多节点：使用适当的 TP/PP 尺寸 fsdp2。
  • MoE 模型与专家并行：fsdp2 且 ep_size > 1（创建独立的 moe_mesh）。
  • 大模型（70B+）：使用 PP + TP 的 fsdp2。
  • 长序列（8K+）：增加 CP（cp_size > 1）。

回答策略选择问题时，先说明所选的 distributed.strategy，然后列举用户必须设置的 YAML 字段。

快速 TP + PP 回答：

  • 使用 strategy: fsdp2；当需要流水线并行时，不要使用 megatron_fsdp。
  • 设置 tp_size 用于张量并行，pp_size 用于流水线并行。
  • 添加 pipeline: 子配置，包含 pp_schedule 和 pp_microbatch_size。
  • 将 dp_size 留空或设置为 none；它被推断为 world_size / (tp_size * pp_size * cp_size)。
  • 尽可能将 TP 保持在快速节点内域中，对于 70B+ 模型使用 PP 跨模型深度。

快速 MoE 专家并行回答：

  • 从 strategy: fsdp2 和 ep_size > 1 开始。
  • 仅当 ep_size > 1 时包含 moe: 子配置；它映射到 MoEParallelizerConfig。
  • 期望在主要 device_mesh 之外另有一个用于专家并行的 moe_mesh。
  • 不要推荐 megatron_fsdp 或 ddp 用于专家并行；megatron_fsdp 不支持 EP。
  • 在完成 MoE EP 回答之前，明确说明 ep_size 必须整除 dp_size * cp_size，且 megatron_fsdp 不支持 EP、PP 或 sequence_parallel。

YAML 配置结构

配方 YAML 中的 distributed 部分直接映射到 recipes/_dist_utils.py 中的 parse_distributed_section()：

distributed:
  strategy: fsdp2           # fsdp2 | megatron_fsdp | ddp
  dp_size: none             # auto-calculated from world_size / (tp * pp * cp)
  dp_replicate_size: none   # FSDP2-only, for HSDP
  tp_size: 1
  pp_size: 1
  cp_size: 1
  ep_size: 1

  # Strategy-specific flags (forwarded to the strategy dataclass):
  sequence_parallel: false
  activation_checkpointing: false
  defer_fsdp_grad_sync: true   # FSDP2 only

  # Sub-configs (optional):
  pipeline:
    pp_schedule: 1f1b
    pp_microbatch_size: 1
    # ... see PipelineConfig fields

  moe:
    reshard_after_forward: false
    # ... see MoEParallelizerConfig fields

dp_size 总是被推断：

dp_size = world_size / (tp_size * pp_size * cp_size)

基础设施流程

initialize_distributed()                       [components/distributed/init_utils.py]
    -> initializes torch.distributed process group and returns DistInfo
YAML distributed section + DistInfo.world_size
    -> parse_distributed_section()          [recipes/_dist_utils.py]
    -> create_distributed_setup_from_config()              [recipes/_dist_utils.py]
        -> DistributedSetup.build()         [components/distributed/config.py]
    -> instantiate_infrastructure()         [_transformers/infrastructure.py]
        -> _instantiate_distributed()       -> FSDP2Manager / MegatronFSDPManager / DDPManager
        -> _instantiate_pipeline()          -> AutoPipeline (if pp_size > 1)
        -> parallelize_fn                   -> MoE parallelizer (if ep_size > 1) or PP wrapper
    -> apply_model_infrastructure()         [_transformers/infrastructure.py]
        -> _shard_pp() or _shard_ep_fsdp()  (applies sharding to the model)

FSDP2 配置

基础 FSDP2（仅数据并行）

distributed:
  strategy: fsdp2
  tp_size: 1
  cp_size: 1

这会自动计算 dp_size = world_size，并通过基于 DTensor 的分片对每个 transformer 块应用 fully_shard()。

FSDP2 与张量并行

将 TP 保持在单个 NVLink 域内（通常是一个节点）：

distributed:
  strategy: fsdp2
  tp_size: 4        # 2, 4 或 8 -- 必须整除每个节点的 GPU 数
  sequence_parallel: true

TP 计划根据模型类型自动选择。如有需要，可通过 Python API 传递自定义计划：

config = FSDP2Config(sequence_parallel=True, tp_plan=my_custom_plan)

FSDP2 与流水线并行

distributed:
  strategy: fsdp2
  pp_size: 2
  pipeline:
    pp_schedule: interleaved1f1b   # 1f1b, gpipe, interleaved_1f1b 等
    pp_microbatch_size: 4
    scale_grads_in_schedule: false

模型必须具有 _pp_plan 属性（在 HF 模型类上设置），AutoPipeline 才能知道如何将层拆分到不同阶段。没有 _pp_plan 的模型不兼容 PP。

FSDP2 与 HSDP（混合分片数据并行）

节点内全分片 + 节点间通过 2D DeviceMesh 复制：

distributed:
  strategy: fsdp2
  dp_replicate_size: 2   # 必须整除 dp_size

约束：dp_replicate_size < dp_size（FSDP2 不支持纯复制而不分片）。

激活检查点

以计算换内存，在反向传播时重新计算激活：

distributed:
  activation_checkpointing: true

这是模型构建/训练行为标志，并非网格拓扑。密集策略从策略配置中读取它；EP/MoE 路径将配方级标志直接传入模型基础设施。

梯度同步延迟

FSDP2 默认将梯度同步延迟到最后一个微批次，以实现通信重叠：

distributed:
  defer_fsdp_grad_sync: true   # 默认

混合精度

FSDP2Config 默认对所有三个精度旋钮使用 bfloat16，通过 MixedPrecisionPolicy(param_dtype=bf16, reduce_dtype=bf16, output_dtype=bf16, cast_forward_inputs=True)。可通过 Python API 覆盖：

from torch.distributed.fsdp import MixedPrecisionPolicy
config = FSDP2Config(
    mp_policy=MixedPrecisionPolicy(param_dtype=torch.float16, reduce_dtype=torch.float32),
)

流水线并行

要求

  1. 模型类必须定义 _pp_plan（一个将模块 FQN 映射到阶段的字典）。
  2. distributed 部分中 pp_size > 1。
  3. 包含调度和微批次大小的 pipeline 子配置。

支持的调度

定义于 PipelineConfig.pp_schedule：

  • 1f1b（一前一后，默认）
  • gpipe
  • interleaved_1f1b / interleaved1f1b
  • looped_bfs
  • dfs
  • v_schedule
  • zero_bubble

示例（8B 模型在 8 个 GPU 上，PP=2 + DP=4）

distributed:
  strategy: fsdp2
  pp_size: 2

  pipeline:
    pp_schedule: interleaved1f1b
    pp_microbatch_size: 4
    scale_grads_in_schedule: false

checkpoint:
  model_save_format: safetensors
  save_consolidated: final

工作原理

AutoPipeline.build() 调用 pipeline_model()，该函数使用模型的 _pp_plan 将模型拆分为阶段，创建 PipelineStage 对象并构建调度。训练期间，schedule.step() 驱动前向和反向通过流水线。

上下文并行

对于长序列（8K+），使用 CP。CP 在序列维度上以 DTensor 形式分片 Q/K/V。

配置

distributed:
  strategy: fsdp2
  cp_size: 2   # 或 4, 8

要求

  • SDPA（Flash Attention 或 Efficient Attention 后端）或 Transformer Engine 注意力。SDPBackend.MATH 与 DTensor 不兼容。
  • 注意力掩码被自动移除；is_causal=True 由 attach_context_parallel_hooks() 注册的前置钩子设置。

工作原理

  1. 模型分片后，apply_model_infrastructure() 在每个模型部件上调用 attach_context_parallel_hooks()（对于非 TE 模型）。
  2. 在每个训练步骤，make_cp_batch_and_ctx() 创建一个 CP 上下文管理器，沿序列维度分片批处理，并设置来自 torch.distributed.tensor.experimental 的 context_parallel()。
  3. 对于 TE 注意力模型，make_cp_batch_for_te() 使用 THD 格式和 TE 的 thd_get_partitioned_indices 进行分片。

CP 与序列打包

CP 支持打包序列。packed_sequence_size 必须能被 cp_size 整除。当使用 TE 时，通过 _shard_thd_chunk_for_te() 按 chunk 分片。

序列打包

将多个序列打包到一个训练样本中以提高效率。

配置

packed_sequence:
  packed_sequence_size: 4096   # 0 = 禁用

step_scheduler:
  local_batch_size: 1          # 打包序列必须为 1

当 packed_sequence_size > 0 时，数据集 collator 将序列打包到该长度。local_batch_size 必须为 1，因为每个“样本”已经是打包批。

MoE 分布式训练

专家并行

设置 ep_size > 1 将专家分布在 GPU 上。这会在主要 device_mesh 旁创建一个独立的 moe_mesh：

distributed:
  strategy: fsdp2
  ep_size: 8
  activation_checkpointing: true

moe_mesh 形状为 (pp_size, ep_shard_size, ep_size)，维度名为 ("pp", "ep_shard", "ep")。

约束：dp_cp_size（= dp_size * cp_size）必须能被 ep_size 整除。

MoE 子配置

distributed:
  strategy: fsdp2
  ep_size: 8
  activation_checkpointing: true

  moe:
    reshard_after_forward: false
    ignore_router_for_ac: false
    wrap_outer_model: true

moe 子部分映射到 MoEParallelizerConfig，仅在 ep_size > 1 时实例化。

完整 MoE 示例（Qwen3-30B-A3B 在 8 个 GPU 上）

distributed:
  strategy: fsdp2
  tp_size: 1
  cp_size: 1
  pp_size: 1
  ep_size: 8
  sequence_parallel: false
  activation_checkpointing: true

MegatronFSDP 的限制

尽管名称如此，megatron_fsdp 并不支持专家并行（ep_size > 1）、流水线并行（pp_size > 1）或 sequence_parallel。如需这些特性，请使用 fsdp2。

并行大小调整指南

稠密模型

模型大小 TP PP CP 策略
< 3B 1 1 1 FSDP2（仅 DP）
3-13B 2-4 1 1 FSDP2 + TP
13-70B 4-8 2-4 1 FSDP2 + TP + PP
70B+ 8 4-8 1 FSDP2 + TP + PP
任意 + 长序列（8K+） 如上 如上 2-8 添加 CP

MoE 模型

MoE 模型比相同总参数量的稠密模型需要更少的 TP，因为每个 token 只有一部分参数是激活的。EP 是主要扩展维度：

模型 TP PP EP 备注
小型 MoE（总 <10B） 1 1 8 仅 EP
中型 MoE（总 10-30B） 1-2 1 8 对共享层使用小 TP
大型 MoE（总 100B+） 1-2 4+ 8-64 PP 用于深度，EP 用于专家

硬件拓扑规则

  • TP 必须保持在单个 NVLink 域内（通常一个节点，通常 8 个 GPU）。
  • 使用 PP 或 DP 进行跨节点扩展。
  • 跨 InfiniBand 的 TP 会严重降低吞吐量。

程序化 API（from_pretrained / from_config）

不使用 YAML 配方时，可通过 Python 配置分布式训练：

from nemo_automodel.components.distributed import (
    DistributedSetup,
    FSDP2Config,
    ParallelismSizes,
    initialize_distributed,
)

dist_env = initialize_distributed("nccl")
distributed_setup = DistributedSetup.build(
    strategy=FSDP2Config(sequence_parallel=True),
    parallelism_sizes=ParallelismSizes(tp_size=2),
    activation_checkpointing=True,
    world_size=dist_env.world_size,
)

或直接传给 from_pretrained：

from nemo_automodel import NeMoAutoModelForCausalLM

model = NeMoAutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-1B",
    distributed_setup=distributed_setup,
)

代码锚点

策略配置数据类：

components/distributed/config.py
    FSDP2Config       -- sequence_parallel, tp_plan, mp_policy, offload_policy,
                         activation_checkpointing, defer_fsdp_grad_sync
    MegatronFSDPConfig -- zero_dp_strategy, overlap_grad_reduce, overlap_param_gather, etc.
    DDPConfig          -- activation_checkpointing only

MeshContext（并行性的单一真实来源）：

components/distributed/mesh.py
    MeshContext  -- device_mesh, moe_mesh
                    Properties: tp_size, pp_size, cp_size, ep_size, dp_size, dp_replicate_size
    MeshAxisName -- PP, DP, DP_REPLICATE, DP_SHARD, DP_SHARD_CP, DP_CP, CP, TP, EP, EP_SHARD

网格上下文和原始网格创建：

components/distributed/config.py
    DistributedSetup.build()      -- builds MeshContext from strategy + parallelism
components/distributed/mesh_utils.py
    _create_device_meshes()       -- routes to FSDP2/MegatronFSDP/DDP raw mesh creation
    _create_fsdp2_device_mesh()   -- shape (pp, dp_replicate, dp_shard, cp, tp) + flattened submeshes
    _create_megatron_fsdp_device_mesh() -- shape (dp, cp, tp)

分布式管理器：

components/distributed/fsdp2.py          -- FSDP2Manager.parallelize()
components/distributed/megatron_fsdp.py  -- MegatronFSDPManager.parallelize()
components/distributed/ddp.py            -- DDPManager

流水线并行：

components/distributed/pipelining/config.py        -- PipelineConfig dataclass
components/distributed/pipelining/autopipeline.py  -- AutoPipeline orchestrator
components/distributed/pipelining/functional.py    -- pipeline_model(), schedule creation
components/distributed/pipelining/hf_utils.py      -- HF model validation for PP

上下文并行：

components/distributed/context_parallel/utils.py
    make_cp_batch_and_ctx()            -- creates CP context manager + shards batch
    create_context_parallel_ctx()      -- wraps torch.distributed.tensor.experimental.context_parallel
    attach_context_parallel_hooks()    -- strips attention_mask, sets is_causal=True
    make_cp_batch_for_te()             -- TE-specific CP batch sharding (THD format)

基础设施编排：

_transformers/infrastructure.py
    instantiate_infrastructure()    -- config objects -> runtime objects
    apply_model_infrastructure()    -- applies sharding, PEFT, checkpoints to model
    _shard_pp()                     -- pipeline parallel path
    _shard_ep_fsdp()                -- EP + FSDP path (non-PP)

YAML 解析：

recipes/_dist_utils.py
    parse_distributed_section()  -- YAML dict -> typed configs + sizes
    create_distributed_setup_from_config()  -- recipe adapter: parse + create DistributedSetup; does not init process group

MoE 配置：

components/distributed/config.py
    MoEParallelizerConfig  -- reshard_after_forward, ignore_router_for_ac, wrap_outer_model, etc.
components/moe/config.py
    MoEConfig              -- n_routed_experts, n_activated_experts, score_func, etc.

常见陷阱

  1. **跨节点 TP 破坏吞吐量。**始终保持 TP 在单个 NVLink 域内。使用 PP 或 DP 进行跨节点扩展。

  2. **PP 需要模型类上的 _pp_plan。**并非所有 HF 模型都有。启用 PP 前检查 validate_hf_model_for_pipeline_support()。

  3. **PP 气泡降低 GPU 利用率。**使用交错调度（interleaved_1f1b）和更小的微批次以减少气泡时间。

  4. **FSDP2 需要支持 DTensor 的状态字典保存。**使用 safetensors 和 save_consolidated: final 进行最终 HF 导出，或使用 save_consolidated: false 加上生成的 model/consolidate.sh 帮助器进行离线导出。

  5. **CP 需要兼容的注意力。**仅限 SDPA（Flash Attention 或 Efficient Attention）或 TE 注意力。SDPBackend.MATH 与 DTensor 不兼容。

  6. **MoE EP 大小必须能整除 dp_size * cp_size。**设备网格创建断言 dp_cp_size % ep_size == 0。

  7. **MegatronFSDP 比 FSDP2 更受限。**它不支持 PP（pp_size > 1）、EP（ep_size > 1）或 sequence_parallel。MeshContext 验证会针对这些组合抛出异常。

  8. **DDP 仅支持数据并行。**不支持 TP、PP、CP、EP 或 HSDP。任何这些都会引发验证错误。

  9. **激活检查点增加计算量。**通过反向传播时重新计算激活来节省内存，但会增加约 30% 的计算开销。

  10. **混合精度策略必须与模型期望匹配。**默认的 bfloat16 策略适用于大多数模型。FP16 模型可能需要自定义 MixedPrecisionPolicy。

  11. 使用 CP 与打包序列时，packed_sequence_size 必须能被 cp_size 整除。

  12. **dp_replicate_size 仅限 FSDP2。**与 megatron_fsdp 或 ddp 一起传递会引发 ValueError。

验证

运行使用所请求策略的最小配方。成功意味着退出码为 0、损失有限、无 NCCL 超时，并且日志输出与预期的 TP/PP/CP/EP 大小匹配。