使用系统内NCCL环境重新编译Pytorch

intro：

费了老大劲,来重新编译pytorch，中间报了无数错误。原生的编译好的pytorch是直接用的其自带NCCL库，并且从外部是不能进行插桩的，因为根本找不到libnccl.so文件。下面记录下重新编译pytorch的过程。指定USE_SYSTEM_NCCL=1。这里不知道为啥，用原生库进行编译的时候，总会爆nvlink相关的错误，如果有知道的大神，可以在评论区赐教，万分感谢。

相关环境配置

Ubuntu 22.04.5 LTS

cuda 11.8+ cudnn 8.9.7

python 3.11.10

torch V2.0.1+ nccl 2.21.5

NVIDIA GeForce RTX 4090 D

gcc (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0

cmake version 3.22.1

如果WSL和cuda以及cudnn不太会配的可以参考：Windows11 + WSL Ubuntu + Pycharm + Conda for deeplearning | 公孙启

以下是查看一些环境使用的bash命令，可以根据本地配置的路径自行修改，关于系统内搭建NCCL环境的方法可以参考：手动配置NCCL环境_nccl 参数设置-CSDN博客。这里是实际上是使用了一个Anaconda的虚拟环境，以下工作都要在虚拟环境内进行：

conda create -n py_pi python=3.11
conda activate py_pi

python --verion
lsb_release -a
nvcc --version
cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR -A 2
cat /usr/local/cuda/include/nccl.h | grep NCCL_MAJOR -A 2

clone pytorch指令如下，如果速度慢或者失败，建议不要再WSL下进行，因为WSL似乎没法直接科学上网，可以在宿主机windows下直接科学上网clone即可，再打开WSL挂载到相应的目录下。或者可以更换镜像源

git clone --branch v2.0.1 --recursive https://github.com/pytorch/pytorch
pip install -r requirements.txt

pytorch源码的编译：

完成以上步骤就可以开始源码编译了，这是一个非常耗时的工作，应该需要一个多小时的时间，具体时间CPU支持的MAX_JOBS的数量有关，我这里写好了一个bash脚本，直接保存运行即可（是在pytorch 的路径下）。这里建议把numpy降级为1.26.3：

pip install numpy==1.26.3

#!/bin/bash# 设置环境变量
export NCCL_ROOT=/usr/local/cuda
export NCCL_INCLUDE_DIR=$NCCL_ROOT/include
export NCCL_LIB_DIR=$NCCL_ROOT/targets/x86_64-linux/lib
export LD_LIBRARY_PATH=$NCCL_LIB_DIR:$LD_LIBRARY_PATH
export LIBRARY_PATH=$NCCL_LIB_DIR:$LIBRARY_PATH
export CPATH=$NCCL_INCLUDE_DIR:$CPATH
export MAX_JOBS=20
#这里的环境变量可以根据你的实际环境去改变,但是如果是按照我以上的流程去配的，使用这个应该没有问题。
#MAX_JOBS数量务必根据你的cpu情况去自行设定，设的太大了内存会报错。我这里是intel i7 14700k 设置为20的话绰绰有余。
# 清理之前的构建
make clean
python setup.py clean
rm -rf build# 编译 PyTorch
USE_GLOO=0 USE_SYSTEM_NCCL=1 NCCL_INCLUDE_DIR=$NCCL_INCLUDE_DIR NCCL_LIB_DIR=$NCCL_LIB_DIR python setup.py develop
#注意这里USE_GLOO=0，否则会报错，我也不知道为啥。

编译过程中可能会有各种各样的问题，相当折磨，所以需要耐心一点。如果出现错误，找到相关的错误日志，然后可以问一下chatgpt或者找一下网上有没有类似的问题。我遇到的问题在上面已经避坑，补充一个可能遇到的问题，虽然我没遇到：解决 libstdc++.so.6: version ‘GLIBCXX_3.4.30‘ not found 问题

如果编译完成了,你可以使用以下的脚本来测试：


import torchprint("torch version",torch.__version__)
print(torch.cuda.is_available(), torch.distributed.is_nccl_available())
print("nccl version:",torch.cuda.nccl.version())
print("cuda version:", torch.version.cuda)       cudnn_version = torch.backends.cudnn.version()
print("cuDNN version:", cudnn_version)
print(torch.cuda.device_count(), torch.cuda.get_device_name(0))

输出类似于这样就可以了：

torch version 2.0.0a0+gite9ebda2
True True
nccl version: (2, 21, 5)
cuda version: 11.8
cuDNN version: 8907
1 NVIDIA GeForce RTX 4090 D

其他

如果在进行一些测试的时候需要torchvision和torchaudio，需要在pytorch的文件夹目录下git clone，然后同样执行setup.py。我这里使用的版本是：

git clone --branch release/0.15 https://github.com/pytorch/vision.git torchvision
git clone --branch release/2.0 https://github.com/pytorch/audio.git torchaudio

然后进到这两个子项目文件夹下执行：