DeepSpeed ZeRO引擎深入

可视化演示

动态演示DeepSpeed ZeRO的核心机制，通过基础/核心/应用/扩展四种模式展示关键概念与交互流程。

显存四巨头

搬家队出发前，队长先把每个人背包里的东西清点了一遍。旧式搬家队每个工人身上有四样东西：要搬运的家具（模型参数）、贴在门上的搬迁贴纸（梯度）、写满户型经验的小本子（优化器状态），以及随时展开坐一下的折叠椅（激活值）。最重的是小本子，因为它记录了这家人所有物件的来历和摆放偏好。

训练大模型时，显存也是由这几类东西组成的。模型参数是网络的权重；梯度是反向传播时算出来的修正量；优化器状态则是 Adam 这类优化器用来记录历史信息的动量与方差，体积往往比参数本身还大好几倍。激活值是前向计算时产生的中间结果，虽然不常被提起，但也会在 moment 上占不少位置。

在混合精度训练下，我们设模型参数量为 \(\Psi\)。一个参数平均要消耗：FP16 参数 2 字节、FP16 梯度 2 字节、优化器状态 12 字节（包括 FP32 主权重 4 字节、Adam 动量 4 字节、Adam 方差 4 字节）。因此，每张 GPU 上的静态显存下限大约为：

所以，训练大模型时最该先处理的，不是家具本身，而是那本厚重的小本子。ZeRO 的切入点，正是从这里开始的。

数据并行之痛

如果搬家队继续沿用老规矩，每位工人都背一整套一模一样的图纸和小本子，那么 N 个工人身上就有 N 份完全相同的重复负重。更麻烦的是，大家经常要停下来互相核对同一本图纸，谁改了一个数字，所有人都得停下来重新对齐。

数据并行（DDP）就是这个老规矩。每张 GPU 上都放一份完整模型，每步算完梯度后，使用 all-reduce 把所有卡的梯度加起来再同步。模型规模不大时，这套方法简单有效；但模型一旦变大，问题立刻爆发：每卡要背 \(16\Psi\) 字节，N 张卡就等于：

这就是所谓的“显存墙”：原本想通过加卡省钱，结果每张卡都成了装不下完整模型的瓶颈。而且，卡数越多，梯度同步的通信量也越大，训练吞吐反而可能下降。队长看到这里明白了：必须打破“人人全量”的老规矩，把重复的负重真正拆开。

ZeRO三阶段

队长开始重新分工。ZeRO 把显存里的三大巨头——优化器状态、梯度、参数——按 GPU 数量 \(N\) 切成 \(N\) 份，每个工人只保管其中一份。第一阶段先切最重的优化器状态，第二阶段连同梯度一起切，第三阶段连参数也切。每多切一类，工人背包就轻一大截。

从技术上看，三个阶段分别对应：

• ZeRO-Stage 1：只切分优化器状态，每卡仍需保存完整参数和梯度，显存约 \(4\Psi + \frac{12\Psi}{N}\)。

• ZeRO-Stage 2：切分优化器状态和梯度，每卡只需一份完整参数，显存约 \(2\Psi + \frac{14\Psi}{N}\)。

• ZeRO-Stage 3：三者全部切分，每卡只持有全模型的 \(1/N\)，显存约：

显存省下来了，代价是通信变多。前向计算需要临时收集被切分的参数，反向传播需要对梯度做 reduce-scatter，参数更新后还要再次广播。阶段越高，每个工人的背包越轻，但互相喊话的次数也越多。因此，实际选择哪一阶段，要看集群网络是否足够快。

混合精度

搬家队为了让背包更轻，决定把实木家具换成轻便的折叠家具（FP16），但为了完整保留每件家具的历史记录，队长还让每个工人保留一本完整的登记簿（FP32 主权重）。家具轻了，账本依旧完整，这样搬家既省力又不会弄丢关键信息。

在混合精度训练中，前向传播和反向传播都使用 FP16 参数与激活值，这样计算更快、显存更省。但优化器更新参数时，必须使用 FP32 主权重，否则精度会快速恶化。FP16 的数值范围较小，梯度在反向传播中很容易变小到“下溢”成 0，为此需要先放大损失：

与 ZeRO 配合时，FP32 主权重和 Adam 状态加起来正好是显存四巨头中的“优化器状态”。ZeRO 把这部分切分到各卡，既保留了混合精度的精度优势，又进一步吃到了显存红利。折叠家具压低体积，登记簿又交给专人分管，两者配合起来效果远超单用其中任意一种。

卸载策略

就算把背包切分到极致，有些工人还是装不下。队长于是启动终极方案：把不常用的“小本子”寄放到公司仓库（CPU 内存甚至 NVMe 磁盘），工人只带当天需要用的几页纸，用完了再回去取。代价是取件时间变长，但背包确实轻了。

这就是 ZeRO-Offload 的思路。它把 Adam 优化器状态甚至梯度、参数从 GPU 显存搬到 CPU 内存或 NVMe 磁盘上，GPU 只保留当前计算必需的碎片。由于 CPU 内存价格远低于 GPU 显存，这让小显存机器也能挑战大模型。但 PCIe 传输速度比 GPU 显存慢得多，每一步训练的总耗时受制于最慢的环节：

