环境准备：GPU驱动与CUDA安装、PyTorch分布式环境配置、NCCL通信库安装与验证

做多卡训练，环境准备是第一步，也是最容易翻车的一步。

我见过太多同学，代码写得漂漂亮亮，结果一跑分布式就报错——要么驱动版本不对，要么NCCL没装好。说白了，底层环境没搭稳，上层再好的代码也跑不起来。

这一章，我们就来把环境彻底搞定。

1. GPU驱动与CUDA安装

先讲驱动。GPU驱动是操作系统和显卡之间的桥梁。没有它，系统根本认不出你的显卡。

我个人习惯用 nvidia-smi 来检查驱动状态。装好后，终端输入：

nvidia-smi

如果能看到显卡列表、驱动版本、CUDA版本，说明驱动装好了。

举个例子，我有一台机器输出是这样的：

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 535.129.03   Driver Version: 535.129.03   CUDA Version: 12.2     |
+-----------------------------------------------------------------------------+

这里要注意：nvidia-smi 显示的 CUDA Version 是驱动支持的最高CUDA版本，不是你实际安装的CUDA Toolkit版本。这两个概念容易混淆。

接下来装CUDA Toolkit。我建议去NVIDIA官网下载runfile安装，别用deb包。为什么？runfile可以自定义安装路径，方便多版本共存。

安装步骤大致如下：

1. 下载对应系统的runfile
2. 执行 sudo sh cuda_xxx_linux.run
3. 取消勾选Driver（驱动已装好，别重复装）
4. 选择安装路径，比如 /usr/local/cuda-11.8

装完后配置环境变量：

export PATH=/usr/local/cuda-11.8/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH

验证CUDA是否装好：

nvcc --version

看到版本号输出，就对了。

2. PyTorch分布式环境配置

PyTorch的分布式训练，核心依赖 torch.distributed 包。装PyTorch时，要选对CUDA版本。

我个人推荐用conda装，依赖管理省心：

conda create -n torch_dist python=3.10
conda activate torch_dist
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

装完后验证：

python -c "import torch; print(torch.cuda.is_available()); print(torch.cuda.device_count())"

输出 True 和显卡数量，就说明PyTorch能认到GPU了。

接下来测试分布式环境是否正常。我一般写个简单脚本：

import torch
import torch.distributed as dist

dist.init_process_group(backend='nccl', init_method='tcp://127.0.0.1:23456', world_size=2, rank=0)
print(f"Rank {dist.get_rank()} initialized")
dist.destroy_process_group()

用 torchrun 启动：

torchrun --nproc_per_node=2 test_dist.py

如果看到两个rank都打印了初始化成功，说明分布式环境没问题。

3. NCCL通信库安装与验证

NCCL是NVIDIA的集体通信库。多卡训练时，梯度同步全靠它。说白了，NCCL就是分布式训练的「高速公路」。

PyTorch自带了NCCL，但版本可能不是最新的。我建议单独装NCCL，特别是用A100、H100这些新卡时。

安装NCCL有两种方式：

• 方式一：用pip装 — pip install nvidia-nccl-cu11（对应CUDA 11.x）
• 方式二：从源码编译 — 适合需要定制化优化的场景

我个人倾向用pip装，省时间。但要注意版本匹配：

CUDA版本	推荐NCCL版本
11.8	2.18.x 或 2.19.x
12.1	2.20.x 或 2.21.x

装完后验证NCCL是否正常工作。我写了个小测试：

import torch
import torch.distributed as dist

dist.init_process_group(backend='nccl', init_method='tcp://127.0.0.1:23456', world_size=2, rank=0)

# 创建一个张量，做all-reduce测试
tensor = torch.tensor([1.0, 2.0, 3.0]).cuda()
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
print(f"After all-reduce: {tensor}")

dist.destroy_process_group()

如果输出是 [2.0, 4.0, 6.0]（两个rank各贡献一份），说明NCCL通信正常。

4. 整体环境验证

最后，把所有组件串起来跑一遍。我写了个综合验证脚本：

import torch
import torch.distributed as dist
import os

# 检查CUDA
assert torch.cuda.is_available(), "CUDA不可用"
print(f"GPU数量: {torch.cuda.device_count()}")

# 检查NCCL
dist.init_process_group(backend='nccl', init_method='tcp://127.0.0.1:23456', world_size=2, rank=0)
print(f"NCCL版本: {torch.cuda.nccl.version()}")
dist.destroy_process_group()

print("环境准备完成！")

跑通了，说明你的环境已经准备好迎接多卡训练了。

想要更多资料，联系可可