分布式训练：DeepSpeed 与 Zero 数据并行

大模型分布式训练：DeepSpeed Zero 优化深度解读

大模型时代来了，模型参数动辄百亿千亿，训练它们可不再是BERT时代那种小打小闹。显存和计算量双双爆炸，带来两大核心挑战：

• 显存的挑战
  ：以A100 80G显存为例，一个175B的GPT-3模型，光参数就要占175B×4bytes≈700GB的空间。加载都费劲，更别说训练了。

• 计算的挑战
  ：175B模型的计算量，叠加海量预训练数据，和BERT时代完全不是一个量级。

这篇文章聚焦大模型分布式训练的核心方案，重点拆解以数据并行为基础的DeepSpeed Zero优化。先梳理一下整体脉络。

三大并行策略

并行技术宏观上可以分为三类：

、和。

1. 数据并行（DP）

：每个GPU worker都持有一份完整的模型副本，各自拿到的数据，独立计算梯度。最后所有worker把梯度汇总求平均，再同步更新模型权重。

：8张卡，每张卡batch_size=4，那么global batch size=32。每张卡用4个样本算梯度，然后计算平均梯度，每张卡同时更新参数。

在小模型时代（比如BERT）数据并行是最常用的方法——模型不大，先上数据并行就对了。

2. 流水线并行与张量并行

模型大到一张卡装不下的时候，就要考虑

和了。两者的本质区别在于不同。

2.1 流水线并行

：把模型的不同层放到不同GPU上。比如一个6层Transformer，前3层放GPU0，后3层放GPU1。如下图所示：

：每个GPU处理一个batch时，必须等前一个GPU完成才能开始，这就产生了——GPU闲置等待。为了缓解这个问题，通常引入，让不同batch的计算交错进行，尽量让气泡变小。

2.2 张量并行

：在层内做切分，把Tensor矩阵运算分散到多个GPU上。下图就是把每一层切成两块：

：存储效率高。

：每次前向和反向计算都会引入额外的通信开销。

DeepSpeed 与 Zero

什么是Deepspeed？

Deepspeed现在已经是训练大模型的标配了。它是微软发布的大规模分布式训练工具，专为大模型设计。而DeepSpeed Zero是它的核心组件，本质上是一种

。

Deepspeed Zero 的核心思想

在

阶段，Zero把训练中占显存的参数分为两部分：

• 模型状态
  ：包括模型参数（fp16/bf16，占2A）、模型梯度（fp16/bf16，2A）、Adam状态（fp32的模型参数备份，fp32的momentum和variance，共4A+4A+4A）。假设模型参数量为A，模型状态总共需要2A+2A+12A=16A字节。一个7B模型，模型状态就要112GB。

• 剩余状态
  ：包括激活值（activation）、临时缓冲区（buffer）和显存碎片。

Zero的优化思想直击要害：针对

做——每张卡只存一部分模型状态，N张卡，每张卡存1/N。但注意，Zero本质上还是数据并行方案，它的基座依然是数据并行。下面这个结论很有用：

• 速度
  ：Zero 1 ＞ Zero 2 ＞ Zero 2 + offload ＞ Zero 3 ＞ Zero 3 + offload

• 显存
  ：阶段0（DDP）＜ 阶段1 ＜ 阶段2 ＜ 阶段2+卸载 ＜ 阶段3 ＜ 阶段3+卸载

DeepSpeed 的优化细节

1. DeepSpeed Zero Stage 1

针对

进行分片。此时每张卡的模型状态显存变成：。以单机8卡、7B模型为例，模型状态部分显存从112GB降到 4×7 + (12×7)/8 = 38.5GB。

2. DeepSpeed Zero Stage 2

在Zero-1基础上，再对

分片。此时每张卡所需显存公式变为：（应与上面类似）。同样以单机8卡7B为例，模型状态部分显存降到 2×7 + (14×7)/8 = 26.25GB。

3. DeepSpeed Zero Stage 3

在Zero-2的基础上，连

也分片。每张卡模型状态显存变为：16×7/8 = 14GB。

4. Zero-offload

核心思想：

。在Zero-2的基础上，把Adam状态和梯度迁移到CPU内存。一般情况下不推荐，因为速度会明显下降。

一点补充

注意，DeepSpeed优化的是模型参数、梯度、Adam状态所占的显存。但正如之前提到的，激活值同样占用巨大显存，此外还有显存碎片。以Qwen 1.5为例，per_gpu_batch_size=1、seq_length=2048、精度bf16时，所需显存约为28.5GB（粗略估值）。即使使用DeepSpeed Zero 3（非常慢），也需要28.5+14=42.5GB。这意味着在大模型时代，如果不引入流水线并行和张量并行，V100-32G和A100-40G训练7B模型都会非常吃力。

最后

大多数情况下，我们优先选择DeepSpeed Zero 2，主要考虑训练效率。DeepSpeed是大模型时代必备工具，值得花时间熟练掌握。目前开源模型迭代飞快，直接使用成熟的微调框架（如LLaMA-Factory）会更高效，不必自己从头造轮子。