3、软件环境搭建：CUDA、cuDNN、NCCL版本匹配，PyTorch/TensorFlow分布式包安装，Docker容器化环境配置

多卡训练的环境搭建，说实话，是很多人翻车的第一站。

我记得刚带团队做分布式训练那会儿，有个同学折腾了两天，最后发现是CUDA和PyTorch版本没对上。嗯，这种坑我踩过不止一次。所以这一章，我把这些年积累的版本匹配经验、安装技巧、还有Docker化的最佳实践，一次性讲清楚。

3.1 CUDA、cuDNN、NCCL：三兄弟的版本哲学

这三者的关系，说白了就是：CUDA是底层驱动，cuDNN是加速库，NCCL是通信库。多卡训练时，NCCL负责卡与卡之间的数据交换，cuDNN负责单卡上的计算加速。

我个人习惯用这张表来对照版本：

CUDA版本	cuDNN推荐版本	NCCL推荐版本	PyTorch对应版本
11.8	8.9.x	2.16.x	2.1.x - 2.3.x
12.1	8.9.x	2.18.x	2.2.x - 2.4.x
12.4	9.0.x	2.21.x	2.4.x+

3.2 NCCL：多卡通信的灵魂

NCCL（NVIDIA Collective Communications Library）是分布式训练的核心。你想想看，8张卡要同步梯度，数据量动辄几百MB，没有高效的通信库根本跑不动。

安装NCCL其实很简单：

# 通过apt安装（推荐）
sudo apt-get install libnccl2 libnccl-dev

# 或者从NVIDIA官网下载deb包
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/nccl-local-repo-ubuntu2204-2.21.5-cuda12.4_1.0-1_amd64.deb
sudo dpkg -i nccl-local-repo-ubuntu2204-2.21.5-cuda12.4_1.0-1_amd64.deb
sudo apt-get update
sudo apt-get install libnccl2 libnccl-dev

3.3 PyTorch分布式包安装

PyTorch的分布式支持，核心就两个包：torch.distributed 和 torch.nn.parallel.DistributedDataParallel。前者负责通信后端，后者负责模型包装。

安装时，我建议直接用官方推荐的pip命令：

# CUDA 12.1版本
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

# CUDA 11.8版本
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

为什么要强调用官方源？因为第三方源经常版本不对。我在项目中遇到过，有人从清华源装了个PyTorch，结果NCCL后端死活起不来。后来换成官方源，一次过。

3.4 TensorFlow分布式包安装

TensorFlow的分布式训练，主要靠 tf.distribute.MirroredStrategy 和 tf.distribute.MultiWorkerMirroredStrategy。安装时注意：

# TensorFlow 2.15+ 推荐
pip install tensorflow[and-cuda]

# 或者指定CUDA版本
pip install tensorflow==2.15.0
pip install nvidia-tensorrt

3.5 Docker容器化：一劳永逸的方案

说实话，手动配环境太容易出问题了。我现在所有项目都用Docker。为什么？因为你可以把整个环境打包，换台机器直接跑，版本完全一致。

我的Dockerfile长这样：

FROM nvidia/cuda:12.4.0-devel-ubuntu22.04

# 安装基础依赖
RUN apt-get update && apt-get install -y \
    python3 python3-pip git wget \
    && rm -rf /var/lib/apt/lists/*

# 安装NCCL
RUN apt-get install -y libnccl2 libnccl-dev

# 安装PyTorch
RUN pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

# 设置环境变量
ENV NCCL_DEBUG=INFO
ENV NCCL_SOCKET_IFNAME=eth0

WORKDIR /workspace
CMD ["bash"]

3.6 验证环境：跑一个简单的分布式测试

环境搭好了，怎么验证？我一般用这个脚本：

import torch
import torch.distributed as dist

# 初始化进程组
dist.init_process_group(backend='nccl')
local_rank = dist.get_rank()
torch.cuda.set_device(local_rank)

# 创建一个张量
tensor = torch.ones(1).cuda() * local_rank

# 执行all_reduce
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)

# 打印结果
print(f'Rank {local_rank}: tensor = {tensor.item()}')

dist.destroy_process_group()

跑这个脚本时，用 torchrun --nproc_per_node=4 test_dist.py 启动。如果每个rank打印出来的tensor值都是6（0+1+2+3），说明环境没问题。

3.7 本章知识体系

下面这张图，是我梳理的软件环境搭建核心逻辑：

从这张图可以看出，整个环境搭建是分层递进的。底层驱动不对，上层框架跑不起来。我个人建议，如果你刚开始接触多卡训练，直接用Docker镜像，省去手动配环境的痛苦。