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 仅支持 简单数据并行。

策略选择

有三种策略，可通过 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：无需分布式配置（当 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_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, 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／MoE 路径会将配方级标志直接传入模型基础设施。

延迟梯度同步

默认情况下，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: final

工作原理

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 的限制

尽管名称中包含 FSDP，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

Mesh 上下文与原始 mesh 创建：

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。 最终导出 HF 模型时，请使用 safetensors 和 save_consolidated: final；也可以使用 save_consolidated: false，并通过生成的 model/consolidate.sh 辅助脚本离线导出。

5. CP 要求使用兼容的注意力机制。 仅支持 SDPA（Flash Attention 或 Efficient Attention）或 TE attention。SDPBackend.MATH 与 DTensor 不兼容。

6. MoE 的 EP 大小必须能整除 dp_size * cp_size。 创建设备 mesh 时会断言 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、loss 为有限值、不发生 NCCL 超时，且日志输出中的 TP/PP/CP/EP 大小与预期相符。