Megatron-FSDP v2 是什么?
为什么 NVLink domain 内外可以用不同的 ZeRO 级别?
先弄清楚:每个概念是什么
- MFSDP v2
- Megatron 自己的 FSDP 第二版。0.19.0 的 release notes 还写着 "MFSDP v2 is under development",0.20.0 通过一个 experimental 的 Megatron Core adapter 接了进来。
- ZeRO 级别
- 配置里叫 sharding strategy,四级是
no_shard(DDP)、optim(ZeRO-1)、optim_grads(ZeRO-2)、optim_grads_params(ZeRO-3)。每级切什么见 DP、ZeRO、FSDP。 - NVLink domain
- 一组互相用 NVLink 直连的 GPU。例:GTP 文档把 GB200 NVL72 叫作 "a 72-GPU NVLink domain"。
- Hybrid FSDP
- 把 DP 的 rank 排成两层。inner 轴
dp_shard是 domain 里的 N 张卡,outer 轴dp_outer是 M 个 domain,每层各选一个 ZeRO 级别。
看图:domain 内一个级别,domain 之间另一个级别
30 秒回答
MFSDP v2 is Megatron's own fully sharded data parallel. It supports no-shard and ZeRO-1, 2 and 3, and puts data-parallel ranks on two axes, each with its own level. The 0.20.0 release notes say the levels can differ inside and across NVLink domains, without saying why. My reading: ZeRO-3 gathers parameters for every layer, cheap on NVLink but expensive across nodes, so you shard fully inside a domain and sync across domains once per step. It has its own FullyShardedOptimizer. In 0.20.0 the adapter is experimental and requires TP and expert TP of one.
两层各切到哪一级
Design doc 不给 HSDP、HFSDP 写特例,每个 DP 轴各自指定 parameter、gradient、optimizer state 怎么放。
| 名字 | outer 轴 | inner 轴 | parameter 和 gradient | optimizer state |
|---|---|---|---|---|
| HSDP | no_shard | ZeRO-3 | inner 上 sharded (N);outer 上 parameter 是 replicated (M),gradient 是 partial (M) | sharded (N),replicated (M) |
| HFSDP | ZeRO-1 | ZeRO-3 | 同上 | sharded (N × M) |
PR #6873 写 "ZeRO-2 and ZeRO-3 on the outer axis come for free"。
手算 7B 模型,model states 是 \(16\Psi\) = 112 GB,2 个 domain 各 8 张卡,N = 8、M = 2。16 张卡全部 ZeRO-3 时每卡 7 GB,但每层的 parameter all-gather 要跨 domain。HSDP 每卡 \(16\Psi/8\) = 14 GB,all-gather 只在 domain 内。HFSDP 每卡 \(4\Psi/8 + 12\Psi/16\) = 3.5 + 5.25 = 8.75 GB。这是估算,不是 Megatron 给的数字。
为什么按 NVLink domain 分层(推断)
0.20.0 的 release notes 只有一句 "The user can specify different ZeRO levels for inter-nvlink domains and intra-nvlink domains.",没有给理由。
ZeRO-3 的每个 unit 在 forward 和 backward 前都要 all-gather 一次 parameter。DGX H100 上 NVLink 单向约 450 GB/s,跨机每卡的网卡是 50 GB/s,差约 9 倍(数字见 TP communication)。所以 inner 轴放在 domain 里切到 ZeRO-3。outer 轴的 gradient reduction 每步只做一次,在最后一个 micro-batch 之后,HFSDP 另外在下一步开头把更新后的 parameter all-gather 回来。PyTorch FSDP 的 HYBRID_SHARD 和 HSDP 是同一个布局,文档给的理由是 "expensive all-gathers and reduce-scatters are only done within a node"。
限制和组合
| 项 | core 0.20.0 release notes 的说法 |
|---|---|
| TP | "the adapter requires TP=1 and expert TP=1" |
| 能组合 | context、expert、pipeline parallel 和 hybrid data-parallel sharding |
| Checkpoint | 用 PyTorch distributed checkpoint。"optimizer checkpointing currently excludes expert-parallel layouts" |
| CUDA Graph | full-iteration capture 只限 "supported dense-model configurations",原因见 CUDA Graph |
| Low precision | MXFP8 compute 和 quantized parameter gather(见 low precision) |
这些限制变得很快,EP 的 checkpoint 限制在 main 上已经修掉了(PR #8230)。
必答点
- MFSDP v2 是 Megatron 自己的 FSDP,0.20.0 里是 experimental 的 adapter,支持
no_shard和 ZeRO-1/2/3 - Hybrid FSDP 有 inner 和 outer 两个轴,各选一级。HSDP 的 outer 是
no_shard,HFSDP 的 outer 是 ZeRO-1,inner 都是 ZeRO-3 - 分层的理由是推断:每层的 parameter all-gather 留在 NVLink 上,跨 domain 每步只同步一次
- Optimizer 是
FullyShardedOptimizer,DistributedOptimizer留给 DDP 和 MFSDP v1 - 0.20.0 的限制:TP=1、expert TP=1;optimizer checkpoint 不含 EP layout;full-iteration CUDA Graph 只限 dense
追问
为什么不复用 DistributedOptimizer?
PR #5865 的理由是 DistributedOptimizer "has strong assumptions about DDP API/classes such as DistributedDataParallel and ParamAndGradBuffer"。MFSDP v1 的复用被 PR 称为 "shallow",大部分 DDP 的代码靠 if use_megatron_fsdp 分支绕过去。
outer 轴由什么配置决定,dense 和 MoE 部分能分开配吗?
PR #6873 写 num_distributed_optimizer_instances "decides only whether an outer axis exists",outer_dp_sharding_strategy "decides what lives on it"。能分开配,release notes 写 "different configs for the dense section and the MoE section",expert 的 mesh 由 expert_num_distributed_optimizer_instances 管(PR #7012)。
TP=1 这个限制会一直在吗?
不能当定论。9 月 12 日合入的 PR #7220 还写着 "TP and PP remain unsupported",9 月 23 日 PR #6486 就加上了 PP。Q4 roadmap 里 MFSDP 的条目是 tuning 和加 Muon、PP/VPP、combined 1F1B,没有 TP。
来源core 0.20.0 release notes(Parallelism、Performance and Memory、Precision and Optimizers 三节)· core 0.19.0 release notes · PR #5865(FullyShardedOptimizer)· PR #6873(HSDP、HFSDP 和两个轴)· MFSDP design doc(Sharding Strategies 的表)· PR #7012 · PR #7075 · PR #7220 · PR #6486 · PR #8230 · MFSDP v2 roadmap #6812 · Q4 roadmap #9232 · PyTorch FSDP 文档(ShardingStrategy.HYBRID_SHARD)· 题库条目在 ml-八股 06 的“零、团队最近在做什么”