所以，卸载策略本质上是“用时间换空间”。在单机多卡或显存非常吃紧的场景下，它能让训练跑起来；但如果网络或 PCIe 带宽不够，训练吞吐会明显下降。仓库能借，但别把喊话通道塞满。

通信优化

工人之间互相借东西时，如果每借一块木板就喊一嗓子，整栋楼都会吵翻天。队长规定：把几十个小物件先装进一个大箱子，再统一喊话搬运。这样一来，喊话次数从几十次变成一次，通讯效率大大提高。

深度学习反向传播会生成成百上千个小梯度张量。如果每个梯度都单独走一次 all-reduce，通信的固定开销会淹没计算本身。通信总耗时可以写成：

DeepSpeed 通过把梯度放进一个个“桶”（bucket），攒到一定大小后再发起通信，从而大幅减少 \(n\)。配置项 reduce_bucket_size 和 allgather_bucket_size 控制桶的大小，overlap_comm 则让通信与反向计算重叠，把等待时间“藏”进计算里。喊话次数少了，搬家速度自然就快了。

调优与排查

一切准备就绪后，队长要检查三件事：背包有没有超限、喊话频道堵不堵、仓库有没有拖后腿。训练调优也是一样，先判断瓶颈到底在显存、通信还是计算，再针对性下药。

调优时，显存不足可以用 zero_optimization.stage 逐步升级，从 1 到 3；仍然不够就开 offload_optimizer 或激活重计算。吞吐下降时，优先确认通信桶大小、NCCL 带宽利用率，再调整 batch size 和梯度累积步数。每一步训练的总时间，大致由三部分组成：

实际操作中，可以先用 nvidia-smi 看显存峰值，用 ds_report 检查 DeepSpeed 版本，再在日志里观察 throughput 与 latency。如果切到更高 ZeRO 阶段后吞吐反而下降，优先怀疑通信瓶颈，而不是盲目加卡。搬家队的目标不是让每班搬得最重，而是让总时间最短：算、传、等三端平衡。

AI 应用场景

Python 演示

import numpy as np def make_params(sizes):    params = {}    for i, size in enumerate(sizes):        params['w' + str(i)] = np.random.randn(size).astype(np.float32)    return params class ZeROStage1:    def __init__(self, rank_gpu, world_size_gpu, params):        self.rank_gpu = rank_gpu        self.world_size_gpu = world_size_gpu        self.params = params        self.grads = {name: np.zeros_like(p) for name, p in params.items()}        self.local_names = list(params.keys())[rank_gpu::world_size_gpu]        self.moments = {}        self.velocities = {}        for name in self.local_names:            self.moments[name] = np.zeros_like(params[name])            self.velocities[name] = np.zeros_like(params[name])        self.memory = self.estimate_memory()     def estimate_memory(self):        total = 0        for name, p in self.params.items():            total += p.nbytes * 2        for name in self.local_names:            total += self.params[name].nbytes * 4        return total     def set_grad(self, name, grad):        self.grads[name] += grad     def all_reduce_grad(self):        for name in self.grads:            self.grads[name] /= self.world_size_gpu     def optimizer_step(self, lr=0.01, beta1=0.9, beta2=0.999):        for name in self.local_names:            g = self.grads[name]            self.moments[name] = beta1 * self.moments[name] + (1 - beta1) * g            self.velocities[name] = beta2 * self.velocities[name] + (1 - beta2) * g * g            self.params[name] -= lr * self.moments[name] / (np.sqrt(self.velocities[name]) + 1e-8) if __name__ == '__main__':    world_size_gpu = 4    params = make_params([1024, 2048, 4096])    workers = [ZeROStage1(rank, world_size_gpu, params) for rank in range(world_size_gpu)]    for worker in workers:        for name in params:            worker.set_grad(name, np.ones_like(params[name]))        worker.all_reduce_grad()        worker.optimizer_step()    print('每个GPU只保存本卡分片的优化器状态:', [w.local_names for w in workers])

1. self.local_names = list(params.keys())[rank_gpu::world_size_gpu]按 GPU 编号切片，决定本卡负责哪些参数的优化器状态。
2. total += p.nbytes * 2每个参数仍保留 FP16 参数和 FP16 梯度的全量副本。
3. total += self.params[name].nbytes * 4只对本卡负责的参数额外加上主权重与 Adam 状态。
4. self.grads[name] /= self.world_size_gpu模拟 all-reduce 后的平均梯度。
5. self.moments[name] = beta1 * self.moments[name] + (1 - beta1) * gAdam 第一动量更新，与 DeepSpeed 内部计算一致。
6. self.velocities[name] = beta2 * self.velocities[name] + (1 - beta2) * g * g第二动量更新，用于缩放参数修正量。
7. self.params[name] -= lr * self.moments[name] / (np.sqrt(self.velocities[name]) + 1e-8)使用分片优化器状态更新参数。

常见误区

面试题