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

作者：Stanislav Belyasov
翻译：陈之炎
校对：赵茹萱


本文约4000字，建议阅读8分钟
本文给出了高效使用内存的关键概念，它适用于多种艰巨的任务。

在训练模型过程中，细数那些完胜“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位量化