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 のエキスパート並列化

ep_size > 1 を指定して strategy: fsdp2 を推奨してください。これにより別個の 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 は単純なデータ並列化のみであることにも言及してください。

ストラテジーの選択

次の 3 つのストラテジーを利用でき、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 を使用してください。
• マルチノード：適切な TP／PP サイズを指定した fsdp2 を使用してください。
• エキスパート並列化を伴う MoE モデル：ep_size > 1 を指定した fsdp2 を使用してください（別個の moe_mesh が作成されます）。
• 大規模モデル（70B 以上）：PP と TP を組み合わせた fsdp2 を使用してください。
• 長いシーケンス（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 の場合に限り、MoEParallelizerConfig に対応する moe: サブ設定を含めてください。
• メインの 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 ベースのシャーディングを通じて、トランスフォーマーブロックごとに fully_shard() が適用されます。

テンソル並列を使用したFSDP2

TPは単一のNVLinkドメイン（通常は1ノード）内に収めます。

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

AutoPipelineがステージ間でレイヤーを分割する方法を把握できるように、モデルには_pp_plan属性（HFモデルクラスに設定）が必要です。_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

これはメッシュトポロジーではなく、モデル構築／トレーニングの動作を指定するフラグです。Dense戦略では戦略設定からこの値を読み取ります。EP／MoEの経路では、レシピレベルのフラグをモデルインフラストラクチャに直接渡します。

勾配同期の延期

FSDP2では、通信をオーバーラップさせるため、デフォルトで最後のマイクロバッチまで勾配同期を延期します。

distributed:
  defer_fsdp_grad_sync: true   # default

混合精度

FSDP2Configでは、MixedPrecisionPolicy(param_dtype=bf16, reduce_dtype=bf16, output_dtype=bf16, cast_forward_inputs=True)を介して、3つの精度設定すべてがデフォルトでbfloat16になります。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（one-forward-one-backward、デフォルト）
• 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のAttentionが必要です。SDPBackend.MATHはDTensorと互換性がありません。
• Attentionマスクは自動的に取り除かれます。attach_context_parallel_hooks()によって登録されたforward pre-hookが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 Attentionモデルでは、make_cp_batch_for_te()がTHD形式と、シャーディング用のTEのthd_get_partitioned_indicesを使用します。

シーケンスパッキング時のCP

CPはパッキングされたシーケンスで動作します。packed_sequence_sizeはcp_sizeで割り切れる必要があります。TEを使用する場合、_shard_thd_chunk_for_te()を介してチャンクごとにシャーディングします。

シーケンスパッキング

効率化のために、複数のシーケンスを1つのトレーニングサンプルにパッキングします。

設定

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 を使用してください。

並列化のサイズ設定ガイドライン

Dense モデル

モデルサイズ	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 モデルは、トークンごとにアクティブになるパラメーターが一部のみのため、総パラメーター数が同程度の Dense モデルよりも少ない 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 ドメイン（1 ノード、通常は 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 対応の状態辞書保存が必要です。 最終的な HF エクスポートには save_consolidated: final を指定して safetensors を使用するか、オフラインエクスポートには save_consolidated: false を指定し、生成された model/consolidate.sh ヘルパーを使用してください。

5. CP には互換性のある Attention が必要です。 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 が発生します。

検証

要求された戦略を実行する最小のレシピを実行してください。成功の条件は、終了コードが 0 であること、損失が有限であること、NCCL タイムアウトが発生しないこと、そしてログ出力が期待される TP/PP/CP/EP サイズと一致することです。