背景
最近炼丹开始用一块 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 的一半 🤣
附上一些数据
十年后捡垃圾拿到3090再试φ(￣∇￣o)
我记得国外有大佬把孤岛危机塞进了3090的24G缓存显存
哈哈，或者买一块 3090 用两三年卖个 7000+，血赚
学习了赞一个
哈哈，欢迎常来哦~
新冠肆虐，注意安全！
好哇，谢谢谢谢啦~
鸽子更新啦？！
没鸽不配称为鸽子
咕咕咕，咕咕，咕！！！！
诶，PyTorch 1.6 以上混合精度训练是自带的哦，不需要 Apex
哈哈，这个知道的。但是对低版本的 PyTorch 上不太适用。貌似目前 apex 和自带的 amp 效率上也有点区别，没细看。