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 の場合にのみ 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 ベースのシャーディングを通じてトランスフォーマーブロックごとに 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 と互換性がありません。

FSDP2 と HSDP（ハイブリッドシャーデッドデータ並列）

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 では、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 とステージの対応を示す dict）を定義する必要があります。
2. 分散セクションで pp_size > 1 に設定する必要があります。
3. スケジュールとマイクロバッチサイズを指定する pipeline サブ設定が必要です。

サポートされているスケジュール

PipelineConfig.pp_schedule で定義されています。

• 1f1b（フォワード 1 回、バックワード 1 回、デフォルト）
• 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() は各モデルパートで （TE モデル以外の場合）attach_context_parallel_hooks() を呼び出します。
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() によって チャンクごとにシャーディングされます。

シーケンスパッキング

複数のシーケンスを 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 はスループットを大幅に低下させます。

コードの参照先

• 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 対応の state dict 保存が必要です。 チェックポイントとの互換性のため、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 が発生します。

検証

要求されたストラテジーを実行する最小のレシピを実行してください。終了コードが 0 で、損失が有限値であり、NCCL タイムアウトがなく、ログ出力に期待される TP/PP/CP/EP サイズが反映されていれば成功です。