mcpField
N
nemo-mbridge-perf-megatron-fsdp

nemo-mbridge-perf-megatron-fsdp

github.com/nvidia开发3.8k1.8k 次安装今天更新
代码仓库

关于 nemo-mbridge-perf-megatron-fsdp

为 Megatron Bridge 训练配置 Megatron FSDP 的技能，适用于需要使用全分片数据并行的开发者。说明启用参数、检查点格式约束和验证方式。

主要功能

  1. 01配置 Megatron FSDP
  2. 02设置数据并行分片策略
  3. 03使用 fsdp_dtensor 检查点格式
  4. 04校验 FSDP 兼容性限制
  5. 05运行双 GPU 功能测试

适用场景

  1. 01为训练配方启用 Megatron FSDP
  2. 02调整分布式训练的分片参数
  3. 03验证 FSDP 训练与检查点保存

安装

复制下面的命令，在终端中运行。

$npx skills add https://github.com/nvidia/skills --skill nemo-mbridge-perf-megatron-fsdp

安装后在 Claude Code 中输入/nemo-mbridge-perf-megatron-fsdp确认已加载。

包含的文件

除 SKILL.md 外，安装时一起下载的参考资料和脚本。

BENCHMARK.md
card.yaml
skill-card.md
skill.oms.sig
evals/
evals.json

详细介绍

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
  • 日志显示最后一次迭代的损失值有限
  • 运行结束时没有出现检查点格式断言

评论

支持 Markdown