本文的灵感来自于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的空闲内存。看起来很不妙,对吧?
然而,可以通过一些方法来尝试解决这个问题,以下是相关内容:
接下来,将详细解读这些技术。
概述
如果模型大于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 来处理这些特殊情况。
增加模型尺寸是获得更佳性能的有效途径。然而,训练大模型时需要存储模型、梯度和优化器的状态(例如,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之上。列出可以使用的全部方法:
把以上方法全部利用起来,查看一下代码:
利用上述所有方法之后,在GPU上实现了对16GB的GPT-2-XL模型微调,绝了!
在本博中,给出了高效使用内存的关键概念,它适用于多种艰巨的任务,如上文所述。将在后续的文章中讨论其他概念。衷心感谢,拨冗阅读本文!
以上是如何在GPU资源受限情况下微调超大模型的详细内容。更多信息请关注PHP中文网其他相关文章!