LLM Lab 实验手记 01：分布式训练第一课——从 init_process_group 到 all_reduce

系列说明：本系列基于作者在 GPU 集群上真实运行过的实验代码（llm_lab/），逐段拆解大模型训练基础设施的核心机制。第 1 篇从分布式训练最不起眼、却决定一切的起点讲起——进程组的初始化。

代码文件：llm_lab/01_ddp/ddp_test.py、ddp_test2.py

1. 实验目标

第一次接触多卡训练的人，第一个问题通常是：多张 GPU 上的程序，是怎么"知道彼此"的？

本实验的目标只有三个：

1. 用最短的代码跑通 NCCL 进程组初始化；
2. 搞清 rank 和 local_rank 到底是什么、为什么必须 set_device(local_rank)；
3. 用 all_reduce 亲眼验证"集合通信 = 多卡数据同步"这件事。

2. 代码逐段拆解

完整代码（ddp_test.py，18 行）：

import os
import torch
import torch.distributed as dist

dist.init_process_group("nccl")          # ① 初始化进程组

rank = dist.get_rank()                   # ② 全局编号
local_rank = int(os.environ["LOCAL_RANK"])   # ③ 节点内编号

torch.cuda.set_device(local_rank)        # ④ 绑定当前卡

x = torch.ones(1).cuda() * rank          # ⑤ 每卡一个不同的数

dist.all_reduce(x)                       # ⑥ 跨卡求和

print(f"Rank {rank}: {x.item()}")        # ⑦ 结果

dist.destroy_process_group()             # ⑧ 清理

① dist.init_process_group("nccl")

这一行是所有多卡 PyTorch 程序的地基。它做了三件事：

• 与所有参与的进程握手，建立进程组（process group）；
• 选定通信后端——nccl 即 GPU 集合通信库（NVIDIA Collective Communications Library），多卡训练的标准后端；
• 从环境变量读取 MASTER_ADDR / MASTER_PORT / WORLD_SIZE / RANK，完成全局拓扑对齐。

一个常见的困惑是"为什么我不传任何参数也能跑？"——因为参数来自环境变量。用 torchrun 启动时，它会自动注入这些变量：

torchrun --nproc_per_node=4 ddp_test.py

②③ rank 与 local_rank

• rank：进程的全局编号，整个训练任务内唯一（0, 1, 2, 3, …），跨节点连续编号。它决定"你是谁"。
• local_rank：进程在当前节点内的编号。它在节点内从 0 重新计数，也是你物理上该使用哪张 GPU 的索引。

⚠️ 关键区别：rank 是通信世界的身份，local_rank 是硬件世界的身份。在多机场景下，rank=8 的进程和 rank=0 的进程都在各自节点的 local_rank=0 上。

④ torch.cuda.set_device(local_rank)

为什么必须写？因为 PyTorch 默认使用 cuda:0。如果不显式切换，8 张卡的进程会全部挤到第 0 张卡上——初始化能过，显存立刻爆。

⑤⑥⑦ ones(1) * rank 与 all_reduce

每张卡构造一个只含自己 rank 值的张量：卡 0 是 [0]，卡 1 是 [1]，卡 2 是 [2]，卡 3 是 [3]。

dist.all_reduce(x) 是集合通信（collective）中最基础的原语：所有进程把自己的 x 贡献出来，做一次全归约，结果广播给所有人。默认归约算子是求和（SUM）。

3. 原理讲解：集合通信与"起跑线"

all_reduce 的语义一句话即可概括：

每个进程输入一份数据，每个进程得到相同的、聚合后的结果。

这是数据并行梯度同步的数学基础：各卡算出的梯度各不相同，all_reduce 把它们求和（或求平均）后广播回去，所有卡就拥有了一致的梯度——这正是 DistributedDataParallel 每步训练背后真正执行的操作。

另一个重要事实：init_process_group 时 NCCL 会建立通信环/树拓扑并做带宽探测，这决定了后续所有集合操作能跑多快。初始化本身也是开销——这也是为什么万卡集群的作业启动（几十万进程同时握手）是系统工程问题。

4. 实验结果

以 4 卡为例，每卡输入 rank 值 [0, 1, 2, 3]，all_reduce(SUM) 后每张卡都应得到：

Rank 0: 6.0
Rank 1: 6.0
Rank 2: 6.0
Rank 3: 6.0

（0 + 1 + 2 + 3 = 6，N 卡时为 N(N−1)/2。这是集合通信确定性的数学结果，可据此核对实验是否跑对。）

如何判断实验成功：四行输出数值一致。只要有一张卡输出不同，说明进程组没对齐或某条链路有问题——在真实集群里，这就是通信故障排查的第一现场。

5. 工程启示

1. NCCL 初始化是万卡训练的"起跑线"：init_process_group 失败/超时的原因通常是 MASTER_ADDR 写错、防火墙拦了 MASTER_PORT、或节点间 RDMA 不可达。集群里海量作业同时启动时，握手风暴本身就能把调度器打挂——大集群常见做法是错峰启动。

2. local_rank 用错是最经典的翻车点：不 set_device(local_rank) 会导致显存抢占；用 rank 当设备号在多机场景会直接越界报错。规范写法永远是 torch.cuda.set_device(local_rank)。

3. 先跑通 18 行的最小实验，再谈优化：生产排障时，我会让 nccl-tests 和这个最小脚本先跑，排除"进程组根本不通"的可能，再进入带宽、拓扑层面的分析——否则调参调的是空气。

4. destroy_process_group 别省略：多轮实验脚本不清理进程组，会残留 NCCL 通信资源，同一节点上连续跑多个实验时可能互相干扰。

下一步：进程组通了，下一步就是把一个真实的模型包上 DDP 跑起来——见 LLM Lab 实验手记 02：DDP 最小训练。