大语言模型基础之‘显存优化‘

上一篇可扩展的训练技术(二)中，我们介绍了零冗余优化器（Zero Redundancy Optimizer, Zero），该技术由DeepSpeed代码库提出，主要用于解决数据并行中的模型冗余技术，即在数据并行训练中，每个GPU上都会存储一份完整的模型参数，这会占用大量的显存。
而Zero技术通过将模型参数分布在多个GPU上，从而减少了显存的使用。

🔍ZeRO有三种逐步细化模型参数和优化器参数的方案：

1. ZeRo-1：仅优化器参数在所有GPU上平摊，模型参数和模型梯度在每张GPU上独立存储。
2. ZeRo-2：优化器参数和模型梯度在所有 GPU 上平摊，模型参数在每张GPU上独立存储。
3. ZeRo-3：优化器参数、模型梯度和模型参数在所有 GPU 上平摊。

---

下面我们举一个实际的例子来解释下各方案之间的区别。

常识：1字节（bytes）= 8位（bits）

LLaMa 7B模型，参数量为70亿，也就是70亿个参数。

模型参数常用的数据类型是FP32和FP16。

FP32：每个参数用32位（bit）浮点数存储，即4字节（Bytes）。
FP16：每个参数用16位（bit）浮点数存储，即2字节（Bytes）。

假设我们使用FP32的数据类型，将7B的权重加载到GPU中，那么需要70亿*4字节=280亿字节=28GB的显存。

假设我们使用FP16的数据类型，将7B的权重加载到GPU中，那么需要70亿*2字节=140亿字节=14GB的显存。

一般模型训练时使用混合精度训练，混合精度训练使用FP16和FP32两种数据类型，FP16用于计算，FP32用于存储，这样既能提高计算速度，又能减少显存占用。(混合精度训练的原理可以参考大语言模型基础之‘可扩展的训练技术‘(二)-混合精度训练)

优化器状态：一般训练时使用Adam优化器，Adam优化器需要维护一份fp32的模型参数（4 bytes * 7B = 28 GB）、fp32的动量参数（4 bytes * 7B = 28 GB）以及fp32的动量二阶矩参数（4 bytes * 7B = 28 GB）。

大模型训练过程中的显存占用情况如下：（由大到小）

1. 优化器参数占用。
   – fp32 模型参数：4 bytes * 7B = 28 GB
   – fp32 动量参数：4 bytes * 7B = 28 GB
   – fp32 动量二阶矩参数：4 bytes * 7B = 28 GB

2. 模型参数占用。
   –fp 16存储，2 bytes * 7B = 14 GB

3. 模型梯度占用。
   – fp 16存储，2 bytes * 7B = 14 GB

4. 激活值占用及其他。
   与batch_size和序列长度有关，与上述三个参数相比，占用的显存较少。

因此，训练：
使用FP16的数据类型，不考虑训练过程中的激活值占用和其他占用的情况下，将7B的权重加载到GPU中，需要28GB（fp32 模型参数）+ 28GB（fp32 动量参数）+ 28GB（fp32 动量二阶矩参数）+ 14GB（fp16 模型参数）+ 14GB（fp16 模型梯度）= 112GB的显存。
推理大约需要14GB的显存。

---

下面我们来看下DeepSpeed ZeRO三个阶段ZeRO-1、ZeRO-2、ZeRO-3三种方案在显存占用上是如何节省显存的。

来自ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

由上图可以看出，在训练大语言模型时，优化器参数占据了显存的大部分，因此，优化器参数的优化是显存优化的重要方向。

ZeRO-1：切分优化器参数。Pos — 将内存使用量减少 4 倍.
优化器参数在所有GPU上平摊，模型参数和模型梯度在每张GPU上独立存储。

每个GPU运行模型的整个前向和反向传播时，它会从其他GPU的参数广播中访问参数以进行更新。

优点：优化器参数的显存占用减少。
缺点：模型参数和模型梯度在每张GPU上独立存储，需要频繁的通信，通信开销大。

ZeRO-2：切分优化器参数和模型梯度。Pos+g — 将内存使用量减少 8 倍。
优化器参数和模型梯度在所有 GPU 上平摊，模型参数在每张GPU上独立存储。

优点：优化器参数和模型梯度的显存占用减少。
缺点：模型参数在每张GPU上独立存储，需要频繁的通信，通信开销大。

ZeRO-3：切分优化器参数、模型梯度和模型参数。Pos+g+p — 将内存使用量减少 Nd （GPU数量）倍。
优化器参数、模型梯度和模型参数在所有 GPU 上平摊。

优点：优化器参数、模型梯度和模型参数的显存占用减少。
缺点：通信开销大。

因此，在训练大语言模型时，我们一般使用ZeRO-3方案，该方案在显存占用上具有较好的优化效果。

---

以上就是关于Deepspeed库中的ZeRO-1、ZeRO-2、ZeRO-3介绍，也希望给有需要的同学一点点帮助。

后边的文章，我将会继续讲些大语言模型的相关基础知识，请各位多多支持。

📢 📢 📢
感谢各位阅读，大家的点赞👍- 关注🔥- 收藏⭐ - 评论📝 四连，都是博主坚持协作、更新高质量博文的最大动力！