GPU 的显存都去哪儿了？

Where Did the GPU Memory Go?

GPT-2（XL）有 15 亿个参数，在 16 位精度下，这些参数大约占用 3GB 的显存（VRAM）。

在 16 位精度下，一个参数占用 2 字节的内存，所以 15 亿个参数会占用 3GB 的显存。你估计在单张 GPU 上训练 GPT-2 至少需要多少显存？

• 优化器 → Adam
• 批大小 → 32
• Transformer 层数 → 48
• 序列长度 → 1000

假设你想在 GPT-2 上做一些尝试，那么上面的配置意味着：显存占用可能在 25GB - 25GB - 0GB 这一区间附近。

答案可能会让你大吃一惊。

在一张 32GB 显存的 GPU 上，勉强才能训练一个 3GB 大小的 GPT-2 模型。

但这怎么可能呢？显存到底都花到哪里去了？让我们来一探究竟。在训练过程中，模型在许多方面都会持续占用显存。

混合精度训练被广泛用于加速模型训练。顾名思义，它的思路是在可行的地方（比如卷积和矩阵乘法）使用较低精度的 float16，并搭配 float32 一起使用——这正是“混合精度”名称的由来。前向传播和反向传播都使用 16 位表示的权重和梯度来完成。Φ 因此，如果模型有 P 个参数，那么：

• Φ 权重将占用 2*P 字节。
• Φ 梯度将占用 2*P 字节。

这里的数字“2”表示每个参数占用 2 字节（16 位）的显存。此外，反向传播结束时的更新仍然在 32 位下进行，以保证计算的有效性。我说的是下面图中圈出的那一步：

Adam 是模型训练中最受欢迎的优化器之一。许多从业者只是因为它流行就用它，却没有意识到：在训练过程中，Adam 会保存两个优化器状态来计算更新——梯度的动量和方差：

Φ 因此，如果模型有 P 个参数，那么这两个优化器状态将占用：

• Φ 4*P 字节用于动量。

• Φ 另外 4*P 字节用于方差。

这里的数字“4”表示每个参数占用 4 字节（32 位）的显存。

最后，如上图所示，最终的更新总是在模型权重的 32 位表示下进行调整。这导致：

• Φ 另外 4*P 字节用于模型参数。

让我们把它们加起来：

Φ 总共是 16*P，也就是 24GB 的显存，这比 16 位参数所用的 3GB 显存高得离谱。而且我们还没把所有开销都算进去。

对于像大语言模型（LLM）这样的大型深度学习模型，激活值在训练期间会占用大量显存。更正式地说，GPT-2 的一个 Transformer 块中计算出的激活值总数为：

因此，跨所有 Transformer 块，总计为：

这是 GPT2-XL 的配置：

算下来大约有 300 亿个激活值。由于每个激活值都用 16 位表示，所有激活值合计占用 60GB 的显存。

借助梯度检查点等技术（在上一章讨论过），可以以多花 25-30% 运行时间为代价，将其降低到大约 8-9GB。这项技术把总显存占用带到大约 32-35GB 的范围——这就是我之前提到的，针对一个区区 3GB 的模型，而且批大小还只有 32。除此之外，还涉及更多额外开销，比如显存碎片化。当已分配的显存块之间存在细小、未使用的空隙时，就会发生这种情况，导致可用显存被低效利用。

由于无法获得连续的显存块，显存分配请求会失败。

在上面的讨论中，我们考虑了一个相对较小的模型——拥有 15 亿参数的 GPT-2（XL），与如今训练的模型规模相比，它微不足道。不过，这段讨论也许能帮助你反思构建 LLM 时固有的挑战。许多人常说，GPT 无非就是在模型里堆叠越来越多的层、把网络做得更大。如果真那么简单，人人都会去做了。从这次讨论中你或许已经明白，这并不只是多加几层那么简单。哪怕只多加一层，都可能带来数 GB 的额外显存需求。多 GPU 训练是这些模型的核心，我们在本书前面的一章中已经讲过。