如何在GPU资源受限情况下微调超大模型

提问：模型大小超过GPU 容量怎么办？ 

本文的灵感来自于Yandex数据分析学院教授的“高效深度学习系统”课程。

预备知识：假设读者已经了解神经网络的前传递和后向传递的工作原理，这对理解本文内容至关重要。文中使用PyTorch作为框架。

开始吧！

当试图使用大型模型(即aka gpt-2-xl)，它带有 5亿多个参数，而你的GPU 资源受限，无法将它安装到GPU上运行，或者在模型训练期间无法实现论文中定义的批大小，此时该怎么办？也许可以选择放弃，使用一个更轻量级版本的模型，或者减小训练的批大小，这样的话，便无法获得论文中描述的训练结果。

但是，有一些技术可以帮助解决上述问题。

下面来讨论一些方法，即如何利用这些方法来微调带有15亿个参数的GPT-2-XL模型。

问题的核心

首先，来了解一下将模型加载到GPU中所需GPU内存问题的实质。

假设模型具有 个FP32（32位浮点）参数，需要在GPU上训练这个模型，例如，运行Adam优化器。

通过计算，结果令人震惊。

假设已有一块带有12 GB内存的NVIDIA GeForce RTX 3060。首先， 1e9个FP32参数约占4 GB的GPU内存。

同样，对于梯度，也将保留相同数量的内存。所以，总共已经保留了8 GB的内存，由于还没有开始训练，也没有加载优化器，加载优化器也同样需要一定数量的内存。Adam优化器需要为每个参数存储第一备份和第二备份，即需要8 GB额外内存。算下来，必须有大约16 GB的GPU内存，才能正确地将模型加载到GPU上，在本文的例子中，GPU只有12 GB的空闲内存。看起来很不妙，对吧？

然而，可以通过一些方法来尝试解决这个问题，以下是相关内容：

• 梯度积累/微批量；
• 梯度检查点；
• 模型并行训练；
• 管道作业；
• 张量并行化
• 混合精度训练；
• 内存卸载；
• 优化器8位量化。

接下来，将详细解读这些技术。

开始

提问：模型比GPU容量大，怎么办？

• 简单模式：无法适配批大小为1
• 专业模式：参数也没办法适配

概述

如果模型大于GPU容量，即便将批大小设为1都不够，那该怎么办呢？有一个解决方案，即设置梯度检查点，下面来看看这个概念。对于一个简单的包含n层的前馈神经网络来说，梯度的计算图如下：

神经网络层的激活对应于用f标记的节点，在正向传递期间，按顺序对所有这些节点进行计算。对应于这些层的激活和参数的损失梯度用b标记的节点表示。在反向传递期间，所有这些节点都以相反的顺序进行计算。f个节点的计算结果用于计算b个节点，因此所有f个节点在向前传递后都保存在内存中。只有当反向传播进展到足够计算出f节点的所有依赖关系时，它才能从内存中擦除。这意味着：简单的反向传播所需的内存随神经网络层数n的变化呈线性增长。

下面是这些节点的计算顺序，紫色阴影圆圈表示在给定时间里需要将哪个节点保存到内存之中。

梯度检查点​

如上所述的简单反向传播在计算方面是最优的：它只计算每个节点一次。但是，如果重新计算节点，可能会节省大量内存。例如，可以简单地重新计算每个节点。执行的顺序和所使用的内存如下图所示：

这种策略在内存方面是最优的。但是，请注意，节点计算的数量进行了n²次缩放，而先前的缩放系数为n：每个n个节点都按n次顺序重新计算。由于计算速度较慢，这种方法并不适用于深度学习。

为了在内存和计算之间取得平衡，需要提出一种策略，允许重新计算节点，但次数不要太频繁。在这里使用这样一种策略：将神经网络激活的一个子集标记为检查点节点。 

在本示例中，选择将第sqrt(n)个节点标记为检查点。这样，检查点节点的数量和检查点之间的节点数量都在sqrt(n)之间，这意味着：所需的内存量也按n的顺序进行了缩放。该策略所需的额外计算量相当于网络单次前向传递所需的计算量。

例程：

在学习了梯度检查点的细节之后，来看看如何在PyTorch中应用这个概念，看起来并不太难： 

梯度累积/微批次

概述

深度学习模型正在越变越大，很难在GPU内存中安装这样大型的神经网络。因此，被迫在训练时选用较小的批大小，它可能导致较慢的收敛和较低的准确性。

什么是梯度累积？

