nemo-mbridge-perf-megatron-fsdp

Megatron FSDP スキル

安定性に関する背景情報と推奨レベルについては、以下を参照してください。

• @docs/training/megatron-fsdp.md
• @skills/nemo-mbridge-perf-megatron-fsdp/card.yaml

有効化

Bridge における最小限の Megatron FSDP オーバーライド：

cfg.dist.use_megatron_fsdp = True
cfg.ddp.use_megatron_fsdp = True
cfg.ddp.data_parallel_sharding_strategy = "optim_grads_params"
cfg.ddp.average_in_collective = False
cfg.checkpoint.ckpt_format = "fsdp_dtensor"

レシピの修正例：

cfg = llama3_8b_pretrain_config()
cfg.dist.use_megatron_fsdp = True
cfg.ddp.use_megatron_fsdp = True
cfg.ddp.data_parallel_sharding_strategy = "optim_grads_params"
cfg.ddp.average_in_collective = False
cfg.checkpoint.ckpt_format = "fsdp_dtensor"
cfg.checkpoint.save = "/tmp/fsdp_ckpts"
cfg.checkpoint.load = None

パフォーマンスハーネスに関する注記：

python scripts/performance/launch.py --use_megatron_fsdp true

コードの参照箇所

Bridge の設定定義：

use_megatron_fsdp: bool = False
"""Use Megatron's Fully Sharded Data Parallel. Cannot be used together with use_torch_fsdp2."""

use_torch_fsdp2: bool = False
"""Use the torch FSDP2 implementation. FSDP2 is not currently working with Pipeline Parallel.
It is still not in a stable release stage, and may therefore contain bugs or other
potential issues."""

Bridge の検証：

if self.dist.use_megatron_fsdp and self.dist.use_torch_fsdp2:
    raise ValueError(...)
...
assert not self.dist.use_tp_pp_dp_mapping, "use_tp_pp_dp_mapping is not supported with Megatron FSDP"
...
assert self.checkpoint.ckpt_format == "fsdp_dtensor", (
    "Megatron FSDP only supports fsdp_dtensor checkpoint format"
)

ランタイムラッパーの選択：

if use_megatron_fsdp:
    DP = FullyShardedDataParallel
elif use_torch_fsdp2:
    DP = TorchFullyShardedDataParallel
else:
    DP = DistributedDataParallel
...
DP(
    config=get_model_config(model_chunk),
    ddp_config=ddp_config,
    module=model_chunk,
    ...
    pg_collection=pg_collection,
)

パフォーマンスハーネスのオーバーライド：

recipe.ddp.use_megatron_fsdp = True
recipe.ddp.data_parallel_sharding_strategy = "optim_grads_params"
recipe.ddp.keep_fp8_transpose_cache = False
recipe.ddp.average_in_collective = False
...
recipe.checkpoint.load = None

注意点

1. 公開レシピでは use_megatron_fsdp が公開されていても、ckpt_format="torch_dist" がデフォルトのままになっていることがよくあります。保存／読み込みが有効な場合は、fsdp_dtensor に切り替えてください。
2. use_torch_fsdp2 は存在しますが、検証済みのブランチでは、_ddp_wrap が pg_collection を渡すため、トレーニング開始前に Bridge が失敗します。
3. CPU オフロードが有効なのは、pipeline_model_parallel_size == 1 かつアクティベーションの再計算が無効な場合のみです。
4. 上流の注意書きによると、Hopper 以前の環境では、FSDP と TP/CP で異なる CUDA_DEVICE_MAX_CONNECTIONS の設定が必要になる場合があります。
5. Megatron FSDP と FSDP2 は併用できません。

検証

既存の 2 GPU 機能スモークテストを使用してください。

CUDA_VISIBLE_DEVICES=0,1 uv run python -m torch.distributed.run --nproc_per_node=2 \
  -m pytest tests/functional_tests/training/test_megatron_fsdp.py::TestMegatronFSDP::test_fsdp_pretrain_basic -v -s

成功条件：

• Pytest が 1 passed と報告する
• ログに最終イテレーションの有限な損失が表示される
• チェックポイント形式のアサーションに失敗せず、実行が完了する