PyTorch源码编译：如何指定自定义NCCL库路径与链接方式

本文将详细介绍如何从源码编译安装PyTorch，并重点演示在编译过程中如何使用自定义版本的NCCL（NVIDIA Collective Communications Library）动态库或静态库，以满足特定环境或性能调优需求。以下是完整的操作步骤与说明。

本文假定您已经完成了GPU驱动与CUDA的基础安装。

安装Anaconda

Anaconda是一个强大的Python虚拟环境管理工具，能够方便地创建和管理多个独立的Python环境，避免包依赖冲突。

下载Anaconda
从官网下载安装脚本，例如 Anaconda3-2024.10-1-Linux-x86_64.sh，并执行安装。

更改Anaconda国内源
为了加速包下载，建议配置国内镜像源。编辑或创建 ~/.condarc 配置文件：

conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge/
conda config --set show_channel_urls yes

更改pip国内源
同样，为pip配置国内源以加速Python包的安装。创建 ~/.pip/pip.conf 文件：

[global]
index-url = https://pypi.tuna.tsinghua.edu.cn/simple/
[install]
trusted-host=https://pypi.tuna.tsinghua.edu.cn

更新
完成配置后，建议更新conda及其包至最新状态。

conda update conda -y
conda update anaconda -y
conda update python -y
conda update --all -y

安装完成后，系统会有一个默认的 base 虚拟环境。您可以在此环境下操作，或创建新的独立环境。

(base) ~# conda env list
# conda environments:
#
base                  *  /root/anaconda3

使用命令安装PyTorch

首先，创建一个名为 pytorch 的虚拟环境。

conda create -n pytorch python=3.9
conda activate pytorch

然后，前往 PyTorch官网 获取适合您CUDA版本的安装命令。例如：

conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

安装成功后，可以验证PyTorch及内置NCCL的版本。官方发布的预编译包通常使用NCCL静态库。

