深度学习专题：模型训练的数据并行（三）

深度学习专题：模型训练的数据并行（三）

ZeRO 系列优化：ZeRO-1、ZeRO-2、ZeRO-3、ZeRO-R

（一）Adam 优化器

传统的 SGD 优化器第 t+1 轮参数更新公式为：
$$w_{t+1}=w_t-lr\cdot g_{t+1}$$
其中 $w_t$ 为第 t 轮参数，$lr$ 为学习率（不变的定值），$g_t$ 为第 t 轮参数的梯度。

这种优化策略容易遇到驻点问题：当模型参数的梯度接近于零（即 $g_{t+1}$ 接近于零）时，参数更新会非常缓慢，甚至停滞不前。

Adam 优化器为了解决驻点问题，在 SGD 优化器的基础上，引入了动量和自适应学习率的概念。其第 t+1 轮参数更新公式为：
$$w_{t+1}=w_t-lr\cdot \frac{m_{t+1}}{\sqrt{v_{t+1}}+ϵ}$$
其中 $m_t$ 为第 t 轮的动量，$v_t$ 为第 t 轮的方差值，$ϵ$ 为一个小的常量（防止除零错误）。特别地，动量与方差的计算公式为：
$$\begin{gathered} m_{t+1}=\frac{{\beta }_1{m}_t+(1-{\beta }_1)g_{t+1}}{1-{\beta }_1^{t+1}} \\ {v}_{t+1}=\frac{{\beta }_2{v}_t+(1-{\beta }_2)g_{t+1}^2}{1-{\beta }_2^{t+1}} \end{gathered}$$
其中 ${\beta }_1$ 和 ${\beta }_2$ 为两个超参数（常设为 0.9 和 0.999）。初值为 $m_0=0$ 和 $v_0=0$。

（二）Adam 优化器数据并行存储代价

按照 Adam 优化器的公式，模型在训练过程中不仅要存储当前轮的参数 $w_t$ 和梯度 $g_t$，还要存储当前轮的动量 $m_t$ 和方差 $v_t$。假设有 N 个参数，在数据并行场景下，与 SGD 优化器比较分析第 t+1 轮 epoch 的反向传播过程：

1. SGD 优化器

SGD 优化器在每个 GPU 中需要额外存储 $2N$ 个浮点数。

第一步：第 k 个 GPU 读取分配给自己的训练数据，读取第 t 轮 epoch 的参数向量 $w_t$ 并进行前向传播、损失计算和反向传播，得到第 t+1 轮 epoch 的梯度向量

第二步：基于 Ring All-Reduce 操作获取其他所有 GPU 上的第 t+1 轮 epoch 的梯度向量，计算梯度向量均值 $g_{t+1}$ 并存储；

第三步：根据聚合后的梯度向量 $g_{t+1}$，更新参数向量 $w_{t+1}$。

2. Adam 优化器

Adam 优化器在每个 GPU 中需要额外存储 $4N$ 个浮点数。

第一步：第 k 个 GPU 读取分配给自己的训练数据，读取第 t 轮 epoch 的参数向量 $w_t$ 并进行前向传播、损失计算和反向传播，得到第 t+1 轮 epoch 的梯度向量；

第二步：基于 Ring All-Reduce 操作获取其他所有 GPU 上的第 t+1 轮 epoch 的梯度向量，计算梯度向量均值 $g_{t+1}$ 并存储；

第三步：读取第 t 轮 epoch 的动量 $m_t$ 和方差 $v_t$ 并根据聚合后的第 t+1 轮 epoch 的梯度向量 $g_{t+1}$ 计算第 t+1 轮 epoch 的动量 $m_{t+1}$ 和方差 $v_{t+1}$ 并存储；

第四步：根据动量 $m_{t+1}$ 和方差 $v_{t+1}$，更新参数向量 $w_{t+1}$。

3. 分析与总结

GPU 需要存储第 t 轮 epoch 的参数向量 $w_t$，因为在第 t+1 轮训练过程中需要根据 $w_t$ 进行前向传播、损失计算和反向传播，从而计算出第 t+1 轮 epoch 的梯度向量 $g_{t+1}$；

GPU 还需要存储第 t 轮 epoch 的动量向量 $m_t$ 和方差向量 $v_t$，因为在第 t+1 轮训练过程中需要根据 $m_t$ 和 $v_t$ 计算出第 t+1 轮 epoch 的动量向量 $m_{t+1}$ 和方差向量 $v_{t+1}$；

上面的两个存储项均涉及到跨 epoch 存储，所以对应的空间在整个训练的生命周期都不能被释放。

