LLM Lab 实验手记 03：手写集合通信的自动微分——AllReduce 如何参与反向传播

代码文件：llm_lab/02_nccl/tp_autograd_comm.py

第 2 篇说 DDP "把梯度同步藏进 backward"，本篇直接拆开看：集合通信到底是怎么嵌进自动微分的？答案是 PyTorch 的 torch.autograd.Function——这也是 Megatron-LM 张量并行中 Row Parallel 通信层、以及各类自定义并行算子的标准实现方式。

1. 实验目标

1. 用 torch.autograd.Function 把 all_reduce 包成一个可微算子；
2. 验证"前向通信，反向也通信"这一张量并行核心机制；
3. 用梯度打印亲眼确认：AllReduce.apply(x) 的输入张量 x.grad 等于全卡梯度之和。

2. 代码逐段拆解

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

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)

class AllReduce(torch.autograd.Function):          # ① 自定义算子
    @staticmethod
    def forward(ctx, x):
        x = x.contiguous()                          # ② 保证内存连续
        dist.all_reduce(x)                          # ③ 前向：跨卡归约
        return x

    @staticmethod
    def backward(ctx, grad_output):                 # ④ 反向也要通信！
        grad_output = grad_output.contiguous()
        dist.all_reduce(grad_output)
        return grad_output

x = torch.ones(4, device="cuda", requires_grad=True)  # ⑤ 输入带梯度

y = AllReduce.apply(x)                                # ⑥ 走自定义算子

loss = y.sum()

loss.backward()                                       # ⑦ 反向传播

if rank == 0:
    print("x.grad =", x.grad)

dist.destroy_process_group()

① torch.autograd.Function

autograd.Function 是 PyTorch 的"自定义可微算子"接口：实现静态方法 forward 与 backward，用 Class.apply(...) 调用。自动微分引擎会把 forward 的输出与 backward 串成计算图——我们写的 backward 就决定了梯度如何回流。

② contiguous() 的隐藏陷阱

all_reduce（及多数 NCCL 集合操作）要求张量内存连续。若上游是转置、切片等产生非连续视图的操作，直接通信会报错或静默错乱。contiguous() 是集合通信算子的"保险带"——生产代码里几乎每个通信算子都有它。

③④ 前向归约 + 反向归约

这是本篇的核心洞察，务必读三遍：

• 前向：输入 x 跨卡求和，每卡拿到相同的聚合值；
• 反向：梯度 grad_output 也要跨卡求和，再回传给输入。

为什么反向必须再通信一次？因为前向的 all_reduce 相当于"多卡共享了同一个输出"。若输出对总损失的梯度是 g，那么每张卡上的输入都对这个输出有贡献，各自的梯度都是 g 的完整副本——但每个输入只应得到"自己的那份"。对 SUM 归约，正确语义是：每个输入的梯度 = 输出梯度。然而在多卡自动微分里，各卡会独立计算出同一份 g；若不通信，后续各卡梯度各不相同，同步就断了。所以反向再做一次 AllReduce：让梯度在所有卡上保持一致。

3. 原理讲解：这就是张量并行的"粘合剂"

AllReduce 自动微分算子在两个地方决定性地出现：

1. DDP 的梯度同步（第 2 篇）：每个参数梯度就绪后，等效于执行"梯度 AllReduce"——和这里的 backward 是同一件事。

2. 张量并行 Row Parallel 的跨卡汇总（第 4 篇将展开）：切行权重后，每卡算出部分和，前向 AllReduce 拼出完整输出；反向时梯度 AllReduce 保证权重各卡梯度一致。一个算子，前后双向通信，就把"切开的矩阵"重新粘成完整的计算——这就是 Megatron-LM 里 f（forward）与 b（backward）通信层。

4. 实验结果

以 4 卡运行，x = ones(4)，每卡先 x * rank 无（本实验 x 各卡相同），前向 AllReduce 后 y = ones(4) * 4，loss = y.sum()，loss.backward() 后 x.grad 每卡都应打印：

x.grad = tensor([4., 4., 4., 4.])

为什么是 4 而不是 1：x.grad = 输出梯度（全 1）× 反向 AllReduce 的归约结果。反向再归约一次，把 4 卡各自的梯度（每卡都是全 1）求和成 4——这就是"梯度跨卡一致且等于全局聚合值"的直接证据。

5. 工程启示

1. 自定义并行算子 = 通信 + autograd：凡是"切分后需要同步"的自定义层（张量并行、专家并行、序列并行的通信点），实现模式都是 autograd.Function + 前向/反向各一次集合通信。看懂本篇，就看懂了 Megatron/DeepSpeed 通信层的骨架。

2. 反向通信的成本与优化：一次 AllReduce 反向意味着梯度通信量与数据量成正比。张量并行中大量小通信会拖垮吞吐，工程上通过算子融合（fused op）、通信与计算重叠、以及选择合适的切分方向（Column 切分反向不需要通信）来回避——这也是第 4 篇"为什么 Megatron 把权重按特定方向切"的设计动机。

3. 验证梯度正确性是并行改造的第一关：任何自定义并行层上线前，先跑"N 卡梯度一致性与单卡对照"测试（本实验的 x.grad 打印即最小形态），再谈性能。

下一步：现在可以真正切矩阵了——第 4 篇用 Column/Row 两种切法实现张量并行，并复刻 Megatron 的 FFN：LLM Lab 实验手记 04：张量并行入门。