深度学习训练模型时，GPU显存不够怎么办？

以下文章来源于：知乎问答

链接：https://www.zhihu.com/question/461811359

编辑来源：CV技术指南

本文仅用于学术分享，如有侵权，请联系后台作删文处理

之前Kaggle有一个Jigsaw多语言毒舌评论分类比赛，当时我只有一张11G显存的1080Ti，根本没法无法训练SOTA的Roberta-XLM-large模型。在这篇文章中，我想分享一些关于如何减少训练时显存使用的技巧，以便你可以用现有的GPU训练更大的深度神经网络。

混合精度训练

第一个可能已经普及的技巧是使用混合精度（mixed-precision）训练。当训练一个模型时，一般来说所有的参数都会存储在VRAM中。很简单，总的VRAM使用量等于存储的参数数量乘以单个参数的VRAM使用量。一个更大的模型不仅意味着更好的性能，而且也意味着更多的VRAM使用量。由于性能相当重要，比如在Kaggle比赛中，我们不希望减少模型的容量。因此减少显存使用的唯一方法是减少每个变量的内存使用。默认情况下变量是32位浮点格式，这样一个变量就会消耗4个字节。幸运的是，人们发现可以在某些变量上使用16位浮点，而不会损失太多的精度。这意味着我们可以减少一半的内存消耗! 此外，使用低精度还可以提高训练速度，特别是在支持Tensor Core的GPU上。

在1.5版本之后，pytorch开始支持自动混合精度（AMP）训练。该框架可以识别需要全精度的模块，并对其使用32位浮点数，对其他模块使用16位浮点数。下面是Pytorch官方文档中的一个示例代码。

梯度积累

第二个技巧是使用梯度积累。梯度累积的想法很简单：在优化器更新参数之前，用相同的模型参数进行几次前后向传播。在每次反向传播时计算的梯度被累积（加总）。如果你的实际batch size是N，而你积累了M步的梯度，你的等效批处理量是N*M。然而，训练结果不会是严格意义上的相等，因为有些参数，如Batch Normalization，不能完全累积。

关于梯度累积，有一些事情需要注意:

1. 当你在混合精度训练中使用梯度累积时，scale应该为有效批次进行校准，scale更新应该以有效批次的粒度进行。

2. 当你在分布式数据并行（DDP）训练中使用梯度累积时，使用no_sync()上下文管理器来禁用前M-1步的梯度全还原，这可以增加训练的速度。

梯度检查点

最后一个，也是最重要的技巧是使用梯度检查点（Gradient Checkpoint）。Gradient Checkpoint的基本思想是只将一些节点的中间结果保存为checkpoint，在反向传播过程中对这些节点之间的其他部分进行重新计算。据Gradient Checkpoint的作者说，在这个技巧的帮助下，他们可以把10倍大的模型放到GPU上，而计算时间只增加20%。Pytorch从0.4.0版本开始正式支持这一功能，一些非常常用的库如Huggingface Transformers也支持这一功能，而且非常简单，只需要下面的两行代码：

在这篇文章的最后，我想分享之前我在惠普Z4工作站上做的一个简单的benchmark。该工作站配备了2个24G VRAM的RTX6000 GPU（去年底升级到2个48G的A6000了），在实验中我只用了一个GPU。我用不同的配置在Kaggle Jigsaw多语言毒舌评论分类比赛的训练集上训练了XLM-Roberta Base/Large，并观察显存的使用量，结果如下。

作者：夕小瑶
链接：

题主提到transformer，那我也针对transformer给点小显存设定下的训练trick吧，好久没去魔改其他模型的结构了

ps：对视觉Transformer了解不深，本文仅从NLP领域的经典Transformer入手来分享一些tricks

从节省（训练）显存方法的有效性/性价比上排个序的话：

1. 砍输入的max len

2. 砍hidden size

3. 砍num layers

4. recompute

5. sharding

6. 小batch size+梯度累积

如果要考虑从预训练模型热启的话，可以去掉2，其他排序不变

为啥最推荐前俩呢，因为不用写代码哇...改改参数，训练显存占用平方级降低，注意是平方级！因此砍一半max seq len和砍一半hidden size比砍一半层数来的有效很多。

只不过砍一半hidden size就不好load预训练模型了，所以优先砍max len吧……

如果非要纠结模型能处理的最大输入长度，可以前90%的训练steps砍输入长度，后10%砍bs艰难finetune下position embedding。这也是预训练常用trick（虽然亲测也会损失一些模型的长文本处理performance）

4和5就是训练大型模型的标配技术了，主流深度学习框架基本都内置了API，调用还是相对方便的，而且不会损失太多训练速度

6就是兜底的办法了，需要写一些代码外加训练慢....

https://www.zhihu.com/question/461811359/answer/1920207852

