环境准备：GPU服务器选型建议、NVIDIA驱动与CUDA安装、NCCL库的安装与验证、Docker环境配置

说实话，做多卡通信这件事，最怕的不是代码写不对，而是环境没搭好。我见过太多人花了两周调代码，最后发现是驱动版本不匹配——那种感觉，嗯，挺崩溃的。所以这一章，咱们把地基打牢。

1. GPU服务器选型建议

选服务器这事儿，我个人的经验是：别只看显卡，要看整机带宽。

你想想看，NCCL通信的核心是跨卡数据传输。如果PCIe带宽不够，或者NVLink没连上，那再好的卡也白搭。

关键指标

• GPU型号：A100、H100、V100 是目前主流。A100支持NVLink 3.0，带宽高达600 GB/s。H100更是到了900 GB/s。
• GPU数量：单机8卡是最常见的配置。4卡也能跑，但通信拓扑会受限。
• NVLink拓扑：8卡A100通常是全互联（All-to-All），但有些廉价方案是环形或混合。我建议你买之前问清楚：是不是全NVLink互联？
• CPU与内存：至少64核、256GB内存。别让CPU成为瓶颈。
• 网络：多机通信需要InfiniBand或RoCE。单机训练的话，网卡可以省一点。

2. NVIDIA驱动与CUDA安装

驱动和CUDA的安装，说白了就是版本匹配游戏。我建议你遵循一个原则：先定CUDA版本，再选驱动版本。

安装步骤

1. 检查当前驱动：nvidia-smi 查看驱动版本和GPU状态。
2. 卸载旧驱动：sudo apt-get purge nvidia-* 或 sudo yum remove nvidia-*。
3. 安装驱动：从NVIDIA官网下载.run文件，或使用包管理器。我个人习惯用.run文件，因为可控性更强。
4. 安装CUDA：同样从官网下载.run文件。注意选择与驱动匹配的版本。
5. 设置环境变量：在 ~/.bashrc 中添加：

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

验证安装

nvidia-smi          # 查看驱动和GPU
nvcc --version      # 查看CUDA版本

3. NCCL库的安装与验证

NCCL是NVIDIA的集体通信库。说白了，它就是让多卡之间高效传数据的工具。

安装方式

• 方式一：pip安装（推荐）

pip install nvidia-nccl-cu12

• 方式二：源码编译（适合定制化需求）

git clone https://github.com/NVIDIA/nccl.git
cd nccl
make -j src.build
make install

验证NCCL

运行NCCL自带的测试程序：

# 下载测试代码
git clone https://github.com/NVIDIA/nccl-tests.git
cd nccl-tests
make

# 运行AllReduce测试（8卡）
./build/all_reduce_perf -b 8 -e 128M -f 2 -g 8

如果看到类似下面的输出，说明NCCL工作正常：

# nThread 1 nGpus 8 minBytes 8 maxBytes 134217728 step: 2(factor) warmup iters: 5 iters: 20
#
#     size       time   algbw   busbw
#      (B)    (usec)  (GB/s)  (GB/s)
       8       12.3    0.65    0.65
      16       12.5    1.28    1.28
      ...
  134217728   234.5   572.3   572.3

4. Docker环境配置

Docker是解决环境一致性的最佳方案。我个人强烈建议你在Docker里跑NCCL，尤其是多机场景。

安装Docker与nvidia-docker

# 安装Docker
curl -fsSL https://get.docker.com | bash

# 安装nvidia-container-toolkit
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
sudo apt-get update
sudo apt-get install -y nvidia-container-toolkit
sudo systemctl restart docker

拉取并运行NCCL镜像

# 拉取官方镜像
docker pull nvcr.io/nvidia/pytorch:23.12-py3

# 运行容器（8卡）
docker run --gpus all -it --rm \
  --shm-size=32g \
  --ulimit memlock=-1 \
  --ulimit stack=67108864 \
  nvcr.io/nvidia/pytorch:23.12-py3

验证Docker内NCCL

在容器内运行：

python -c "import torch; print(torch.cuda.nccl.version())"

如果输出类似 (2, 18, 3)，说明NCCL已就绪。

知识体系总览

下面这张图概括了本章的核心逻辑。你可以把它当作环境搭建的检查清单：