单台服务器上的多 GPU：拓扑、NCCL 与 torchrun

先查看拓扑

GPU 之间的连接方式决定了哪种并行方式成本更低。选择之前先看一看：

nvidia-smi topo -m        # GPU、网卡与 CPU 之间的连接矩阵
nvidia-smi topo -p2p r    # 哪些 GPU 对支持点对点读取
amd-smi topology          # 适用于 AMD MI355X

CPU Affinity 和 NUMA Affinity 列显示了哪些核心离每块 GPU 最近。当 CPU 侧成为瓶颈时，把数据加载工作绑定到这些核心上。

用 torchrun 启动 2、4 或 8 GPU 任务

torchrun 会为每块 GPU 启动一个进程，并为每个进程分配 RANK、LOCAL_RANK 和 WORLD_SIZE。在开始长时间任务之前，先用一个简短脚本测试配置：

# ddp_check.py
import os
import torch
import torch.distributed as dist

local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
dist.init_process_group("nccl")

x = torch.ones(1, device="cuda") * dist.get_rank()
dist.all_reduce(x)          # 对所有 GPU 上的 rank 编号求和
print(f"rank {dist.get_rank()} of {dist.get_world_size()}: sum = {x.item():.0f}")
dist.destroy_process_group()

torchrun --standalone --nproc-per-node=gpu ddp_check.py   # 每个可见 GPU 一个进程
torchrun --standalone --nproc-per-node=4 train.py         # 四个进程
CUDA_VISIBLE_DEVICES=4,5,6,7 torchrun --standalone --nproc-per-node=4 train.py   # 仅使用 4 到 7 号 GPU

每个 rank 都应打印出相同的总和 0 + 1 + … + (n − 1)：8 块 GPU 上的结果是 28。CUDA_VISIBLE_DEVICES 可以让两个独立任务共享一台服务器，各自使用自己的 GPU。torchrun 会为每个进程设置 OMP_NUM_THREADS=1，除非您自己另行设置；如果训练进程内的 CPU 工作较慢，可以调高这个值。

NCCL 要点

NCCL 是 PyTorch 的 nccl 后端背后所用的库。它会自动检测拓扑，并自行在 NVLink、PCIe 点对点、共享内存或网络之间做出选择，默认设置就能满足大多数任务。以下变量可以显示并调试它的选择结果：

NCCL_DEBUG=INFO torchrun --standalone --nproc-per-node=gpu ddp_check.py 2>&1 | grep -E 'via (P2P|SHM|NET)' | head

用 nccl-tests 测量总线带宽

nccl-tests 用于测量互联实际能提供的性能。它需要来自 NVIDIA 仓库的 CUDA 工具包和 NCCL 开发文件（参见 NVIDIA 驱动与 CUDA）；NGC 的 PyTorch 容器已经自带这两者。

sudo apt install -y libnccl2 libnccl-dev
git clone https://github.com/NVIDIA/nccl-tests.git
cd nccl-tests
make -j                                          # CUDA_HOME 默认为 /usr/local/cuda
./build/all_reduce_perf -b 8 -e 2G -f 2 -g 8     # -g：服务器中的 GPU 数量

-b 和 -e 设置最小和最大的消息大小，-f 设置不同大小之间的倍数，-g 设置 GPU 数量。每一行输出对应一种消息大小：

• algbw：数据大小除以耗时。
• busbw：按 GPU 数量修正后的 algbw（对于 all-reduce，乘以 2(n − 1)/n）。它反映的是路径中最慢链路（NVLink、PCIe 或 CPU 插槽之间的链路）的速度，是在不同服务器之间进行比较的指标。
• #wrong 必须始终为 0。运行结束时会输出一行 Avg bus bandwidth。

请关注大消息的行。在服务器状态正常时先记录一次结果，以便日后性能发生变化时可以对比。

PCIe 显卡与 NVLink 平台

• 8 块 GPU 的 H100、H200、B200 和 B300 服务器（HGX）：每块 GPU 都通过 NVSwitch 以完整的 NVLink 速度连接到其他每一块 GPU，Hopper 上每块 GPU 为 900 GB/s，Blackwell 上为 1.8 TB/s。八块 GPU 一起做张量并行效果很好。必须运行 Fabric Manager（参见 NVIDIA 驱动与 CUDA）。在这些型号的 2 GPU 和 4 GPU 服务器上，用 nvidia-smi topo -m 查看 GPU 之间的连接方式。
• MI355X：GPU 之间通过 AMD Infinity Fabric 连接，RCCL 会使用它（见下文）。
• RTX PRO 6000 和 L40S：仅支持 PCIe（Gen5 和 Gen4）。GPU 之间能否进行点对点数据交换取决于显卡型号和平台；可以用 nvidia-smi topo -p2p r 检查。
• GeForce RTX 5090、4090 和 5080：仅支持 PCIe，没有 NVLink，点对点通常也不可用，因此 NCCL 会通过主机内存传输数据。这样也能工作，只是通信开销更大。

在 PCIe 服务器上，优先使用数据并行（DDP 或 FSDP），并让每一步的计算量足够大，同时用梯度累积来降低通信频率。张量并行在每一层都要交换数据，最能从 NVLink 中获益。

如果任务在 PCIe 服务器上的第一次集合通信操作处挂起，先用 NCCL_P2P_DISABLE=1 运行一次。如果这样能解决问题，说明点对点路径存在问题：请通过联系表单反馈。

更大的模型：FSDP 与 DeepSpeed

• DDP（DistributedDataParallel）：每块 GPU 都保存一份完整的模型和优化器状态副本。当一块 GPU 的显存足以容纳全部内容时，这种方式最简单也最快。
• FSDP2（torch.distributed.fsdp.fully_shard）：将参数、梯度和优化器状态分片到各个 GPU 上，因此模型只需要在所有 GPU 上合计放得下即可。启动方式与 DDP 一样使用 torchrun；参见 PyTorch 的 FSDP2 教程。
• DeepSpeed ZeRO：第 1 到第 3 阶段依次对优化器状态、梯度、参数进行分片，并可选择将部分内容卸载到 CPU 内存或 NVMe。pip install deepspeed；参见入门指南。
• 张量并行与流水线并行（例如 Megatron 风格）：将各层或若干层拆分到不同 GPU 上。用于最大规模的模型；张量并行需要高速互联。

Hugging Face Accelerate 和 Transformers 的 Trainer 会把 DDP、FSDP 和 DeepSpeed 封装在一个配置文件之后。如果是做推理而不是训练，请参见推理服务。

AMD MI355X：RCCL

在 MI355X 上，同样的代码无需修改即可运行。PyTorch 的 nccl 后端使用的是 RCCL，也就是 AMD 对 NCCL API 的实现；torchrun 的用法完全相同，RCCL 读取的也是同样的 NCCL_* 变量，例如 NCCL_DEBUG=INFO。可以用 HIP_VISIBLE_DEVICES 选择 GPU；也支持 CUDA_VISIBLE_DEVICES。

amd-smi topology                           # GPU 之间的连接
git clone https://github.com/ROCm/rccl-tests.git
cd rccl-tests
make GPU_TARGETS=gfx950                    # 如果 ROCm 不在 /opt/rocm，请添加 HIP_HOME=...
./build/all_reduce_perf -b 8 -e 2G -f 2 -g 8

构建 rccl-tests 最简单的方式是使用 rocm/pytorch 容器，其中已经安装好了 ROCm 和 RCCL。安装步骤和 GPU 选择方式参见 MI355X 上的 ROCm。