记一次在 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 只已被捕捉
    • stevenlyf Apple Safari | 604.1 iPhone 14_7_1

      3090现在淘宝我看都快2w了学长买的时候多少钱,自己买的么Σ(‘◉⌓◉’)(。•ˇ‸ˇ•。)

      • 千千 Chrome | 92.0.4515.159 Windows 10

        我当时联系的时候商家说是 1w2,的确涨价涨的太快了。我也没买,这个是租的,呜呜,当时就应该买下来的,算是理财产品

    • Trouble Maker Chrome | 88.0.4324.182 Windows 10

      博主的博客网站是自己写的吗,我也梦想着有一个属于自己的博客网站,真的好喜欢这种可以diy的东西,很有新鲜感显得与众不同,同时做完以后也成就感爆棚!别人的博客网站总是别人的,自己在自己写的博客网站是发表文章真的感受是不一样的,真的做的太好看了。我觉得你的生活状态就是我想要的,努力工作,同时也作者自己感兴趣的东西,努力赚钱的同时也有着自己的生活,真的太好了

      • 千千 Chrome | 89.0.4389.114 Windows 10

        基于 word press 搭建的,如果你也想搭建的话网上有一键式建站教程,像阿里云、腾讯云之类的也提供相应的服务,搭建起来特别容易,不过 DIY 成自己喜欢的样子不一定轻松了。
        哈哈,学生时代第一次接触的确感觉很与众不同,成就感特别爆棚,每天都想上去看一看。不过工作或者读研以后,自己的精力放在别处时便更新的也慢了,也不想做太多的修改。所以趁着自己有时间还有精力可以好好学着做做,以后可能就没有那种心情了。(☆ω☆)