|
|
深度学习训练模型时，GPU显存不够怎么办？
关注者
428
被浏览
462,164
52 个回答
感觉这个问题是年经贴(´-ω-`)
题主提到transformer，那我也针对transformer给点小显存设定下的训练trick吧，好久没去魔改其他模型的结构了qaq
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就是兜底的办法了，需要写一些代码外加训练慢....
当然了，以上trick可以叠加混合使用，拿好不谢(^-^)
此外，还有一些不方便公开的内部黑科技就不写出来啦，感兴趣的小伙伴欢迎私信发我简历，组内还有实习head count噢！
坐标：北京，百度科技园
方向：NLP、搜索
对于大模型在小显存的 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 倍大小的模型。
import megengine as mge
mge.dtr.eviction_threshold = "5GB" # 设置显存阈值为 5GB
mge.dtr.enable() # 开启 DTR 显存优化关于显存阈值，请根据显卡容量和模型大小设定。同时，DTR 显存优化也支持多卡训练，详细可参考 教程 。
目前 MegEngine 仍然在针对 DTR 算法做进一步的优化，例如：去掉 DTR 中的显存阈值提升显存利用率和性能，将 DTR 实现进静态图中提升性能等等，这些特性预计将在下一个版本中发布。
更多关于 MegEngine 动态图显存优化的实现和效果，可查看：