xformers pip 安装

最近复现的一些仓库用 xformers 的很多，xformers 安装的版本需要对应 CUDA 版本 和 pytorch 版本。

而且在 arm 的 aarch64 下安装所有和 CUDA 相关的库都不是非常方便，这里记录一下。

参考：
https://github.com/facebookresearch/xformers
https://blog.csdn.net/x1131230123/article/details/139231686

---

确定版本：

首先要确定版本，xFormers 高度绑定 pytorch 和 CUDA 版本，基本上一一对应。

如果仓库的 requirements.txt 给了，那就最好，只要保证其它库版本一致。

如果 pytorch 有改动，那么要求找一下 xFormers 对应的版本。

比如打开这个：
https://github.com/facebookresearch/xformers/tree/v0.0.28

选择tag，在 README.md 可以看到

二分查找你的 pytorch 需要对应什么版本的 xformers，比如 pytorch==2.1.0 对应 0.0.22.post7

安装：

如果 x86，直接 pip 安装即可，记得带上 cuda 版本：

pip3 install -U xformers==0.0.28 --index-url https://download.pytorch.org/whl/cu118

如果 arm，或者 x86 下找不到可用的预编译的特定 cuda 和 pytorch 版本的包，

就需要从源码安装：

pip install ninja

pip install -v -U git+https://github.com/facebookresearch/xformers.git@v0.0.28#egg=xformers

@ 后面的填 tag 。

验证安装：

python -m xformers.info

官方使用文档：
https://facebookresearch.github.io/xformers/

常见报错：

RuntimeError: Error compiling objects for extension：

参考：
https://github.com/facebookresearch/xformers/issues/748

解决方法：引入环境变量：

export TORCH_CUDA_ARCH_LIST="8.0 8.6 9.0+PTX"
export MAX_JOBS=4

MAX_JOBS=4 是因为安装时可能会内存不够，实测 =4 时内存占用在 40 GB 左右
还不够可以试一下设成 1，但是会非常慢，应该要好几个小时

TORCH_CDA_ARCH_LIST 设为实际的算力，下面的指令可以看计算能力：

python -c "import torch; print(torch.cuda.get_device_capability())"
# sample output : (8, 0)

那么就是：

export TORCH_CUDA_ARCH_LIST=8.0

也可能需要指定 gcc 的版本，我在某个 issue 上看到 10.3 的 gcc 比较好，尝试了一下 11 以上会直接报错说不支持：

module unload compilers/gcc
module load compilers/gcc/10.3.1