AI陪读ZeRO：万卡训练的第一块基石

问题：显存为什么不够用

训练大模型，显存是第一个拦路虎。一个 1000 亿参数（100B）的模型，单是模型参数就要 400GB（每个参数 4 字节），还没算梯度、优化器状态、中间激活值。

光 Optimizer States（Adam 的两个动量）就要 800GB——比整张 A100 80GB 的卡还多 10 倍。

100B 参数模型显存去哪了

• 模型参数：~40GB
• 梯度：~40GB
• Optimizer states：~80GB
• 激活值：~100GB+
• 总计：~260GB+

80GB 的 A100？单卡跑不起来。

原来的解法：各有代价

Data Parallel（数据并行）

每张卡复制一份完整模型，各自算不同数据的梯度，最后汇总。简单，但每张卡都存全量——显存浪费，等于没解决根本问题。

Model Parallel（模型并行）

把模型拆开放到不同卡上。论文里测了 Megatron-LM 跑 40B 参数跨两个 DGX-2 节点，每张 V100 卡只跑出 5 Tflops——不到硬件峰值的 5%，效率极低。原因：跨节点通信带宽不够，层间通信等待时间太长。

ZeRO 核心方案

Zero Redundancy Optimizer——"零冗余"的意思：每张卡只存 1/N 的数据，不重复，消除显存冗余，用通信换显存。分三个阶段，递进增强：

三个阶段对比

ZeRO-1：只分片 Optimizer States → 显存减少 4 倍，通信几乎不变

ZeRO-2：+ 分片 Gradients → 显存减少 8 倍，通信几乎不变

ZeRO-3：+ 分片 Parameters → 线性减少（64卡=64倍），通信+50%

ZeRO-3 的关键意义

论文原文计算：万亿参数模型用 Adam 混合精度训练，optimizer states + gradients + parameters ≈ 16TB。16TB ÷ 1024 张 GPU = 每卡 16GB，A100 32GB 完全够用。

结论：ZeRO 让万亿参数在 1024 张卡上成为可能。

实测数据（论文原文）

• 400 张 GPU 训练 100B 参数，吞吐量 15 Petaflops
• 相比当时最好方案：模型规模 8 倍，性能 10 倍
• ZeRO-2 训 13B 参数模型，不需要模型并行
• 最终成果：Turing-NLG，170 亿参数，当时世界最大语言模型

三个核心记忆点

1. 三个阶段按需开

显存不够先 ZeRO-1，还不够 ZeRO-2，再不够才 ZeRO-3。通信代价逐级增加，不要一上来就开满。

2. ZeRO-3 最强但通信最重

适合机器间 NVLink 带宽充足的环境（单节点内）。跨节点通信带宽受限的场景，用 ZeRO-1 或 ZeRO-2 更稳。

3. ZeRO = DeepSpeed

微软把 ZeRO 做进了 DeepSpeed 框架，现在工业界用的就是 DeepSpeed ZeRO。这是万卡训练的事实标准。

ZeRO 在整个训练栈里的位置

万卡训练是个系统工程，ZeRO 只是显存优化这一层：

硬件层（NVLink / NCCL）

↓

ZeRO — 节省 4-64 倍显存

↓

LAMB — 万卡大 Batch 稳定收敛

↓

框架（DeepSpeed / Megatron）

ZeRO 让万卡能跑起来，LAMB 让万卡跑得稳。两个合在一起，才是真正可用的万卡训练方案。

"

类比：Data Parallel 像每个学生都复印一本教材，浪费纸（显存）。ZeRO 像把一本教材拆成 N 份，每人拿 1/N，既完成了学习（训练），又省了资源——最后通过交换笔记（通信）保证每个人都学到了全部内容。

💡 核心洞察

ZeRO-1/2/3 是递进关系，通信换显存，资源有限时按需选择。不要一上来开 ZeRO-3，先 ZeRO-1，

看显存够不够，再逐步升级。

· · ·

— END —