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

前言

• 显存的挑战：当前常见的GPU 显存如 A100 是 80G 显存，而如 175B 的 GPT-3 模型需要 175B * 4bytes 即 700GB 模型参数空间，连加载都困难更不要说训练了。
• 计算的挑战：175B模型所需要的计算量也是非常庞大的，并且叠加预训练数据量的增大，所需的计算量与BERT时代完全不可同日而语。

三大并行策略

并行技术可以整体划分为两类：数据并行（DP），流水线并行(PP)以及张量并行(TP)。

1 数据并行(DP)

• 原理：数据并行简单来说，就是每个GPU worker都有一个完整模型的副本，然后获取 per_gpu_batch_size 的数据，每个worker在自己那批数据上计算梯度。最后所有worker 对梯度汇总进行平均最后更新模型权重。
• 举例：假设有 8 张卡，per_gpu_batch_size=4，那么 global batch size=8*4=32，此时每张卡都保存一个模型，然后每张卡通过 4 个数据来分别计算梯度，之后计算梯度均值，最后每张卡再进行参数更新。

2 流水线并行(PP)与张量并行(TP)

2.1 流水线并行

• 思想：将模型的不同层放到不同的 GPU 上。比如：一个6层的trasformer ，那么就可以将前3层放在一个GPU上，后3层放在另一个GPU上。如下图所示：

• 缺点：每个GPU在处理相同batch时都需要等待前一个GPU完成计算才能开始处理，这样容易产生GPU空闲时间，即流水线气泡。为了改善这种情况，通常使用流水线并行（PP）来通过让不同的batch数据的计算重叠。

2.2 张量并行

• 思想：对层内进行分割，实际上是将Tensor 矩阵运算分配到不同的 GPU 上。如下图就是将每一层分为两部分：

• 优点：张量并行存储效率较高，
• 缺点：每次前向和后向计算时会引入额外通信开销。

DeepSpeed 与 Zero

什么是Deepspeed？

Deepspeed Zero 的核心思想

• 模型状态：模型参数(fp16/bf16，2A)，模型梯度(fp16/bf16，2A)，Adam状态(fp32的模型参数备份，fp32的momentum和fp32的variance，4A+4A+4A)。假设模型参数量为 A，则共需要：2A + 2A + (4A + 4A + 4A) = 16A 字节存储，如一个7B的模型，模型状态就需要 16 * 7 = 112GB
• 剩余状态：除了模型状态之外的显存占用，包括激活值（activation）、各种临时缓冲区（buffer）以及无法使用的显存碎片（fragmentation）。

zero的优化思想为：针对模型状态部分，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

针对 Adam 状态进行分片，此时每张卡模型状态所需显存变成了。我们以单机8卡，7B模型为例，此时模型状态部分所需显存 降低到 4 * 7 + (12*7)/8 = 38.5GB。

1. Deepspeed Zero Stage 2

在Zero-1 的基础上，针对模型梯度再次分片，此时每张卡模型状态所需显存变成了  同样以单机8卡，7B模型为例，此时模型状态部分所需显存降低到 2*7+(14*7)/8=26.25GB。

1. Deepspeed Zero Stage 3

在Zero-2的基础上，对模型参数也进行分片，此时每张卡模型状态所需显存变成了。

1. Zero-offload

一点补充

最后

参考