Megatron-FSDP v2 是什么?
为什么 NVLink domain 内外可以用不同的 ZeRO 级别?

英文问法What is Megatron-FSDP v2, and which ZeRO levels can it shard to? · Why would you pick different sharding levels inside and across NVLink domains? · Why does it have its own optimizer instead of reusing the distributed optimizer?

考过推断,非面经。没有任何面经问过这道题 · 依据是团队自己的材料:core 0.20.0 release notes 的 Parallelism 一节、PR #5865 和 #6873 的描述、MFSDP design doc、roadmap #6812 · 这些功能在 0.20.0 标着 experimental,下面每条限制都写明是哪个版本的说法

先弄清楚:每个概念是什么

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 之间另一个级别

2 个 NVLink domain,每个画 4 张卡,N = 4,M = 2 NVLink domain 0 NVLink domain 1 GPU0 GPU1 GPU2 GPU3 GPU4 GPU5 GPU6 GPU7 inner 轴 dp_shard:ZeRO-3 每层的 parameter all-gather,走 NVLink inner 轴 dp_shard:ZeRO-3 每层的 parameter all-gather,走 NVLink outer 轴 dp_outer:no_shard(HSDP)或 ZeRO-1(HFSDP) 每步一次的 gradient reduction,走跨 domain 的网络
图 1:把 inner 轴对到 domain 内是按 release notes 那句话读的,PR 和 design doc 没有提 NVLink。

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 和 gradientoptimizer state
HSDPno_shardZeRO-3inner 上 sharded (N);outer 上 parameter 是 replicated (M),gradient 是 partial (M)sharded (N),replicated (M)
HFSDPZeRO-1ZeRO-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 Graphfull-iteration capture 只限 "supported dense-model configurations",原因见 CUDA Graph
Low precisionMXFP8 compute 和 quantized parameter gather(见 low precision)

这些限制变得很快,EP 的 checkpoint 限制在 main 上已经修掉了(PR #8230)。

必答点

追问

为什么不复用 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 的“零、团队最近在做什么”