LLM Lab 实验手记 01：分布式训练第一课——从 init_process_group 到 all_reduce
系列说明：本系列基于作者在 GPU 集群上真实运行过的实验代码（
llm_lab/），逐段拆解大模型训练基础设施的核心机制。第 1 篇从分布式训练最不起眼、却决定一切的起点讲起——进程组的初始化。代码文件：
llm_lab/01_ddp/ddp_test.py、ddp_test2.py
1. 实验目标
第一次接触多卡训练的人，第一个问题通常是：多张 GPU 上的程序，是怎么"知道彼此"的？
本实验的目标只有三个：
- 用最短的代码跑通 NCCL 进程组初始化；
- 搞清
rank和local_rank到底是什么、为什么必须set_device(local_rank)； - 用
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. 工程启示
NCCL 初始化是万卡训练的"起跑线"：
init_process_group失败/超时的原因通常是MASTER_ADDR写错、防火墙拦了MASTER_PORT、或节点间 RDMA 不可达。集群里海量作业同时启动时，握手风暴本身就能把调度器打挂——大集群常见做法是错峰启动。local_rank用错是最经典的翻车点：不set_device(local_rank)会导致显存抢占；用rank当设备号在多机场景会直接越界报错。规范写法永远是torch.cuda.set_device(local_rank)。先跑通 18 行的最小实验，再谈优化：生产排障时，我会让
nccl-tests和这个最小脚本先跑，排除"进程组根本不通"的可能，再进入带宽、拓扑层面的分析——否则调参调的是空气。destroy_process_group别省略：多轮实验脚本不清理进程组，会残留 NCCL 通信资源，同一节点上连续跑多个实验时可能互相干扰。
下一步：进程组通了，下一步就是把一个真实的模型包上 DDP 跑起来——见 LLM Lab 实验手记 02：DDP 最小训练。