记一次在 RTX 3090 上安装 APEX

背景

最近炼丹开始用一块 RTX 3090 (24 G),因为代码里用 ALBERT-base-v2 处理了很多东西导致显存爆炸,于是开始谋求可以节约显存的办法。

网上的一些方法例如及时 del 无用的参数,测试过基本显存占用没有什么变化,也不确定是不是用的方法不对。

之前已经实现了梯度累计,的确可以用 batch_size = 1 占用的显存模拟出 batch_size = n 的情况,但损失的当然是一定的时间,以及如果你用到了 bn 层可能会让模型的效果轻微降低(毕竟多个样本可以一定程度上降低 bias)。

暂时没找到其他有用的方法节约显存,于是想试试混合精度用起来如何。

本文记录一次在 RTX 3090 + pytorch 1.7.0 + cuda 11.1 上的 APEX 安装经历。


合适的 nVidia 驱动与合适的 pytorch

首先确保你的电脑 nVidia 驱动版本支持 RTX 3090,这一点可以在 https://www.nvidia.cn/geforce/drivers/ 查到。

如下图,RTX 3090 最低需要的驱动版本是 455.23,给电脑安装显卡驱动的时候就不要装更低的版本了。

img

 

如下图所示,安装完成后执行 nvidia-smi 就可以看到这块显卡了,也能看到 Driver Version: 455.23.04。

img

 

然后便是 pytorch 之类的环境,这里推荐使用 conda 来维护所有不同的环境。

30 系列显卡是新一代架构,新驱动不支持 cuda 9 以及 cuda 10,所以必须安装 cuda 11。

目前 pytorch 官网上表示使用 conda 最新可安装的 cudatoolkit 是 11.0 的版本。

conda install pytorch cudatoolkit=11.0 -c pytorch

img

 

好了,至此基于 cuda 11.0 编译的 pytorch 已经安装完成。接下来是 APEX 的安装。


APEX 安装记录 & CUDA

APEX 官方给的安装其实挺容易的,但实际操作起来可能出现各种坑。

以下是安装 APEX 的方法,记得在执行 pip 时切换到目标虚拟环境(conda activate envs_name)。

git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --no-cache-dir --global-option="--pyprof" --global-option="--cpp_ext" --global-option="--cuda_ext" ./

 

假设你的电脑上没有安装过完整的 cuda(例如仅仅在虚拟环境里安装了 cudatoolkit),那最后一步大概率会说没有找到 nvcc。

nvcc 是 cuda 中提供的一个编译工具,它并不包含在专门针对 pytorch / tensorflow 运行所提供的 cudatoolkit 里。

因此实际上我们还需要 nvcc,我们还需要安装完整的 cuda。

 

重要的是,针对以上,我们选择哪个版本的 cuda 才能成功安装 apex。

因为刚刚我们提到,在虚拟环境里安装 cudatoolkit 11.0 即可正常使用 RTX 3090。那此时 cuda 版本是否也直接安装为 11.0 即可呢?

这里有一份 cuda 版本与 nVidia 显卡驱动版本的对照表:https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html

img

 

我们电脑上的 nVidia 驱动版本目前是 455.23,看这张对照表实际上应该安装 cuda 11.1 及以上才可行。(但为什么 cudatoolkit 11.0 就可以了呢?这个没具体了解,如果有了解的小伙伴可以相互讨论)

根据表上的信息,需要安装 cuda 11.1 的完整版,具体的安装方法也可以找别人的教程,我这里直接搬官方的安装指令过来(适用于 Ubuntu)。

安装的过程中可以只选 cuda 部分(假设之前已安装过 nVidia 驱动可不选)。

wget https://developer.download.nvidia.com/compute/cuda/11.1.0/local_installers/cuda_11.1.0_455.23.05_linux.run
sudo sh cuda_11.1.0_455.23.05_linux.run

 

安装完毕,配置好环境变量(如果需要)后使用 nvcc -V 查看 cuda 版本信息。

如果有类似于下面的这张图,那么恭喜你安装完成。

img

 

最后一步,安装 APEX,还是上文中的那三条指令。其实你在执行指令的时候会报错,报错信息如下:

img

 

这就是刚才所纠结的问题,我们 conda 虚拟环境中的 pytorch 是基于 cuda 11.0 编译的,但此时安装的完整版 cuda 是 11.1,版本号匹配不上。

