将 Adam 梯度同步（Reduce-Scatter）移入反向 Hook：modded-nanogpt 分布式优化器每步提速约 0.7 秒的工程实践

• 人工智能
• 大模型
• 预训练
• 分布式训练
• 模型优化
• 深度学习

本篇文章以 modded-nanogpt 仓库 track_1_short 2025-10-31 记录 为核心，剖析一次纯性能优化改动：把DistAdam优化器step()内的梯度 reduce-scatter 集体通信操作提前到反向传播的register_post_accumulate_grad_hook中，并配合step()逆序迭代参数，让后层参数的梯度同步更早启动、更早完成，最终将一次 8xH100 训练运行的总耗时缩短约 0.7 秒。读完本文，你将理解 PyTorch 分布式优化器里"异步通信 + 依赖顺序"的调度思想，并掌握用 profiler trace 验证通信-计算重叠的方法。

背景：为什么梯度同步会成为每步的瓶颈

在 modded-nanogpt 的 speedrun 训练（torchrun --standalone --nproc_per_node=8 train_gpt.py，目标是将 GPT-2 small 在 8 张 H100 上训练到 3.28 FineWeb val loss）中，模型参数以分片（sharded）方式分布在多个 GPU 上，每个优化器步骤都需要：

1. reduce-scatter：把各 rank 上完整梯度聚合、按 rank 切片分发回各 GPU（每个 rank 只持有参数的一个纵向分片）；
2. 用本地梯度分片做 Adam 更新；
3. all-gather：把更新后的参数分片收集还原成完整参数。

尽管dist.reduce_scatter_tensor(..., async_op=True)返回的是异步操作，旧实现中这些 collectives 都集中在每个训练步的末尾触发。step()按参数顺序逐个启动 reduce-scatter 后立即wait()，于是通信与主 GPU 流上的计算几乎没有重叠，通信延迟完全暴露在关键路径上。在 8 卡规模下，这一串串行的通信操作占用的墙钟时间相当可观，是"90 秒训练 124M 模型"这种极限场景里必须抠掉的每一毫秒之一。

改动一：把 Reduce-Scatter 从 step() 移入 Backward Hook

该 PR 的核心理念：既然 reduce-scatter 只依赖梯度就绪，而梯度的就绪顺序天然由反向传播决定（后层参数先得到梯度），那就让"启动 reduce-scatter"这个动作顺着反向传播的节奏去执行，而不是等 forward/backward 全部结束后再在优化器里统一发起。

实现上，DistAdam.__init__里为每个参数注册了钩子（对应记录目录下训练脚本快照中的源码，如 b725e7bc…txt 的DistAdam类）：

self.should_sync = False self._reduce_scatter_hooks = [] self._reduce_scatter_futures = {} self.register_backward_hooks() def register_backward_hooks(self): for group in self.param_groups: params: list[Tensor] = group["params"] for param in params: hook = param.register_post_accumulate_grad_hook(self._sync_gradient) self._reduce_scatter_hooks.append(hook) @torch.compile @torch.no_grad() def _sync_gradient(self, param): if not self.should_sync: return grad = param.grad rank_size = grad.shape[0] // self.world_size grad_slice = torch.empty_like(grad[:rank_size]) self._reduce_scatter_futures[param] = ( dist.reduce_scatter_tensor(grad_slice, grad, op=dist.ReduceOp.AVG, async_op=True).get_future(), grad_slice )

几个值得注意的实现细节：

• register_post_accumulate_grad_hook保证梯度就绪：它只在参数的梯度完成累加（post-accumulate）之后触发，因此钩子内读取的param.grad是最终可用于同步的完整梯度，不会出现只同步了部分微批梯度的竞态；
• should_sync开关：用于训练早期（warmup、graph 捕获阶段）或特定步骤跳过同步，避免钩子在不需要同步梯度时执行不必要的启动动作；
• async_op=True+get_future()：reduce-scatter 以异步方式启动，返回的 future 与本地梯度切片一起存入_reduce_scatter_futures字典，供step()稍后按参数逐个wait()；
• @torch.compile编译钩子函数：_sync_gradient被 torch.compile 处理，减少每次钩子调用的 Python 开销，这也是 speedrun 场景下所有热点代码的通用处理方式。

由于反向传播（loss.backward()）会从输出层向输入层依次执行，后层参数对应的钩子先被调用，它们的 reduce-scatter 也就先被启动——这正是接下来step()改动要利用的顺序信息。

