mcpField
N
nemo-mbridge-perf-megatron-fsdp

nemo-mbridge-perf-megatron-fsdp

github.com/promptingcompany开发039 次安装17 周前更新
代码仓库

关于 nemo-mbridge-perf-megatron-fsdp

nemo-mbridge-perf-megatron-fsdp 是用于在 Bridge 训练任务中配置和检查 Megatron FSDP 的技能，涵盖必要覆盖项、检查点设置、约束和功能冒烟测试。

主要功能

  1. 01Megatron FSDP 配置指导
  2. 02检查点格式验证
  3. 03兼容性约束说明
  4. 04双 GPU 冒烟测试

适用场景

  1. 01在 Bridge 配方中启用 Megatron FSDP
  2. 02排查检查点配置错误

安装

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

$npx skills add https://github.com/promptingcompany/nv-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