Nvidia驱动更新后深度学习框架可以调用共享GPU内存进行模型训练

1.背景

记录一个在工作中偶然发现的情况。接触过深度学习的朋友的都知道，在使用GPU（显卡）训练模型时，影响模型训练的关键因素有两个：算力和显存。其中，算力影响模型训练的速度，显存决定了每个batch喂给GPU训练的数据量的上限。

因为工作内容的原因，经常需要将高分辨率的图像数据作为数据集，对模型进行训练。为了提高模型对高分辨率图像特征的学习能力，每个batch只喂一张图给GPU。经常在跑模型的时候都是把显存拉满了，偶尔还会报OOM（显存溢出）的错误，着实头痛，因为大显存的显卡都比较贵，经费有限，没有使用像RTXA6000这样的48GB显存以上的显卡（目前使用的最大显存为24GB，型号：RTX 4090）。有的朋友可能想到了多卡并联，但是多卡并联不能把一张图分给多个显卡同时训练（目前没有找到相关的解决方案，也许有）。

2.问题发现

在前一段时间更新了Nvidia的显卡驱动，版本从531升级到了536。之后训练模型的时候，突然就发现当显存溢出之后可以调用主板的内存了，发现问题的时候比较疑惑，并不知道是更新了驱动之后的影响（因为之前也经常更新驱动）。后面经过排除法，又用其他计算机做了测试，才发现是显卡驱动版本更新带来的优化。得知这个结果后还是比较高兴的，这样在显存偶然溢出的时候就不会报OOM的错误了，而且现在有ddr5的高频内存支持下，对训练速度要求不是特别高的情况下，可以暂时用共享GPU内存作为补充。

3.结论

后面又在不同的显卡和计算机上做了测试，目前的结果如下：
（1） 在Nvidia的10系、20系、30系、40系的显卡上更新驱动到536版本后均可实现共享GPU内存进行模型训练，536以上版本测试了随机测试了几个也可以（应该536版本以上驱动均支持该功能）。
（2） 在Nvidia的RTX A和quadro上的显卡上测试536的驱动不能实现调用共享GPU内存进行模型训练（主要测试了RTx A2000 、RTX A5000显卡等Ampere架构的显卡），而且该系列的多数显卡最新版本只有536的驱动。

关于536版本驱动的详细跟新说明链接如下，有兴趣的朋友可以去看看：https://us.download.nvidia.cn/Windows/536.99/536.99-win11-win10-release-notes.pdf

4.补充说明

（1） 以上测试均在win10平台中进行
（2） 深度学习的框架用的pytorch 1.11.0
（3） cuda版本11.1

有什么其他疑问欢迎大家在评论区讨论，互相学习！