GPU 还需要临时存储第 t+1 轮 epoch 的梯度向量 $g_{t+1}$，因为 GPU 需要使用 Ring All-Reduce 算法在所有 GPU 之间归约梯度向量并计算均值。如果是 SGD 优化器，这个梯度向量会直接用于更新参数；如果是 Adam 优化器，这个梯度向量会用于更新动量和方差。不同的是，完成对应的更新操作后，梯度向量就不需要继续存储了，对应的空间就可以被释放了。所以梯度向量只需要临时存储，而不需要跨 epoch 存储。

（三）ZeRO-1 优化原理

为了减少 Adam 优化器的存储代价，ZeRO-1 提出了一种优化策略：​​将优化器状态在多个 GPU 之间进行分区存储​​，而不是在每个 GPU 上完整保存所有优化器状态。

在数据并行训练中，每个 GPU 通常保存完整的模型参数、梯度和优化器状态。ZeRO-1 观察到，优化器状态（如 Adam 中的动量和方差）在训练过程中主要参与参数更新计算，但不同参数的优化器状态之间是相互独立的。因此，可以将优化器状态按照参数维度进行分区，每个 GPU 只需要存储自己负责的那部分动量向量 $m_{t+1}$ 和方差向量 $v_{t+1}$ 并更新对应的参数，然后通过 GPU 的 All-Gather 算法同步所有更新后的参数。

1. 步骤分析

准备步骤：假设模型有 N 个参数，有 P 个 GPU 负责并行训练

• 每个 GPU 负责管理 $M=\lceil N/P\rceil$ 个参数
• 第 k 个 GPU 将负责第 $s_k=1+(k-1)M$ 到 $e_k=kM$ 个参数

第一步：第 k 个 GPU 读取分配给自己的训练数据，读取第 t 轮 epoch 的参数向量 $w_t$ 并进行前向传播、损失计算和反向传播，得到第 t+1 轮 epoch 的梯度向量；

第二步：基于 Ring All-Reduce 操作获取其他所有 GPU 上的第 t+1 轮 epoch 的梯度向量，计算梯度向量均值 $g_{t+1}$ 并存储；

第三步：根据聚合后的梯度向量 $g_{t+1}$，更新自己负责的那一部分参数的动量和方差；

第四步：每个 GPU 根据自己负责的那部分更新后的动量和方差，更新自己负责的那部分参数向量，然后释放梯度向量 $g_{t+1}$ 占用的存储空间；

第五步：每个 GPU 与其他所有 GPU 通过 All-Gather 算法通信，得到更新后的参数向量 $w_{t+1}$。

2. 存储优化分析

• 传统数据并行：每个 GPU 需要存储 $4N$ 个浮点数；
• ZeRO-1 优化：每个 GPU 负责管理 $M=\lceil N/P\rceil$ 个参数，所以每个 GPU 需要存储 $2M+2N$ 个浮点数。当 $M\ll N$ 时，存储量趋近于 $2N$，存储代价可以显著减少。

（四）ZeRO-2 优化原理

ZeRO-2 在 ZeRO-1 的基础上进一步优化，提出了​​梯度分区存储​​的策略，其核心通信模式是 ​​Reduce-Scatter​​ 而不是 All-Reduce。

在传统数据并行中，梯度同步使用 All-Reduce，每个 GPU 最终都获得完整的平均梯度。ZeRO-2 改为使用 Reduce-Scatter，每个 GPU 只获得部分梯度的聚合结果，与优化器状态的分区相对应。

1. 步骤分析

准备步骤：假设模型有 N 个参数，有 P 个 GPU 负责并行训练

• 每个 GPU 负责管理 $M=\lceil N/P\rceil$ 个参数
• 第 k 个 GPU 将负责第 $s_k=1+(k-1)M$ 到 $e_k=kM$ 个参数

第一步：第 k 个 GPU 读取分配给自己的训练数据，读取第 t 轮 epoch 的参数向量 $w_t$ 并进行前向传播、损失计算和反向传播，得到第 t+1 轮 epoch 的梯度向量；

第二步：基于 Reduce-Scatter 操作获取其他所有 GPU 上的第 t+1 轮 epoch 的梯度向量的第 $s_k$ 到 $e_k$ 个元素，计算梯度向量均值 $g_{t+1}$ 的第 $s_k$ 到 $e_k$ 个元素并存储，然后释放梯度向量均值 $g_{t+1}$ 的 $s_k$ 到 $e_k$ 之外的元素占用的存储空间；

第三步：根据聚合后的梯度向量 $g_{t+1}$ 的 $s_k$ 到 $e_k$ 个元素，更新自己负责的那一部分参数的动量和方差；