改动二：step() 逆序遍历 Param Groups 与参数

把 reduce-scatter 提前到钩子里之后，step()不再负责启动通信，而只需要"等待-更新-再收集"。为了让等待动作也能利用"后层先同步完成"的顺序，DistAdam.step()改为逆序遍历 param_groups 和组内参数：

@torch.compile @torch.no_grad() def step(self): rank = dist.get_rank() all_gather_futures: list[torch.Future] = [] for group in reversed(self.param_groups): beta1, beta2 = group['betas'] eps = group['eps'] wd = group['weight_decay'] for param in reversed(group['params']): if param not in self._reduce_scatter_futures: continue fut, g_slice = self._reduce_scatter_futures[param] fut.wait() rank_size = param.shape[0] // self.world_size p_slice = param[rank * rank_size:(rank + 1) * rank_size] lr = group['lr'] * getattr(param, "lr_mul", 1.0) state = self.state[param] exp_avg = state["exp_avg"] exp_avg_sq = state["exp_avg_sq"] state["step"] += 1 t = state["step"] # weight decay if wd != 0: eff_weight_decay = lr * wd * getattr(param, "wd_mul", 1.0) p_slice.mul_(1 - eff_weight_decay) # update running averages exp_avg.mul_(beta1).add_(g_slice, alpha=1 - beta1) exp_avg_sq.mul_(beta2).addcmul_(g_slice, g_slice, value=1 - beta2) # bias corrections bias1 = 1 - beta1 ** t bias2 = 1 - beta2 ** t # compute step denom = exp_avg_sq.sqrt().add_(eps) step_size = lr * (bias2 ** 0.5 / bias1) update = exp_avg.div(denom).mul_(step_size) p_slice.add_(other=update, alpha=-1.0) all_gather_futures.append(dist.all_gather_into_tensor(param, p_slice, async_op=True).get_future()) self._reduce_scatter_futures.clear() torch.futures.collect_all(all_gather_futures).wait()

为什么逆序有效，README 给出了两条依据：

1. 参数层面：DistAdam.__init__中按参数张量的 shape 建立 param_groups（每个 shape 一组），而 GPT 模型中后层（更靠近输出）的参数位于参数列表末尾。反向传播先为它们计算梯度、钩子先为它们启动 reduce-scatter，因此它们的 future 也最早完成。step()逆序遍历时先wait()这些最早完成的 future，等待时间最短，后续计算也能更早与通信并行；
2. 组层面：param_groups 中第一个 group 对应第一次遇到的 shape（即靠近输入层的形状），后层参数落在更靠后的 group 里，所以reversed(self.param_groups)保证先处理后层参数的 group。

最终效果是：通信的启动顺序（由反向传播驱动）与消费顺序（由 step 逆序驱动）对齐，通信完成的等待链不再出现"先等一个晚启动、后等一个早启动"的错位，整条等待关键路径被压缩。

收益与验证：0.7 秒与显著性检验

README 给出了同机对比数据。本 PR 实施后：

import scipy.stats import torch losses = [3.2775, 3.2776, 3.2777, 3.2780, 3.2781, 3.2775, 3.2786, 3.2774, 3.2751, 3.2739] times = [140.909, 140.872, 140.743, 140.743, 140.747, 140.809, 140.728, 140.784, 140.862, 140.934] print("p=%.4f" % scipy.stats.ttest_1samp(losses, 3.28, alternative="less").pvalue) # p=0.0001 print("losses:", torch.std_mean(torch.tensor(losses))) # losses: (std=0.0015, mean=3.2771) print("time:", torch.std_mean(torch.tensor(times))) # time: (std=0.0760, mean=140.8131)

上一 PR 在同一台机器上的基线：

import scipy.stats import torch times = [141.654, 141.413, 141.467, 141.516] print("time:", torch.std_mean(torch.tensor(times))) # time: (std=0.1033, mean=141.5125)

可解读为：

• 最终训练时间从约141.51 秒降至约 140.81 秒，改善约 0.7 秒；
• 最终 loss 均值 3.2771（std 0.0015），对目标 3.28 做单侧单样本 t 检验 p=0.0001，说明该改动没有以牺牲收敛为代价——纯粹是系统层面的提速；
• 计时标准差从 0.1033 缩小到 0.0760，运行稳定性也有提升。

