SuperAPI AI 工具集

nemo-mbridge-perf-megatron-fsdp

部署与运维

查看项目

更新于 10-16 09:51

技能介绍

用途:指导在 Megatron-Bridge 中启用 Megatron 的全分片数据并行(Megatron FSDP),面向训练性能调优与配置排错。

主要能力:提供最小化配置覆盖、示例 recipe 修正、性能测试开关、代码锚点(Bridge 配置定义、校验逻辑、运行时包装器选择)以及验证方法。

必要用法:最小覆盖需设置 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"。性能测试可用:python scripts/performance/launch.py --use_megatron_fsdp true。

限制与陷阱:Megatron FSDP 与 use_torch_fsdp2 互斥;不支持 use_tp_pp_dp_mapping;仅支持 fsdp_dtensor 检查点格式,若启用保存/加载需从默认 torch_dist 切换;CPU offloading 仅在 pipeline_model_parallel_size == 1 且关闭激活重计算时有效;Hopper 及更早架构上 FSDP 与 TP/CP 可能要求不同的 CUDA_DEVICE_MAX_CONNECTIONS。验证使用 2 卡功能冒烟测试: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,成功标准为 1 passed、末次迭代损失有限且无检查点格式断言。