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

共 4063字,需浏览 9分钟

 ·

2022-08-26 14:15


大数据文摘授权转载自数据派THU
作者:Stanislav Belyasov 
翻译:陈之炎 
校对:赵茹萱

在训练模型过程中,细数那些完胜“CUDA 内存出错..”报错的提高内存效率技术。
 
提问:模型大小超过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模型微调,绝了!

结论

在本博中,给出了高效使用内存的关键概念,它适用于多种艰巨的任务,如上文所述。

将在后续的文章中讨论其他概念。

衷心感谢,拨冗阅读本文!

原文标题:
How to fine tune VERY large model if it doesn’t fit on your GPU
原文链接:https://medium.com/@bestasoff/how-to-fine-tune-very-large-model-if-it-doesnt-fit-on-your-gpu-3561e50859af文:Stanislav Belyasov


点「在看」的人都变好看了哦!
浏览 41
点赞
评论
收藏
分享

手机扫一扫分享

分享
举报
评论
图片
表情
推荐
点赞
评论
收藏
分享

手机扫一扫分享

分享
举报