LLM Lab 实验手记 02：DDP 最小训练——一行包装跑通多卡

代码文件：llm_lab/01_ddp/mini_train.py（配套性能版 mini_train_profile.py）。在第 1 篇跑通进程组之后，本实验让一个真实模型在多卡上完成训练循环。

1. 实验目标

1. 把第 1 篇的进程组初始化接进一个真实训练循环；
2. 用 DistributedDataParallel 一行包装，跑通 5 步多卡训练；
3. 拆解 DDP 在"训练循环"背后替我们做了什么。

2. 代码逐段拆解

完整代码（mini_train.py，37 行）：

import os
import torch
import torch.nn as nn
import torch.optim as optim
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)

model = nn.Linear(1024, 1024).cuda()
model = nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])   # ① 一行包装

optimizer = optim.Adam(model.parameters(), lr=1e-3)                           # ② 优化器

x = torch.randn(32, 1024).cuda()
y = torch.randn(32, 1024).cuda()

for step in range(5):                                                         # ③ 训练循环
    optimizer.zero_grad()

    out = model(x)

    loss = ((out - y) ** 2).mean()

    loss.backward()                                                           # ④ 反向
    optimizer.step()                                                          # ⑤ 更新

    if rank == 0:
        print(f"Step {step}, Loss {loss.item()}")

dist.destroy_process_group()

① DistributedDataParallel(model, device_ids=[local_rank])

这是整个实验的核心一行。DDP 初始化时做了两件在分布式训练中至关重要的事：

• 参数广播：rank 0 把自己的初始权重广播给所有进程，保证起点一致；
• ReducHook 注册：在每个参数的梯度上挂一个同步钩子，反向传播时自动执行梯度 AllReduce。

② 优化器放在 DDP 包装之后

一个非常容易踩的坑：必须先包 DDP，再建优化器。因为 model.parameters() 在 DDP 包装后仍然可用，而如果先建优化器再包装，优化器持有的是包装前的参数引用，DDP 的参数广播就无法同步到优化器状态上。

③④⑤ 训练循环里的"隐形通信"

循环看起来和单卡一模一样——这正是 DDP 的设计目标：上层无感，通信藏在 backward 里。loss.backward() 触发时，每个参数的梯度在计算完成后自动执行一次 AllReduce，于是 optimizer.step() 时每张卡拿到的都是全局一致的梯度，更新后的参数自然保持一致。

3. 原理讲解：梯度同步的两种时机

DDP 默认使用异步的梯度 AllReduce：反向传播按参数逆序执行，每个参数梯度就绪后立即开始通信，与后续层的前向计算重叠——这被称为 gradient bucketing（梯度分桶）：小梯度先合并成桶，一次通信处理一个大桶，降低通信次数、提高带宽利用率。

对比朴素做法（同步完所有梯度再统一通信），分桶重叠能让通信耗时几乎"隐身"在计算中——这是 DDP 与早期 DataParallel 拉开吞吐差距的关键原因之一。

4. 实验结果

以 4 卡运行 5 步为例，预期输出形态（rank 0 打印）：

Step 0, Loss 1.0396
Step 1, Loss 0.9983
Step 2, Loss 0.9597
Step 3, Loss 0.9231
Step 4, Loss 0.8885

（损失随步数稳定下降；具体数值因随机初始化而不同，但趋势与量级应一致。）

如何判断实验成功：

• 所有卡训练正常结束且仅 rank 0 打印；
• 损失逐步下降——说明梯度同步链路是通的（若各卡梯度不同步，损失会乱跳甚至发散）；
• 想进一步确认同步正确性，可以在 step 后手动对比各卡的 model.module.weight，4 卡应当逐元素相等。

性能版（mini_train_profile.py，Linear(4096, 4096)）会把每步拆成三段计时：

Step 0
 forward : 3.21 ms
 backward: 5.87 ms
 update  : 0.42 ms

5. 工程启示

1. "一行包装"不等于"零成本"：DDP 的梯度 AllReduce 每步通信量为 2× 模型参数量。对 1024×1024 的小模型无所谓，但对 70B 模型就是每步 ~280 GB 的集群流量——这正是第 11 章《NCCL 网络调优》要解决的事。

2. device_ids 与 CUDA_VISIBLE_DEVICES 的配合：在 Slurm 环境里常先设 CUDA_VISIBLE_DEVICES=$SLURM_LOCALID 再做 set_device(0)；两种写法等价，但不要两套混用，否则设备编号错位是排查半小时起步的坑。

3. 训练损失"乱跳"先怀疑梯度同步：多卡训练发散，第一时间验证各卡梯度一致性（torch.distributed.all_reduce 对比），而不是调学习率——这通常是同步问题而非优化问题。

4. 从 5 步到万卡，框架不变，工程变：mini_train.py 里的循环结构与生产训练框架（Megatron-LM、DeepSpeed）逐行对应，差异只在并行策略、通信重叠、Checkpoint 与故障恢复的工程深度上。

下一步：DDP 的梯度同步靠 AllReduce——但 AllReduce 是怎么"参与反向传播"的？第 3 篇手写一个 AllReduce 自动微分函数来验证：LLM Lab 实验手记 03：手写集合通信的自动微分。