分布式训练之ZeRO学习面试题

• 一、什么是 3D 并行？ 3D 并行可以让大型模型以非常有效的方式进行训练 二、3D 并行策略有哪些？
• DataParallel （DP）
• TensorParallel （TP）
• PipelineParallel （PP） 2.1 DataParallel （DP）

• 介绍：假设有 N 张卡，每张卡都保存一个模型，每一次迭代（iteration/step）都将 batch 数据分割成 N 个等大 小的 micro-batch，每张卡根据拿到的 micro-batch 数据独立计算梯度，然后调用 AlIReduce 计算梯度均值， 每张卡再独立进行参数更新。
• 举例说明： #假设模型有三层：LO，LI，L2 # 每层有两个神经元＃两张卡

GPUO： GPU1： 2.2 TensorParallel （TP）

• 介绍：每个张量都被分成多个块，因此不是让整个张量驻留在单个 GPU 上，而是张量的每个分片都驻留在 其指定的 GPU 上。在处理过程中，每个分片在不同的 GPU 上分别并行处理，最终结果在步骤结束时同步。 这也被称作横向并行。
• 举例说明： #假设模型有三层：LO，LI，L2 # 每层有两个神经元＃两张卡 GPUO：

2.3 PipelineParallel （PP）

• 介绍：模型在多个 GPU 上垂直（层级）拆分，因此只有模型的一个或多个层放置在单个 GPU 上。每个 GPU 并行处理管道的不同阶段，并处理一小部分批处理。
• 举例说明： #假设模型有 8 层

GPUO 三.为什么需要 ZeRO？ 虽然 DataParallel （DP）因为简单易实现，所以目前应用相比于其他两种广泛，但是由于 DataParallel （DP）需要 每张卡都存储一个模型，导致显存大小成为制约模型规模的主要因素。

cO

GPU1

既然每张卡都存储一个模型会增加模型训练过程中的显存占用，那么是否可以让每行卡训练 1/N 的模型参 数，然后合并起来就是一个完整模型呢？这样，随着卡数的增加，每张卡用于模型训练的显存占用将减低，能 够训练的模型也就越大。 如今训练大模型离不开各种分布式并行策略，ZeRO 系列技术就是一种显存优化的数据并行方案，旨在训练超大 规模的语言模型。

• 四、ZeRO 的核心思想是什么？ 去除数据并行中的冗余参数，使每张卡只存储一部分模型状态，从而減减少显存占用。 五、ZeRO 显存如何分配？ ZeRO 将模型训练阶段中每张卡的显存内容分为两类：
• 模型状态：包括参数、梯度和优化器状态，其中优化器状态占比 75%。
• 剩余状态：除了模型状态之外的显存占用，包括激活值、各种临时缓冲区以及无法使用的显存碎片。 来看一个例子，GPT-2 含有 1.5B 个参数，如果用 fp16 格式（混合精度），只需要 3GB 显存，但是模型状态实际上 需要耗费 24GB！所以模型状态就成了头号显存杀手，它也是 ZeRO 的重点优化对象。而其中优化器状态又是第一 个要被优化的。 六、.ZeRO 优化策略是怎么样？ 针对模型状态的存储优化（去除冗余），ZeRO 使用的方法是分片，即每张卡只存 1/N 的模型状态量，这样系统 内只维护一份模型状态。 ZeRO 具有三个主要的优化阶段（ZeRO-1，ZeRO-2，ZeRO-3），它们对应于优化器状态（optimizer states）、梯度（gradients）和参数（parameters）的分片。累积启用时：
• 优化器状态分区（Pos）-内存减少 4 倍，通信量与数据并行性相同
• 添加梯度分区（Postg）内存減少 8 倍，通信量与数据并行性相同
• 添加参数分区（Postg+p）-内存减少与数据并行度 Nd 成线性关系。例如，拆分为 64 个 GPU （Nd =64）内 存将减少到 1/64。GPU 通信量略有增加 50%。

Baseline

注：图中 Memory Consumption 第二列给出了一个示例：k=12，中=7.5B,Nd=64，可以看到随着 ZeRO 阶段深入，显存优化相当明显。

GPU 24 GPU 16

GPU O

Model Parallel

GPU 20 GPU 12 GPU 28 GPU

Pipeline Parallel

Memory Consumption Volume

Formulation Specific Example 120GB 31.4G8

16.6G8

1.9G8

1.5x N. • Optimizer States Gradients

一张卡训不了大模型，根因是显存不足，ZeRO-Offload 则将训练阶段的某些模型状态下放（offload）到内存以 及 cPU 计 t 算，即显存不足，内存来补。相比于昂贵的显存，内存廉价多了。 下图是某一层的一次迭代过程，使用了混合精度训练，前向计算（FWD）需要用到上一次的激活值 （activation）和本层的参数（parameter），反向传播（BWD）也需要用到激活值和参数计算梯度，当我们用 Adam 优化器进行参数更新时，假设模型参数量是 M，在混合精度训练的前提下，边的权重要么是 2M （fp16），要么是 4M （fp32）。为了不降低计算效率，将前两个节点放在 GPU，后两个节点不但计算量小还需 要和 Adam 状态打交道，所以放在 CPU 上，Adam 状态自然也放在内存中，为了简化数据图，将前两个节点融合 成一个节点 FWD-BWD Super Node，将后两个节点融合成一个节点 Update Super Node，沿着 gradient 16 和 parameter 16 把数据流图切分为两部分，分布对应 GPU 和 CPU

activation 16

GPU 2M

activation 16 FWD-BWD 4M Super Node activation 16| parameter 16 parameter 16

CPU

float2half

2M •2M

gradient 16

'Update Super Node

parameter 32 parameter 32 parameter 32 Param update momentum 32 Param update float2half variance 32 variance 32 variance 32

• 七、ZeRO Offload 后的计算流程是怎么样？
• 在 GPU 上面进行前向和后向计算，将梯度传给 CPU。同时为了提高效率，可以将计算和通信并行起来，GPU 在反向传播阶段，可以待梯度值填满 bucket 后，一边计算新的梯度一边将 bucket 传输给 CPU，当反向传播结 束，CPU 基本上已经有最新的梯度值了；
• cpu 进行参数更新，再将更新后的参数传给 GPU。 Step i+1 P update FWD & BWD Computation stream： GPU CPU GPU GPU->CPU Swapping stream： GPU->CPU