在训练神经网络时，通常会将数据分批量处理，神经网络预测批处理标签，用于计算相对于实际目标的损失。接下来，执行反向传递计算出梯度，更新模型权值。梯度累积对训练过程的最后一步进行了修正：在继续下一个小批之前，保存梯度值，并将新的梯度添加到之前保存的梯度中，用这种方法取代更新每个小批的网络权重。只有在模型处理了几个小批次后，才会更新权重。梯度积累模拟了一个更大的批大小，如果想在一个小批中使用64张图像，如果批大小超过了8，则会报“CUDA内存出错…”。在这种情况下，可以使用8批图像，并在模型处理64/8=8批后更新一次权重。如果你从这8个批次中积累每一个梯度，结果将是（几乎）相同的，这样便能够执行训练啦！

例程：

没有梯度累积的标准训练环通常为： 

在PyTorch中，梯度累积可以很容易地完成。模型利用accumulation_steps处理完成小批之后，便可以执行优化。还可以利用accumulation_steps根据损失函数的性质来划分运行损失： 

真漂亮，对吗？当调用loss.backward() 时计算梯度，并由PyTorch累积，直到调用optimizer.zero_grad()时停止。

重点

某些网络体系结构使用专用的批处理操作，如BatchNorm，当使用相同的批大小时，结果可能会略有不同。

混合精度训练

概述​

混合精度训练是指将部分或全部FP32参数转换为更小的格式，如FP16、TF16（浮点张量）或BF16(浮点字节)。

主要优势

混合精度训练的主要优势是：

• 减少内存使用；
• 性能提速（更高的算术强度或更小的通信占用）；
• 使用专用硬件进行更快地计算。

目前只对第一个优势感兴趣——减少内存的使用量，来看看如何使用PyTorch模型实现它。

例程：

结果，在完成.half()操作之后，模型变小了2倍。将模型转换为不同的格式(即BF16，TF16)后的缩放损失，将在后续的文章中讨论。有些操作在FP16中是无法完成的，如Softmax。PyTorch可利用torch.autocast 来处理这些特殊情况。

8位优化器

增加模型尺寸是获得更佳性能的有效途径。然而，训练大模型时需要存储模型、梯度和优化器的状态(例如，Adam的指数平滑和及先前梯度的平方和)，所有这些都存储在数量有限的可用内存之中。

将32位优化器降到8位优化器，将数值的范围从2³²减少到仅2⁸=256，会对优化器预留的内存数量产生巨大的影响。

研究人员提出了一种新的8位Adam优化器，论文作者在文中这么说： “它将32位的性能维持到部分原始内存中”。

8位优化器有三个组成部分：（1）块级量化，隔离异常值，将误差均匀地分配给每一个比特；（2）动态量化，高精度地量化小值和大值；（3）稳定的嵌入层，以提高词嵌入优化模型的稳定性。

有了这些组件，可直接使用8位状态执行优化。将8位优化器状态量化为32位，执行更新，然后再将状态量化为8位进行存储。在寄存器中逐元素进行8位到32位的转换，无需慢速复制到GPU内存或额外的临时内存中执行量化和去量化。对于GPU来说，这意味着8位优化器要快于常规的32位优化器。

来看看使用8位Adam之后，鼓舞人心的结果：

可以看出，使用量化的Adam可以节省大约8.5 GB的GPU内存，看起来相当棒！

理解了它的可用性之后，再来看看如何用python实现它。

由Facebook提供的Bitsandbytes 包是一个围绕CUDA自定义函数的轻量级包装器，封装了 8位优化器和量化函数，利用它可以实现8位Adam的使用。

例程：

如上所述，量化优化器的使用非常简单，结果也不错。

综合上述全部方法，对GPU上的GPT-2-XL进行微调。

最后，在掌握了上述方法之后，利用这些方法来解决实际问题，对拥有15亿个参数的GPT-2-XL模型进行微调。显然，无法将它加载到12 GB内存的NVIDIA GeForce RTX 3060 GPU之上。列出可以使用的全部方法：

• 梯度检查点；
• 混合精度训练(我设了一个技巧：使用相同模型的两个样本。首先，用.half将它加载到GPU上，将其命名为gpu_model；其次，在CPU上，将其命名为cpu_model。评估好GPU模型之后，将 gpu_model的梯度加载到cpu_model中，运行optimizer.step()，将更新后的参数加载到gpu_model上)；
• 使用batch_size=64，minibatch_size=4的梯度累积，需要通过 accumulation_steps来缩放损失；
• 8位Adam优化器。

把以上方法全部利用起来，查看一下代码： 

利用上述所有方法之后，在GPU上实现了对16GB的GPT-2-XL模型微调，绝了！

结论

在本博中，给出了高效使用内存的关键概念，它适用于多种艰巨的任务，如上文所述。将在后续的文章中讨论其他概念。衷心感谢，拨冗阅读本文!