nemo-automodel-distributed-training

NeMo AutoModel 中的分布式训练

目的

NeMo AutoModel 使用原生 PyTorch 分布式训练。 所有并行策略都通过单个 MeshContext 对象进行协调，该对象包含设备网格、策略配置和轴名称。

说明

对于概念性的分布式训练问题，请直接根据本技能中的快速模式作答，无需检查仓库。先说明策略选择，然后仅列出与问题相关的 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 仅支持简单的数据并行。

策略选择

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

策略	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：无需分布式配置（当 world_size=1 时，FSDP2Manager 会跳过并行化）。
• 单节点多 GPU：使用 fsdp2（默认选项）。仅在需要尽可能简单的设置时使用 ddp。
• 多节点：使用 fsdp2，并设置合适的 TP/PP 规模。
• 使用专家并行的 MoE 模型：使用 fsdp2 并设置 ep_size > 1（会创建独立的 moe_mesh）。
• 大型模型（70B+）：使用 fsdp2，并结合 PP + TP。
• 长序列（8K+）：添加 CP（cp_size > 1）。

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

TP + PP 快速回答：

• 使用 strategy: fsdp2；需要流水线并行时不要使用 megatron_fsdp。
• 设置 tp_size 以启用张量并行，设置 pp_size 以启用流水线并行。
• 添加包含 pp_schedule 和 pp_microbatch_size 的 pipeline: 子配置。
• 将 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_setup.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)

基础设施流程

YAML distributed section
    -> parse_distributed_section()          [recipes/_dist_setup.py]
    -> setup_distributed()                  [recipes/_dist_setup.py]
        -> create_device_mesh()             [components/distributed/device_mesh.py]
        -> MeshContext(...)                  [components/distributed/mesh.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, or 8 -- must divide GPUs per node
  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, etc.
    pp_microbatch_size: 4
    scale_grads_in_schedule: false

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

使用 HSDP（混合分片数据并行）的 FSDP2

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

distributed:
  strategy: fsdp2
  dp_replicate_size: 2   # must divide dp_size

限制：dp_replicate_size < dp_size（FSDP2 不支持不分片的纯复制）。

激活检查点

通过在反向传播期间重新计算激活，以计算换取内存：

distributed:
  activation_checkpointing: true

对于非 EP 模型，此设置会传递给策略配置；对于 EP 模型，则从 MeshContext.activation_checkpointing 读取。

梯度同步延迟

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

distributed:
  defer_fsdp_grad_sync: true   # default

混合精度

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. 分布式配置中 pp_size > 1。
3. 必须提供包含调度策略和微批次大小的 pipeline 子配置。

支持的调度策略

定义于 PipelineConfig.pp_schedule：

• 1f1b（一次前向、一次反向，默认）
• gpipe
• interleaved_1f1b / interleaved1f1b
• looped_bfs
• dfs
• v_schedule
• zero_bubble

示例（8 张 GPU 上运行 8B 模型，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: true

工作原理

AutoPipeline.build() 会调用 pipeline_model()，使用模型的 _pp_plan 将模型拆分为多个阶段，创建 PipelineStage 对象，并构建调度策略。训练期间，schedule.step() 会驱动流水线中的前向传播和反向传播。

上下文并行

长序列（8K+）使用 CP。CP 将 Q/K/V 作为 DTensor 沿序列维度分片。

配置

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

要求

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

工作原理

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() 按块进行分片。

序列打包

为提高效率，将多个序列打包到单个训练样本中。

配置

packed_sequence:
  packed_sequence_size: 4096   # 0 = disabled

step_scheduler:
  local_batch_size: 1          # must be 1 for packed sequences

当 packed_sequence_size > 0 时，数据集整理器会将序列打包到该长度。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 示例（8 张 GPU 上的 Qwen3-30B-A3B）

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 更少，因为每个词元只有一部分参数处于激活状态。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。
• TP 跨 InfiniBand 会严重降低吞吐量。

代码定位

• components/distributed/config.py：FSDP2Config、MegatronFSDPConfig、DDPConfig。
• components/distributed/mesh.py：MeshContext、策略映射和网格大小。
• components/distributed/device_mesh.py：设备网格和 moe_mesh 创建。
• components/distributed/pipelining/config.py：PipelineConfig 字段。
• components/moe/config.py：MoEParallelizerConfig 和 MoEConfig。
• recipes/_dist_setup.py：YAML 解析和分布式设置。

注意事项

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 的状态字典保存方式。 请搭配 save_consolidated: true 使用 safetensors，以确保检查点兼容性。

5. CP 要求使用兼容的注意力机制。 仅支持 SDPA（Flash Attention 或 Efficient Attention）或 TE attention。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。

验证

运行能覆盖所需策略的最小 recipe。成功的标准是退出代码为 0、损失值有限、没有 NCCL 超时，并且日志输出中的 TP/PP/CP/EP 大小符合预期。