第四步：每个 GPU 根据自己负责的那部分更新后的动量和方差，更新自己负责的那部分参数向量，然后释放梯度向量 $g_{t+1}$ 的第 $s_k$ 到 $e_k$ 的元素占用的存储空间；

第五步：每个 GPU 与其他所有 GPU 通过 All-Gather 算法通信，得到更新后的参数向量 $w_{t+1}$。

2. 存储优化分析

• ZeRO-1 优化：每个 GPU 需要存储 $2M+2N$ 个浮点数，当 $M\ll N$ 时，存储量趋近于 $2N$，存储代价可以显著减少；
• ZeRO-2 优化：每个 GPU 需要存储 $3M+N$ 个浮点数，当 $M\ll N$ 时，存储量趋近于 $N$，存储代价可以显著减少。

（五）ZeRO-3 优化原理

ZeRO-3 在 ZeRO-2 的基础上进一步优化，提出了​​参数分区存储​​的策略，每个 GPU 只存储自己负责的那部分参数，这需要在训练的前向传播和反向传播之前通过 All-Gather 通信收集全部参数，从而达成一种通信换存储的优化效果。

1. 步骤分析

准备步骤：假设模型有 N 个参数，有 P 个 GPU 负责并行训练

• 每个 GPU 负责管理 $M=\lceil N/P\rceil$ 个参数
• 第 k 个 GPU 将负责第 $s_k=1+(k-1)M$ 到 $e_k=kM$ 个参数，且只存储这些参数。

第一步：第 k 个 GPU 通过 All-Gather 通信获得第 t 轮 epoch 的参数向量 $w_t$，然后读取分配给自己的训练数据，完成前向传播，计算得到损失值；

第二步：进行反向传播，计算得到第 t+1 轮 epoch 的梯度向量，然后释放第 t 轮参数向量的 $s_k$ 到 $e_k$ 之外的元素占用的存储空间；

第三步：基于 Reduce-Scatter 操作获取其他所有 GPU 上的第 t+1 轮 epoch 的梯度向量的第 $s_k$ 到 $e_k$ 个元素，计算梯度向量均值 $g_{t+1}$ 的第 $s_k$ 到 $e_k$ 个元素并存储，然后释放梯度向量均值 $g_{t+1}$ 的 $s_k$ 到 $e_k$ 之外的元素占用的存储空间；

第四步：根据聚合后的梯度向量 $g_{t+1}$ 的 $s_k$ 到 $e_k$ 个元素，更新自己负责的那一部分参数的动量和方差；

第五步：每个 GPU 根据自己负责的那部分更新后的动量和方差，更新自己负责的那部分参数向量，然后释放梯度向量 $g_{t+1}$ 的第 $s_k$ 到 $e_k$ 的元素占用的存储空间；

2. 存储优化分析

• ZeRO-3 优化：每个 GPU 需要存储 $4M$ 个浮点数，当 $M\ll N$ 时，存储量趋近于 $0$，存储代价可以显著减少。

（四）ZeRO-3 层间流水线优化

由于层间流水线优化可以将多个层的前向传播和反向传播并行执行，而每个层只需要存储自己负责的那部分参数，因此 ZeRO-3 与层间流水线优化结合可以进一步减少存储代价。

1. 前向传播和反向传播的计算资料

考虑下面的模型：

• 输入输出都是标量，分别记作 $x$ 和 $y$；
• 模型有 4 个层，其中前三个层是全连接层加上 ReLU 激活函数，最后一个层是全连接层；
• 每个层的参数数量都是 2 个，分别是 权重 $w$ 和 偏置 $b$；
• 模型的前向传播：
  • 第 1 层：$h_1=\text{ReLU}(w_1\times x+b_1)$
  • 第 2 层：$h_2=\text{ReLU}(w_2\times h_1+b_2)$
  • 第 3 层：$h_3=\text{ReLU}(w_3\times h_2+b_3)$
  • 第 4 层：$y=w_4\times h_3+b_4$
  • 损失函数：$L=\frac{1}{2}(y-x{)}^2$
