DeepSpeed 与 ZeRO 详解

一句话介绍

• DeepSpeed：微软开发的深度学习优化库，主要用于训练大模型。
• ZeRO（Zero Redundancy Optimizer）：DeepSpeed 中的核心显存优化技术，分为 ZeRO-1、ZeRO-2、ZeRO-3 三个阶段。

为什么需要 ZeRO？

训练大模型（如 LLama 70B）时，GPU 显存经常不够用。除了模型参数本身，显存里还要存梯度、优化器状态（如 Adam 的动量和方差）。
传统分布式训练（如 DDP）会让每个 GPU 都保存完整模型、梯度、优化器状态 → 显存浪费严重。

ZeRO 的核心思想

分片（Sharding）：不再让每个 GPU 都存一份完整数据，而是将数据切开，分散到不同 GPU 上。使用时通过高速通信临时取用。

根据“切什么数据”，分为三个等级：

ZeRO 三阶段对比图

各阶段详解

ZeRO-1：只切分优化器状态

• 保留在每张 GPU 上的数据：完整的模型参数 + 完整的梯度
• 切分的数据：优化器状态（如 Adam 的一阶动量和二阶动量）
• 效果：显存占用约为原来的 1/4（以 Adam 优化器为例）
• 通信量：与普通 DDP 相同，没有额外开销
• 适用场景：模型刚刚超出单卡显存一点点，且希望通信开销最低

ZeRO-2：切分优化器状态 + 切分梯度

• 保留在每张 GPU 上的数据：完整的模型参数
• 切分的数据：优化器状态 + 梯度
• 效果：显存占用约为原来的 1/8
• 通信量：与普通 DDP 相同
• 适用场景：最推荐的通用场景，在显存节省和通信效率之间取得很好的平衡

ZeRO-3：切分优化器状态 + 切分梯度 + 切分模型参数

• 保留在每张 GPU 上的数据：无（所有三类数据都被切分）
• 切分的数据：优化器状态 + 梯度 + 模型参数
• 效果：显存占用随 GPU 数量线性下降（N 张卡 → 1/N 的显存）
• 通信量：显著增加（约为普通 DDP 的 1.5~2 倍）—— 用通信时间换取显存空间
• 适用场景：超大模型（如 70B+）或只有小显存 GPU 的环境

总结对比表