不过先别着急重新下载完整版的 cuda 11.0,因为这里也有坑,后文说。

img

 

本着 cuda 大版本差异大,小版本差异小的心态,我们直接把 apex setup.py 文件中关于版本号的校验改掉,应该就可以正常安装了。

如下,找到 get_cuda_bare_metal_version 函数,其中的 bare_metal_major 是主版本,即 11;而 bare_metal_minor 是次要版本,我们需要改的就是这个,直接在下面让它等于 0 即可。

这样 pytorch 是基于 cuda 11.0 编译的,而我们完整版的 cuda 11.1 也会在 nvcc -V 识别中返回 11.0 的版本号。

校验一致,可以正常安装。

img

 

这里假如你在安装 apex 时提示类似于 apex Error : Given no hashes to check 137 links for project 'pip': discarding no candidates 时,确保不是 cuda / cudnn 版本的问题后,可通过以下方法解决。

git clone https://github.com/NVIDIA/apex
cd apex
python setup.py install

其他坑

上一节中我们提到:

不过先别着急重新下载完整版的 cuda 11.0,因为这里也有坑,后文说。

我尝试安装了 cuda 11.0,然后在编译安装 apex 时版本号的确校验过了,但是会有如下错误:

nvcc fatal : Unsupported gpu architecture 'compute_86'

猜测是 cuda 11.0 不支持 RTX 3090 这张卡,从这一步中想要继续也许可行,具体的方法可自行查阅。


RTX 3090 vs TITAN RTX vs V100 (16 GB)

这里放一组使用 RTX 3090、TITAN RTX 以及 V100 (16 GB) 在单纯 ALBERT-base 下训练与微调速度的测试对比。

实验数据仅作参考,并非权威数据。实验参数直接搬当时代码里的参数过来,max_seq_length 是允许 tokenizer 返回的最大句子长度。

args = {
    'batch_size': 4,
    'test_batch_size': 4,
    'accumulation_steps': 1,
    'lr': 0.001,
    'fine_tune_lr': 0.000005,
    'adam_epsilon': 0.000001,
    'weight_decay': 0.01,
    'epochs': 4,
    'fine_tune_epochs': 2,
    'max_seq_length': 512,
}

 

实验数据(train 和 finetune 阶段是什么可不必在意)

显卡 train 阶段 finetune 阶段
RTX 3090 11.84 it/s 4.32 it/s
TITAN RTX 6.47 it/s 2.39 it/s
V100 (16 GB) 7.34 it/s 2.69 it/s

看上去 RTX 3090 真的很香,价格仅是 TITAN RTX 的一半 🤣

 

附上一些数据

img

img


  • 15 只已被捕捉
    • NIMITIZ Chrome | 87.0.4280.88 Windows 10

      十年后捡垃圾拿到3090再试φ( ̄∇ ̄o)
      我记得国外有大佬把孤岛危机塞进了3090的24G缓存显存

      • 千千 Chrome | 86.0.4240.193 Windows 10

        哈哈,或者买一块 3090 用两三年卖个 7000+,血赚

    • 站元素主机 Mozilla FireFox | 83.0 Windows 10

      学习了赞一个

      • 千千 Chrome | 86.0.4240.193 Windows 10

        哈哈,欢迎常来哦~

    • 择偶网 搜狗浏览器 Windows 7

      新冠肆虐,注意安全!

      • 千千 Chrome | 86.0.4240.193 Windows 10

        好哇,谢谢谢谢啦~

    • Buok Chrome | 86.0.4240.198 Windows 10

      鸽子更新啦?!

      • 千千 Chrome | 86.0.4240.193 Windows 10

        没鸽不配称为鸽子

        • Buok Chrome | 86.0.4240.198 Windows 10

          咕咕咕,咕咕,咕!!!!

    • 熊猫小A Chrome | 87.0.4280.67 Mac OS X

      诶,PyTorch 1.6 以上混合精度训练是自带的哦,不需要 Apex

      • 千千 Chrome | 86.0.4240.193 Windows 10

        哈哈,这个知道的。但是对低版本的 PyTorch 上不太适用。貌似目前 apex 和自带的 amp 效率上也有点区别,没细看。