• 模型的反向传播：

  • 输出层：

    • $\frac{\partial L}{\partial y}=y-x$

  • 第 4 层：

    • $\frac{\partial L}{\partial w_4}=\frac{\partial L}{\partial y}\times \frac{\partial y}{\partial w_4}=(y-x)\times h_3$
    • $\frac{\partial L}{\partial b_4}=\frac{\partial L}{\partial y}\times \frac{\partial y}{\partial b_4}=y-x$

  • 第 3 层：

    • $\frac{\partial L}{\partial h_3}=\frac{\partial L}{\partial y}\times \frac{\partial y}{\partial h_3}=(y-x)\times w_4$
    • $\frac{\partial L}{\partial w_3}=\frac{\partial L}{\partial h_3}\times \frac{\partial h_3}{\partial w_3}=\frac{\partial L}{\partial h_3}\times h_2\times \mathbb{I}(h_3>0)$
    • $\frac{\partial L}{\partial b_3}=\frac{\partial L}{\partial h_3}\times \frac{\partial h_3}{\partial b_3}=\frac{\partial L}{\partial h_3}\times \mathbb{I}(h_3>0)$

  • 第 2 层：

    • $\frac{\partial L}{\partial h_2}=\frac{\partial L}{\partial h_3}\times \frac{\partial h_3}{\partial h_2}=\frac{\partial L}{\partial h_3}\times w_3\times \mathbb{I}(h_3>0)$

    • $\frac{\partial L}{\partial w_2}=\frac{\partial L}{\partial h_2}\times \frac{\partial h_2}{\partial w_2}=\frac{\partial L}{\partial h_2}\times h_1\times \mathbb{I}(h_2>0)$

    • $\frac{\partial L}{\partial b_2}=\frac{\partial L}{\partial h_2}\times \frac{\partial h_2}{\partial b_2}=\frac{\partial L}{\partial h_2}\times \mathbb{I}(h_2>0)$

  • 第 1 层：

    • $\frac{\partial L}{\partial h_1}=\frac{\partial L}{\partial h_2}\times \frac{\partial h_2}{\partial h_1}=\frac{\partial L}{\partial h_2}\times w_2\times \mathbb{I}(h_2>0)$

    • $\frac{\partial L}{\partial w_1}=\frac{\partial L}{\partial h_1}\times \frac{\partial h_1}{\partial w_1}=\frac{\partial L}{\partial h_1}\times x\times \mathbb{I}(h_1>0)$

    • $\frac{\partial L}{\partial b_1}=\frac{\partial L}{\partial h_1}\times \frac{\partial h_1}{\partial b_1}=\frac{\partial L}{\partial h_1}\times \mathbb{I}(h_1>0)$

分析可知：

（1）第 m 层的前向传播，需要的计算资料有：

• $w_m$
• $b_m$
• $h_{m-1}$

其中，$h_{m-1}$ 是前一层的输出，由前一层计算结束后保存；$w_m$ 和 $b_m$ 是本层的参数，如果参数是分区存储的，那么参数需要通过 All-Gather 通信获取。

（2）第 m 层的反向传播，需要的计算资料有：

• $\frac{\partial L}{\partial h_{m+1}}$
• $w_{m+1}$
• $h_{m-1}$

其中，$\frac{\partial L}{\partial h_{m+1}}$ 是后一层的梯度，由后一层计算结束后保存；$h_{m-1}$ 是前一层的输出，由前一层计算结束后保存；$w_{m+1}$ 是下一层的参数，需要通过 All-Gather 通信获取。

2. 前向传播和反向传播的层间流水线并行执行

基于上面的讨论，我们可以让前向传播和反向传播进行层间流水线并行执行，即通信过程和计算过程重叠。

假设有 4 块 GPU，每块 GPU 负责存储 1 个层的参数，每个层的参数数量都是 2 个；此外每块 GPU 都负责 1/4 的 mini batch size 的训练数据进行训练。

前向传播重叠过程：

• 第一步：第 1 层参数的广播通信。由于第 1 层参数存储在 GPU-1 上，因此基于 All-Gather 通信将第 1 层参数广播到所有 GPU 上。
• 第二步：第 1 层前向传播的计算，第 2 层参数的广播通信。同时完成这两件事：①在每个 GPU 上计算第 1 层的前向传播，并保存结果 $h_1$，然后删除上一步广播到的参数（除 GPU-1 外）；②由于第 2 层参数存储在 GPU-2 上，因此基于 All-Gather 通信将第 2 层参数广播到所有 GPU 上。
• 第三步：第 2 层前向传播的计算，第 3 层参数的广播通信。同时完成这两件事：①在每个 GPU 上计算第 2 层的前向传播，并保存结果 $h_2$，然后删除上一步广播到的参数（除 GPU-2 外）；②由于第 3 层参数存储在 GPU-3 上，因此基于 All-Gather 通信将第 3 层参数广播到所有 GPU 上。
• 第四步：第 3 层前向传播的计算，第 4 层参数的广播通信。同时完成这两件事：①在每个 GPU 上计算第 3 层的前向传播，并保存结果 $h_3$，然后删除上一步广播到的参数（除 GPU-3 外）；②由于第 4 层参数存储在 GPU-4 上，因此基于 All-Gather 通信将第 4 层参数广播到所有 GPU 上。
• 第五步：第 4 层前向传播的计算。在每个 GPU 上计算第 4 层的前向传播，并保存结果 $y$，然后删除上一步广播到的参数（除 GPU-4 外）（删除操作可选）。

