问：模型大小超过GPU容量怎么办？这篇文章的灵感来自Yandex数据分析学院讲授的“高效深度学习系统”课程。预备知识：假设读者已经了解神经网络的前向和反向传递是如何工作的，这对于理解本文的内容至关重要。本文使用PyTorch作为框架。开始吧！当尝试使用具有超过5亿个参数的大型模型（又名gpt-2-xl）时，您的GPU资源受到限制以使其适合在GPU上运行，或者在模型训练期间无法实现论文中定义的batchsize，此时怎么办？也许你可以选择放弃，使用更轻量级的模型，或者减小训练批量大小，在这种情况下，你将无法获得论文中描述的训练结果。但是，有一些技术可以帮助解决上述问题。让我们讨论一些方法，如何使用这些方法对具有15亿个参数的GPT-2-XL模型进行微调。问题的核心首先，让我们了解将模型加载到GPU所需的GPU内存问题的核心。假设模型有FP32（32位浮点）参数，你需要在GPU上训练这个模型，例如运行Adam优化器。通过计算，结果是惊人的。假设您已经拥有一块内存为12GB的NVIDIAGeForceRTX3060。首先，1e9FP32参数占用大约4GB的GPU内存。同样，对于梯度，将保留相同数量的内存。因此，总共预留了8GB的内存。由于还没有开始训练，还没有加载优化器，加载优化器也需要一定的内存。Adam优化器需要为每个参数存储第一个备份和第二个备份，这需要8GB的额外内存。计算一下，您必须有大约16GB的GPU内存才能将模型正确加载到GPU上，而在本文的示例中，它只有12GB的可用内存。看起来很糟糕，对吧？然而，有几种方法可以尝试解决这个问题，这里是相关的：梯度累积/微批处理；梯度检查点；模型并行训练；流水线作业；张量并行混合精度训练；内存卸载；.接下来，将详细解释这些技术。开始发问：模型大于GPU容量怎么办？简单模式：无法将batchsize适配为1Professional模式：参数也无法适配概述如果模型大于GPU容量，即使将batchsize设置为1也不行，怎么办？有一个解决方案，就是设置梯度检查点，我们来看看这个概念。对于一个简单的n层前馈神经网络，梯度的计算图如下：神经网络各层的激活值对应于标有f的节点，在前向传播过程中，对所有这些节点按顺序排列。与这些层的激活和参数对应的损失梯度由标记为b的节点表示。在向后传递期间，所有这些节点都以相反的顺序进行评估。f个节点的计算结果用于计算b个节点，所以前向传播后所有f个节点都保存在内存中。只有当反向传播进行到足以计算f节点的所有依赖项时，它才会从内存中删除。这意味着：简单反向传播所需的内存随着神经网络层数n线性增长。下面是这些节点的计算顺序，紫色阴影圆圈表示在给定时间需要将哪个节点保存在内存中。梯度检查点如上所述的简单反向传播在计算上是最优的：它只计算每个节点一次。但是，如果重新计算节点，可能会节省大量内存。例如，可以简单地重新计算每个节点。执行顺序和使用的内存如下图所示：这种策略在内存方面是最优的。但是请注意，节点计算的数量被缩放了n2次，而之前的缩放因子是n：每n个节点按顺序重新计算n次。由于计算速度慢，这种方法不适合深度学习。为了在内存和计算之间取得平衡，需要提出一种策略，允许重新计算节点，但不要太频繁。这里使用的策略是将神经网络激活的子集标记为检查点节点。在此示例中，选择将第sqrt(n)个节点标记为检查点。这样，checkpoint节点数和checkpoint之间的节点数都在sqrt(n)之间，也就是说：需要的内存量也是按照n的量级缩放的。该策略所需的额外计算量相当于网络的单次前向传递所需的计算量。例程：了解了梯度检查点的细节后，让我们看看如何在PyTorch中应用这个概念。看起来并不太难：梯度累积/微批次概述在GPU内存中安装如此庞大的神经网络。因此，您在训练期间被迫使用较小的batchsize，这可能会导致收敛速度变慢和准确性降低。什么是梯度累积？在训练神经网络时，数据通常是分批处理的，神经网络预测批标签，用于计算相对于实际目标的损失。接下来，执行反向传递以计算梯度并更新模型权重。梯度累积修改了训练过程的最后一步：在进行下一个mini-batch之前，保存梯度值并将新的梯度添加到先前保存的梯度中，而不是为每个mini-batch更新网络权重。只有在模型处理了几个小批量之后才会更新权重。梯度累积模拟更大的批量。如果要小批量使用64张图片，如果batchsize超过8，就会报“CUDAmemoryerror...”。在这种情况下，您可以使用8批次图像，并在模型处理完64/8=8批次后更新一次权重。如果你从这8个批次中累积每个梯度，结果将（几乎）相同，并且可以进行训练！套路：没有梯度累积的标准训练循环通常是：在PyTorch中，梯度累积可以很容易地完成。模型用accumulation_steps处理完mini-batch后，就可以进行优化了。您还可以使用accumulation_steps根据损失函数的性质划分运行损失：非常漂亮，对吧？梯度在调用loss.backward()时计算并由PyTorch累积，直到调用optimizer.zero_grad()。重要一些网络架构使用专门的批处理操作，例如BatchNorm，当使用相同的批处理大小时，结果可能会略有不同。混合精度训练概述混合精度训练是指将部分或全部FP32参数转换为更小的格式，例如FP16、TF16（浮点张量）或BF16（浮点字节）。主要优点混合精度训练的主要优点是：减少内存使用；更快的性能（更高的算术强度或更小的通信足迹）；使用专用硬件进行更快的计算。目前只对第一个优势感兴趣——减少内存使用，让我们看看如何使用PyTorch模型实现这一点。套路：结果，.half()操作完成后，模型变小了2倍。将模型转换为不同格式（即BF16、TF16）后的缩放损失将在后续文章中讨论。有些操作无法在FP16中完成，例如Softmax。PyTorch可以利用torch.autocast来处理这些特殊情况。8位优化器增加模型大小是获得更好性能的有效方法。然而，训练大型模型需要存储模型的状态、梯度和优化器（例如，亚当的指数平滑和以及先前梯度的平方和），所有这些都在有限的可用内存中。将32位优化器降级为8位优化器，将取值范围从232缩小到恰好2?=256，对优化器保留的内存量影响巨大。研究人员提出了一种新的8位Adam优化器，作者在论文中说：“它将32位性能保持到部分原始内存中”。8位优化器包含三个组件：(1)块级量化，隔离异常值并将错误均匀分布到每一位；（2）动态量化，对小值和大值进行高精度量化；(3)stableEmbedding层提高wordembedding优化模型的稳定性。使用这些组件，可以直接使用8位状态进行优化。将8位优化器状态量化为32位，执行更新，然后将状态量化为8位存储。寄存器中的逐元素8位到32位转换，无需缓慢复制到GPU内存或额外的临时内存来执行量化和反量化。对于GPU，这意味着8位优化器比常规的32位优化器更快。看看使用8-bitAdam后的鼓舞人心的结果：可以看出使用量化的Adam可以节省大约8.5GB的GPU内存，看起来相当不错！了解了它的可用性之后，我们来看看如何用python实现它。Facebook提供的Bitsandbytes包是对CUDA自定义函数的轻量级封装，封装了8位优化器和量化函数，用它来实现8位Adam的使用。Routine：上面说了quantizedoptimizer使用起来很简单，效果也不错。结合以上所有方法，在GPU上对GPT-2-XL进行微调。最后，在掌握以上方法后，利用这些方法解决实际问题，微调15亿参数的GPT-2-XL模型。显然，无法将其加载到具有12GBRAM的NVIDIAGeForceRTX3060GPU之上。列出所有可以使用的方法：gradientcheckpointing；混合精度训练（我有一个技巧：使用相同模型的两个样本。首先，将其加载到GPU上。half，命名为gpu_model；其次，在CPU上，命名为cpu_model。评估GPU模型后，加载gpu_model到cpu_model的梯度，运行optimizer.step()，并将更新后的参数加载到gpu_model上）；使用batch_size=64,minibatch_size=4的梯度累加需要通过accumulation_steps对loss进行缩放；8位Adam优化器。使用以上所有方法并查看代码：使用以上所有方法后，在GPU上实现了16GBGPT-2-XL模型的微调，绝！结论在此博客中，介绍了有效使用内存的关键概念，这些概念适用于如上所述的各种困难任务。其他概念将在后续文章中讨论。非常感谢您花时间阅读这篇文章！