这是一个"零算法改动、纯工程收益"的典型案例：模型、数据、超参都不变，只调整通信的调度时机，就白赚了约 0.7 秒。

Profiler Trace 分析：通信与计算的重叠

README 将改动前后的两套 trace 文件随记录一起提交，并注明可用 perfetto trace viewer 打开查看（注意：trace 为外部工具格式，这里仅按仓库记录说明其结论）：

• 当前实现（改动前）：modded-nanogpt-current-gpu-rank-00-chrome-trace.json.gz
• Hook 实现（改动后）：modded-nanogpt-hook-gpu-rank-00-chrome-trace.json.gz

改动前（Current Implementation）：第一个 reduce-scatter 在DistAdam.step()开始时才启动。观察 GPU 流可以发现，初始 reduce-scatter 与主 GPU 流的计算没有重叠——通信是在整段计算结束后的串行时段里执行的。

改动后（Hook Implementation）：第一个 reduce-scatter 由第一个（后层参数的）钩子启动，时间点大幅提前。在 GPU 流视图上可以看到 reduce-scatter 与主 GPU 流上的计算产生了重叠——反向传播还在进行时，后层参数的梯度同步已经并行展开。

这正是提速的机理：通信从"计算结束后的串行尾巴"变成了"计算进行中的并行流水"。训练时总耗时由最长的一条链决定，当 NCCL 的 reduce-scatter 与反向传播的算子在同一时间窗内执行时，关键路径被显著压缩。

实现要点与可复用的工程结论

综合 README 与记录目录中训练脚本快照的源码实现（DistAdam完整实现可见 8fb212eb…txt、b725e7bc…txt 等文件），这套手法可以总结为几条可复用的经验：

1. 异步 collectives 不等于自动重叠：async_op=True只是把"等待"从启动点推迟到你wait()的地方；真正决定重叠与否的是启动时机。把启动尽可能提前到依赖满足的最早时刻（这里是梯度就绪的瞬间），才能让通信躲进计算的阴影里；
2. 反向传播本身是天然的顺序信号：register_post_accumulate_grad_hook不仅保证梯度完整，还免费提供了"后层先就绪"的时序，适合做梯度类通信的调度；
3. 消费顺序要与生产顺序对齐：step()逆序遍历、先 wait 后层参数，避免"等待链错位"，让每个wait()都尽量落在对应 future 已完成的时刻；
4. 用统计检验证明性能结论：多跑几次取均值/方差、对 loss 做 t 检验，确认提速不是噪声、也没有牺牲质量——这是记录中体现的方法论，也是可持续复现优化的前提；
5. 用 profiler trace 佐证机理：光有墙钟时间不够，还要能解释"快在哪"——trace 里看流的重叠关系是标准手段，两个版本的 trace 文件都随记录提交，方便对照复现。

从源码看后续演进

值得一提的是，这一思路在后来的版本中继续演进。当前仓库的优化器实现在 track_1_short/optim/anvil.py（AnvilAndAdam），其文档字符串明确写道："Gradient communication is explicitly scheduled rather than hook-driven"（梯度通信采用显式调度而非钩子驱动），即通信改由scatter_order显式编排、更新在work_order中执行（见 anvil.py 与_launch_reduce相关实现）。而训练入口 train_gpt.py 中TrainingManager.step_optimizers负责在固定 cadence 上推进优化器（Adam 参数只在奇数步更新，is_adam_step，见 track_1_short/training.py）。

从代码结构看，这种演进是自然的：随着 ANVIL 优化器引入多路（twin-rail）通信、bank 更新和更复杂的参数分组，钩子驱动的隐式顺序难以覆盖所有调度约束，显式调度成为更可控的方案。但本记录的价值在于它清晰地演示了底层原理——只要把通信启动提前到依赖就绪的瞬间，并让消费顺序匹配生产顺序，就能从系统层面白赚时间——这一原则在显式调度版本中同样成立。

如果你在自己的分布式训练里遇到"通信总是出现在计算尾部"的瓶颈，不妨从本记录出发：找到梯度就绪的最早时机（post-accumulate hook 是一个很好的切入点），把 reduce-scatter 的启动提前，再用 profiler 确认重叠是否真实发生。这套"启动提前 + 顺序对齐 + trace 验证"的组合拳，比盲目加大带宽或换用 ring 算法更便宜、也更立竿见影。

• 人工智能
• 大模型
• 预训练
• 分布式训练
• 模型优化
• 深度学习