反向传播重叠过程：

• 第一步：输出层反向传播的计算，第 4 层反向传播的计算，第 4 层参数的广播通信。同时完成这两件事：①在每个 GPU 上计算输出层和第 4 层的反向传播，并保存结果 $\frac{\partial L}{\partial y}$，$\frac{\partial L}{\partial w_4}$，$\frac{\partial L}{\partial b_4}$；②由于第 4 层参数存储在 GPU-4 上，因此基于 All-Gather 通信将第 4 层参数广播到所有 GPU 上（如果前向传播时没有进行删除操作，则此操作省略）。
• 第二步：第 3 层反向传播的计算，第 3 层参数的广播通信，第 4 层参数的更新。同时完成这三件事：①在每个 GPU 上计算第 3 层的反向传播，并保存结果 $\frac{\partial L}{\partial h_3}$，$\frac{\partial L}{\partial w_3}$，$\frac{\partial L}{\partial b_3}$，然后删除上一步广播到的参数（除 GPU-4 外）；②由于第 3 层参数存储在 GPU-3 上，因此基于 All-Gather 通信将第 3 层参数广播到所有 GPU 上；③基于 Reduce-Scatter 操作将第 4 层梯度收集到 GPU-4 上，然后计算第 4 层优化器状态，更新第 4 层参数。
• 第三步：第 2 层反向传播的计算，第 2 层参数的广播通信，第 3 层参数的更新。同时完成这三件事：①在每个 GPU 上计算第 2 层的反向传播，并保存结果 $\frac{\partial L}{\partial h_2}$，$\frac{\partial L}{\partial w_2}$，$\frac{\partial L}{\partial b_2}$，然后删除上一步广播到的参数（除 GPU-3 外）；②由于第 2 层参数存储在 GPU-2 上，因此基于 All-Gather 通信将第 2 层参数广播到所有 GPU 上；③基于 Reduce-Scatter 操作将第 3 层梯度收集到 GPU-3 上，然后计算第 3 层优化器状态，更新第 3 层参数。
• 第四步：第 1 层反向传播的计算，第 2 层参数的更新：在每个 GPU 上计算第 1 层的反向传播，并保存结果 $\frac{\partial L}{\partial h_1}$，$\frac{\partial L}{\partial w_1}$，$\frac{\partial L}{\partial b_1}$，然后删除上一步广播到的参数（除 GPU-2 外）；基于 Reduce-Scatter 操作将第 2 层梯度收集到 GPU-2 上，然后计算第 2 层优化器状态，更新第 2 层参数。
• 第五步：第 1 层参数的更新：在每个 GPU 上计算第 1 层优化器状态，更新第 1 层参数。

（五）ZeRO-R 优化

ZeRO-R（ZeRO-Redundancy）主要针对非参数存储进行优化，包括激活值、临时缓冲区和内存碎片等问题。

1. 激活值存储优化

在前向传播中，需要保存中间激活值用于反向传播，对于大模型和大 batch size，激活值显存占用可能过大。

传统方法：每个GPU保存完整激活值

• GPU-0: [a1, a2, a3, a4]
• GPU-1: [a1, a2, a3, a4]
• GPU-2: [a1, a2, a3, a4]
• GPU-3: [a1, a2, a3, a4]

ZeRO-R：激活值分区存储

• GPU-0: [a1]
• GPU-1: [a2]
• GPU-2: [a3]
• GPU-3: [a4]

2. 临时缓冲区优化

All-Reduce、All-Gather等操作需要临时缓冲区，这些缓冲区可能很大，且需要频繁分配释放。

传统方法：每次操作动态分配缓冲区

ZeRO-R：

• 固定大小的持久化缓冲区
• 每次操作复用这块缓冲区，避免频繁分配释放

3. 内存碎片优化

由于张量的频繁分配和释放，GPU显存会出现碎片化，导致即使总空闲内存足够，也无法分配大张量。

ZeRO-R 解决方案：

• 首先尝试在现有空闲块中分配大张量
• 如果不存在足够大的空闲块，则尝试合并相邻的小碎片，形成大碎片
• 如果合并后仍无法满足需求，则触发垃圾回收，释放所有未使用的内存