记一次在 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，给电脑安装显卡驱动的时候就不要装更低的版本了。

 

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

 

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

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

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

conda install pytorch cudatoolkit=11.0 -c pytorch

 

好了，至此基于 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

 

我们电脑上的 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 版本信息。

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

 

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

 

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

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

 

本着 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 的版本号。

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

 

这里假如你在安装 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 的一半 🤣

 

附上一些数据