训练一个 70B 模型，你有 8 张 80GB 的 A100。NVIDIA 工程师跟你说：用 Megatron-LM。微软的人说：用 DeepSpeed。PyTorch 的人说：用 FSDP。三个人都拍胸脯保证自己能省 80% 显存、提速两倍、稳定不掉点。
我当年就在这个岔路口浪费了整整一周，起因是我天真地以为这三个框架是同一个问题的三种解法，挑个快的就行。后来才发现，它们解决的问题根本不在一个维度。选错框架，比选错显存策略更致命。
张量并行：把同一个矩阵拆开让多张卡一起算
先说 Megatron-LM。它来自 NVIDIA 2019 年的论文，核心是张量并行。一个 Transformer 里的 Linear 层，本质是矩阵乘法。普通数据并行是每张卡存一份完整权重、各自算各自的 batch；Megatron 反过来，把权重矩阵本身拆成几块，分别放在不同 GPU 上，计算时每张卡只做自己那块局部计算，再通过 all-reduce 把结果拼起来。
具体到实现，Megatron 会把 QKV 和 output 这些 Linear 的权重按行或列切开。每层 forward 至少做两轮 all-reduce 通信。一张卡跑一轮完整前向，等于把整个模型的通信链路都走了一遍。因此张量并行非常依赖高带宽——单机内的 NVLink 是它的主场，跨节点用它就是自杀式体验。
我在一个 30B 模型上试过 Megatron，改了一整天模型代码，把 attention、linear 全部替换成张量并行版本，终于跑通，但通信占比高得吓人。后来换成 FSDP，十分钟烧通，速度还差不多。这个经历让我记住了：张量并行解决的是单个权重矩阵过大、一张卡放不下的问题，而不是显存优化的一般解。
ZeRO 和 FSDP：不拆计算，拆冗余
DeepSpeed 的 ZeRO 来自微软 2020 年的论文，它的思路完全不同。数据并行训练时，每张卡都保存一份完整的参数、梯度和优化器状态。拿 10B 参数模型举例，混合精度下模型状态总占用约 160GB：fp16 参数 20GB，fp16 梯度 20GB，fp32 参数副本 40GB，Adam 的一阶二阶动量 80GB。8 张卡就是 8 份完全重复的存储。ZeRO 的洞察是：这些冗余没必要存在。
于是它把优化器状态（ZeRO-1）、梯度（ZeRO-2）、参数（ZeRO-3）按卡数均分。计算的前向/反向过程中，先 all-gather 取出完整参数，算完一块就释放，等需要时再取。每一张卡并不持有完整模型副本，但整组卡的算力仍然协同完成了一次完整训练。
FSDP 可以理解为 ZeRO-3 在 PyTorch 生态中的原生实现。它把模型按层递归包装，每一层是一个分片单元。默认情况下，每个 FSDP 实例持有该层参数的一个 shard，前向传播时先 all-gather 凑齐这一层的完整权重，计算完立刻释放，反向传播时再对梯度做 reduce-scatter 分片累积。用代码表达就是这么几行：
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
model = FSDP(model) # 等价于 ZeRO-3 的分片策略顺带一提，FSDP 的每一步通信都发生在不同层之间，因此可以和计算天然重叠——前向算第 1 层时，正好去通信第 2 层的参数。这也是它在单机 8 卡上速度不输 Megatron 的原因之一。
一张表看清三者的适用边界
下表是我根据源码、论文和实测整理的差异，横向看就能明白为什么有人喊着 "无脑上 FSDP"，也有人坚持 Megatron 不可替代。
| 维度 | Megatron-LM (TP) | DeepSpeed ZeRO-3 | FSDP |
|---|---|---|---|
| 切分对象 | 权重矩阵内部的维度 | 参数、梯度、优化器状态按卡均分 | 参数、梯度、优化器状态按层均分 |
| 通信模式 | 每层多次 all-reduce | forward all-gather，backward reduce-scatter | 与前同，但按层粒度频繁发出 |
| 通信量级 | 约等于张量大小 × 层数 × 2 | 约等于模型大小（比 DP 多个零头） | 与 ZeRO-3 同级 |
| 跨节点扩展 | 差，受限于带宽 | 优秀，可到数千卡 | 优秀 |
| 易用性 | 需要改模型定义 | 命令行集成，HF ecosystem 友好 | PyTorch 原生 API，几行代码搞定 |
| 额外能力 | 流水线并行、序列并行 | CPU/NVMe offload、通信压缩 | torch.compile 集成、meta device |
注意，FSDP 和 ZeRO-3 在通信量上并没有本质差别，差别在实现粒度。FSDP 以模块为单位递归分片，能够跟 PyTorch 的计算流调度器重叠得更细；DeepSpeed 则在通信压缩、offload 和批处理优化上做得更老辣。
如果把三种方案放在一起，你会发现它们的通信开销跟模型结构的关系完全不同。张量并行的通信量正比于 hidden_size × num_layers，几乎不随 batch size 变化；ZeRO/FSDP 的通信量正比于参数总量，和 batch size 无关。所以当 batch 很小时，张量并行更亏；当 batch 很大时，两者都还可以。这个规律是我被拷打了好几次才总结出来的。
实战选型：我踩过的坑和你的行动路线
如果你现在要训一个模型，可以按下面顺序拍板：
- 模型能完整塞进单卡（比如 7B 的 fp16 约 14GB，80G 卡无压力）——直接用普通 DDP，或者加一个 ZeRO-1/2 就好，没必要引入分片通信。
- 单卡装不下，但单机 8 卡能装下——优先用 FSDP 或 ZeRO-3。它们的通信模式对 NVLink 友好，同时以后想扩到多机也不用改代码。
- 如果模型大到单机上怎么分都放不下（单个权重矩阵超过单卡显存），必须上 Megatron 的张量并行——只有它能切开单个矩阵。
- 多节点场景下，用 FSDP/ZeRO-3 做横向扩展，必要时与 Megatron 的 TP+PP 混合。NVIDIA 和微软已经把这个组合做成了 Megatron-DeepSpeed。
另一个实用的坑：如果你只有一两张卡但想训大模型，DeepSpeed 的 ZeRO-Offload 可以把优化器状态放到 CPU 内存里，这是 FSDP 短期没有的对标功能。
四个常见误区
FSDP 就是 ZeRO-3
方向相同，但实现视角不同。FSDP 递归包装每个子模块，能在更细粒度上调度通信与计算；ZeRO 作为训练引擎，用统一的抽象管理参数分区，并额外提供 offload、checkpoint 等能力。
越省显存越快
分片降低了显存占用，但换来的是额外的 all-gather/reduce-scatter。在通信带宽足够时性能几乎无损，否则反而变慢。这就是为什么单机小模型不如 DDP。
张量并行总能省显存
它能把单个矩阵的存储切开，但同时要存储激活值和临时通信缓冲，实际显存收益不一定比 ZeRO 大。而且你必须为它付出沉重的代码修改成本。
这三个框架只能选一个
现实训练中经常同时用两三种。TP 处理单个权重矩阵，PP 处理层间，ZeRO/FSDP 处理跨节点的冗余状态，它们互补。
我的最终判断
如果你要快速跑通一个实验，而且愿意接受 PyTorch 原生 API，FSDP 是最低摩擦的选择。如果你在微软生态、需要 CPU offload 或者复杂 checkpoint 机制，DeepSpeed 更成熟。如果你要冲千亿模型，并且对每一层手写优化有足够耐心，Megatron-LM 给你的控制力最强。
但别把这三者当成对立的。今天你要做一个生产级训练系统，大概率是 FSDP 或 DeepSpeed 打底，再根据模型特征叠加 Megatron 的算子级并行。理解了它们各自的切分哲学，你就能在任何一个新框架出现时一眼看清它到底在做哪一层优化。
原创文章，作者：guanweilu，如若转载，请注明出处：https://guanweilu.cn/article/1370.html