⚡

MegatronFSDP训练调优技能Skill nemo-mbridge-perf-megatron-fsdp

本技能提供在 Megatron-Bridge 中启用 Megatron FSDP（全分片数据并行）的操作指南，涵盖配置项、代码锚点、常见陷阱及验证方法。适用于使用 FSDP 替代 DDP 进行分布式训练，或排查相关 OOM 和性能回退问题。关键词：Megatron FSDP、全分片数据并行、分布式训练、性能调优、Megatron-Bridge、FSDP 配置、检查点格式。

Megatron-Core训练 0 次安装 0 次浏览 更新于 9/7/2026
名称 nemo-mbridge-perf-megatron-fsdp
描述 操作指南，用于在Megatron-Bridge中启用Megatron FSDP，包括配置旋钮、代码锚点、陷阱和验证。
开源协议 Apache-2.0 when_to_use: 使用基于FSDP的数据并行而不是DDP，或将OOM或回归追踪到FSDP配置更改；‘use_megatron_fsdp’、‘data_parallel_sharding_strategy’、‘sharded data parallel’、‘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存在，但在验证分支上，Bridge在训练前仍然失败，因为_ddp_wrap传递了pg_collection。
  3. 仅当pipeline_model_parallel_size == 1且禁用激活重计算时，CPU卸载才有效。
  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
  • 日志显示最后一次迭代的损失为有限值
  • 运行结束，没有出现检查点格式断言