(pytorch) ~# python
Python 3.9.20 (main, Oct  3 2024, 07:27:41)
[GCC 11.2.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.__version__
'2.5.1'
>>> torch.cuda.nccl.version()
(2, 21, 5)

通过源码安装PyTorch

如果您需要更灵活的配置或特定的NCCL版本，从源码编译是更好的选择。此处以 pytorch-2.1.2 版本为例（其默认依赖的NCCL版本为 nccl-2.18.6），具体步骤可参考官方指南。

下载源码
克隆仓库并确保子模块同步更新。

git clone --recursive https://github.com/pytorch/pytorch
cd pytorch
# 如果是更新现有代码库
git submodule sync
git submodule update --init --recursive

安装依赖
编译前需要安装一些必要的工具和库。

conda install cmake ninja
pip install -r requirements.txt
conda install mkl mkl-include
conda install -c pytorch magma-cuda121

编译PyTorch
默认情况下，编译脚本会使用PyTorch源码中自带的NCCL子模块，并以静态库形式链接。

MAX_JOBS=32 CUDA_HOME=/usr/local/cuda python setup.py develop

编译成功后，在基础环境中验证版本。此时使用的NCCL版本是源码内置的。

(base) ~# python
Python 3.10.12 (main, Jul  5 2023, 18:54:27) [GCC 11.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.__version__
'2.1.2+cu121'
>>> torch.cuda.nccl.version()
(2, 18, 6)
>>>

使用自定义NCCL版本
通过设置环境变量，可以指示编译系统使用您自己编译或下载的NCCL库，这在深度学习模型训练的环境定制中非常有用。

• NCCL_INCLUDE_DIR: 指定自定义NCCL头文件路径。
• NCCL_LIB_DIR: 指定自定义NCCL库文件路径。
• USE_SYSTEM_NCCL: 设为 1，告知系统使用自定义NCCL。
• USE_STATIC_NCCL: 设为 1 表示静态链接，不设置或设为 0 表示动态链接。

示例1：使用自定义NCCL版本，并以静态方式链接

NCCL_INCLUDE_DIR="/root/nccl-2.18.6/build/include" NCCL_LIB_DIR="/root/nccl-2.18.6/build/lib"  USE_SYSTEM_NCCL=1 USE_STATIC_NCCL=1 MAX_JOBS=32 CUDA_HOME=/usr/local/cuda python setup.py develop

示例2：使用自定义NCCL版本，并以动态方式链接

NCCL_INCLUDE_DIR="/root/nccl-2.18.6/build/include" NCCL_LIB_DIR="/root/nccl-2.18.6/build/lib"  USE_SYSTEM_NCCL=1 MAX_JOBS=32 CUDA_HOME=/usr/local/cuda python setup.py develop

运行分布式任务验证

为了验证自定义NCCL编译的PyTorch能否正常工作，我们可以在两台机器上部署，并运行一个简单的分布式集合通信任务。

脚本准备
机器A（rank 0）上的脚本 machine_a.py：

#!/usr/bin/env python3
import os
import torch
import torch.distributed as dist
from datetime import timedelta

t=timedelta(seconds=100)
tensor = torch.randn(100, 100).cuda()
print("before", tensor)
dist.init_process_group(backend='nccl', init_method='tcp://172.10.16.136:12345', rank=0, world_size=2, timeout=t)
dist.all_reduce(tensor)
print(tensor)

机器B（rank 1）上的脚本 machine_b.py：

#!/usr/bin/env python3
import os
import torch
import torch.distributed as dist
from datetime import timedelta

t=timedelta(seconds=100)
tensor = torch.randn(100, 100).cuda()
print("before", tensor)
dist.init_process_group(backend='nccl', init_method='tcp://172.10.16.136:12345', rank=1, world_size=2, timeout=t)
dist.all_reduce(tensor)
print(tensor)

执行顺序
首先在机器A上启动脚本，然后在机器B上启动脚本。

输出结果
机器A输出（节选）：

before tensor([[-0.4122,  1.5498,  0.5335,  ..., -0.5380,  0.0839,  0.2863],
        [-0.2936,  0.2634, -1.4883,  ..., -0.3352,  0.1396, -0.3077],
        [ 1.9145, -1.3494,  1.3779,  ...,  0.6003, -0.2153, -0.0569],
        ...,
        [ 0.0897, -0.7549,  0.4063,  ..., -0.5300,  0.1050, -0.5825],
        [-0.2933,  0.8490,  1.5020,  ..., -0.1208, -1.4131,  0.1459],
        [-1.1359, -1.4827, -0.6804,  ..., -1.0101, -1.4612, -1.4856]],
       device='cuda:0')
tensor([[ 0.8489,  2.5036,  1.8025,  ..., -0.7820,  0.1438,  0.2032],
        [ 0.6158,  0.6582, -1.0910,  ..., -0.0594, -0.2070, -2.2146],
        [ 2.3681, -0.8484,  0.9504,  ...,  1.1521, -0.2830,  0.0953],
        ...,
        [-0.5323,  0.9372, -0.3793,  ..., -0.9586,  0.3927, -0.0074],
        [ 0.0589,  1.5489,  1.7448,  ..., -1.1915, -0.3582,  1.5681],
        [ 0.2856, -0.8354, -1.0480,  ..., -1.0900, -2.5273, -1.6689]],
       device='cuda:0')

机器B输出（节选）：

before tensor([[ 1.2610,  0.9538,  1.2690,  ..., -0.2441,  0.0599, -0.0831],
        [ 0.9094,  0.3947,  0.3973,  ...,  0.2758, -0.3466, -1.9069],
        [ 0.4536,  0.5011, -0.4275,  ...,  0.5518, -0.0677,  0.1523],
        ...,
        [-0.6220,  1.6921, -0.7856,  ..., -0.4286,  0.2877,  0.5751],
        [ 0.3523,  0.6999,  0.2428,  ..., -1.0707,  1.0549,  1.4222],
        [ 1.4216,  0.6473, -0.3676,  ..., -0.0799, -1.0661, -0.1834]],
       device='cuda:0')
tensor([[ 0.8489,  2.5036,  1.8025,  ..., -0.7820,  0.1438,  0.2032],
        [ 0.6158,  0.6582, -1.0910,  ..., -0.0594, -0.2070, -2.2146],
        [ 2.3681, -0.8484,  0.9504,  ...,  1.1521, -0.2830,  0.0953],
        ...,
        [-0.5323,  0.9372, -0.3793,  ..., -0.9586,  0.3927, -0.0074],
        [ 0.0589,  1.5489,  1.7448,  ..., -1.1915, -0.3582,  1.5681],
        [ 0.2856, -0.8354, -1.0480,  ..., -1.0900, -2.5273, -1.6689]],
       device='cuda:0')

可以看到，两台机器执行 all_reduce 操作后，输出得到了完全一致的张量结果，这证明了基于自定义NCCL编译的PyTorch分布式通信功能工作正常。

通过以上步骤，您应该能够成功地从源码编译安装PyTorch，并根据自身需求灵活地指定NCCL库的版本和链接方式。如果在实践过程中遇到问题，欢迎到云栈社区的相应板块与更多开发者交流讨论。