对于大模型在小显存的 GPU 上训练，通常需要调小 batch size，但是这就意味着更低的准确率。在保持 batch size 不变的情况下正常训练，通常只有两种方法：升级装备（用金钱换显存），用时间换显存。说到性价比，现在一块 32G 的 V100 至少要 6W 起步，而一块 11G 的 1080Ti 大约是 6K 左右，你需要多花 10 倍的金钱才可以让显存容量提升到原来的 3 倍左右。

但是，有了 MegEngine v1.4 之后，你可以不用升级装备，就能训练相对原来 3 倍大小的模型。虽然要花 3 倍左右的训练时间，但是 MegEngine 实现了“从 0 到 1”的过程——原来训练不了的模型，现在可以训练了！

对于 8 张 V100 的训练场景，相比于升级 8 张卡，使用 MegEngine 的 DTR 优化后只需要 8 张 1080Ti，省下了将近 50W，可谓性价比超值。

是什么让 MegEngine 做到了超高性价比训练大模型呢？

MegEngine 引入了 DTR 技术，通过额外计算减少显存占用，从而让小显存也能训练大模型。在 2080Ti 上，ResNet-50、ShuffleNet 等网络的最大 batch size 可以达到原来的 3 倍以上。

无需修改模型代码，只要在训练代码中加入如下几行代码，就可以训练显卡容量的 3~4 倍大小的模型。

关于显存阈值，请根据显卡容量和模型大小设定。同时，DTR 显存优化也支持多卡训练，详细可参考 教程。

目前 MegEngine 仍然在针对 DTR 算法做进一步的优化，例如：去掉 DTR 中的显存阈值提升显存利用率和性能，将 DTR 实现进静态图中提升性能等等，这些特性预计将在下一个版本中发布。

作者：王磊

链接：

前面的很多问答都是从超参或者算法的角度给出了方案，其实从AI框架的角度也可以做一些事情来更高效地利用GPU/NPU设备内存，从而实现用更少的GPU/NPU节点训练更大参数量的深度学习模型的目的。

方案1：框架对生命周期不重叠的数据进行内存复用

MindSpore在其1.1 版本中采用约束编程技术，设计开发了并行拓扑内存约束优化器SOMAS，对并行流信息与数据依赖进行聚合分析，得到算子间祖先关系，构建张量全局生命期互斥约束，再使用多种启发式算法求解获得最优的内存静态规划，实现逼近理论极限的内存复用，可大幅降低多流大型计算图训练内存使用量，进而提升训练的Batch Size数量和训练吞吐量FPS。

以BERT-Large数据并行为例 :采用SOMAS优化后Batch-Size数量可以从160（MindSpore 0.6 Beta）提升到208(MindSpore 1.1)，整体提升30%，优化后的内存占用减少20%-50%。

方案2：框架将部分暂时不需要的数据提前释放，需要时再通过重计算的方式重新获取数据

MindSpore框架采用反向模式的自动微分，根据正向图计算流程来自动推导出反向图，正向图和反向图一起构成了完整的计算图。在计算某些反向算子时，需要用到一些正向算子的计算结果，导致这些正向算子的计算结果需要驻留在内存中，直到依赖它们的反向算子计算完，这些正向算子的计算结果占用的内存才会被复用。这一现象推高了训练的内存峰值，在大规模网络模型中尤为显著。

为了解决这个问题，MindSpore在1.7版本提供了重计算的功能，可以不保存正向算子的计算结果，让这些内存可以被复用，然后在计算反向算子时，如果需要正向的结果，再重新计算正向算子。

从上图可以看出，开启重计算功能后的GPT-3内存占用峰值减少了25%，同时过程中内存占用整体比较稳定。

方案3：框架将部分暂时不需要的数据拷贝至CPU侧，需要时再将该内存拷贝回设备侧，从而达到内存复用的效果

在MindSpore的1.6版本中，我们针对昇腾NPU引入了Memory Swap的能力，Memory Swap可以将部分数据暂存于host侧，帮助用户在不需要增加设备的前提下训练更大规模的模型。后续版本也会提供对于GPU的Memory Swap支持。

针对一张静态计算图，Memory Swap统计并分析其内存使用情况，生成swap策略。在计算执行的同时，会按照既定策略进行内存的拷入拷出以保证计算图执行过程中不会遇到Device内存不够的情况。（MindSpore 1.6 版本Memory Swap仅支持Ascend硬件非下沉执行方式，设置环境变量GRAPH_OP_RUN=1; DISABLE_ASCEND_MIDRT=1后自动生效）。

上图为在单机(8卡)和两机(16卡)的昇腾910设备（Host内存1.5T）上使用并行+异构+Swap，分别可以进行2080亿、4050亿参数量的MoE结构鹏程·盘古模型的训练， 为仅使用并行+异构能够训练的模型规模的三倍左右。

从MindSpore AI框架以上三种方案可以看出，在单卡支持大模型的训练方面，其实框架也有很多事情可以做。