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

人工智能大模型预训练分布式训练模型优化深度学习【免费下载链接】modded-nanogptNanoGPT (124M) in 90 seconds项目地址https://gitcode.com/GitHub_Trending/mo/modded-nanogpt点击查看免费下载本篇文章以 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_node8 train_gpt.py目标是将 GPT-2 small 在 8 张 H100 上训练到 3.28 FineWeb val loss中模型参数以分片sharded方式分布在多个 GPU 上每个优化器步骤都需要reduce-scatter把各 rank 上完整梯度聚合、按 rank 切片分发回各 GPU每个 rank 只持有参数的一个纵向分片用本地梯度分片做 Adam 更新all-gather把更新后的参数分片收集还原成完整参数。尽管dist.reduce_scatter_tensor(..., async_opTrue)返回的是异步操作旧实现中这些 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, opdist.ReduceOp.AVG, async_opTrue).get_future(), grad_slice )几个值得注意的实现细节register_post_accumulate_grad_hook保证梯度就绪它只在参数的梯度完成累加post-accumulate之后触发因此钩子内读取的param.grad是最终可用于同步的完整梯度不会出现只同步了部分微批梯度的竞态should_sync开关用于训练早期warmup、graph 捕获阶段或特定步骤跳过同步避免钩子在不需要同步梯度时执行不必要的启动动作async_opTrueget_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, alpha1 - beta1) exp_avg_sq.mul_(beta2).addcmul_(g_slice, g_slice, value1 - 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_(otherupdate, alpha-1.0) all_gather_futures.append(dist.all_gather_into_tensor(param, p_slice, async_opTrue).get_future()) self._reduce_scatter_futures.clear() torch.futures.collect_all(all_gather_futures).wait()为什么逆序有效README 给出了两条依据参数层面DistAdam.__init__中按参数张量的 shape 建立 param_groups每个 shape 一组而 GPT 模型中后层更靠近输出的参数位于参数列表末尾。反向传播先为它们计算梯度、钩子先为它们启动 reduce-scatter因此它们的 future 也最早完成。step()逆序遍历时先wait()这些最早完成的 future等待时间最短后续计算也能更早与通信并行组层面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, alternativeless).pvalue) # p0.0001 print(losses:, torch.std_mean(torch.tensor(losses))) # losses: (std0.0015, mean3.2771) print(time:, torch.std_mean(torch.tensor(times))) # time: (std0.0760, mean140.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: (std0.1033, mean141.5125)可解读为最终训练时间从约141.51 秒降至约 140.81 秒改善约 0.7 秒最终 loss 均值 3.2771std 0.0015对目标 3.28 做单侧单样本 t 检验 p0.0001说明该改动没有以牺牲收敛为代价——纯粹是系统层面的提速计时标准差从 0.1033 缩小到 0.0760运行稳定性也有提升。这是一个零算法改动、纯工程收益的典型案例模型、数据、超参都不变只调整通信的调度时机就白赚了约 0.7 秒。Profiler Trace 分析通信与计算的重叠README 将改动前后的两套 trace 文件随记录一起提交并注明可用 perfetto trace viewer 打开查看注意trace 为外部工具格式这里仅按仓库记录说明其结论当前实现改动前modded-nanogpt-current-gpu-rank-00-chrome-trace.json.gzHook 实现改动后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 等文件这套手法可以总结为几条可复用的经验异步 collectives 不等于自动重叠async_opTrue只是把等待从启动点推迟到你wait()的地方真正决定重叠与否的是启动时机。把启动尽可能提前到依赖满足的最早时刻这里是梯度就绪的瞬间才能让通信躲进计算的阴影里反向传播本身是天然的顺序信号register_post_accumulate_grad_hook不仅保证梯度完整还免费提供了后层先就绪的时序适合做梯度类通信的调度消费顺序要与生产顺序对齐step()逆序遍历、先 wait 后层参数避免等待链错位让每个wait()都尽量落在对应 future 已完成的时刻用统计检验证明性能结论多跑几次取均值/方差、对 loss 做 t 检验确认提速不是噪声、也没有牺牲质量——这是记录中体现的方法论也是可持续复现优化的前提用 profiler trace 佐证机理光有墙钟时间不够还要能解释快在哪——trace 里看流的重叠关系是标准手段两个版本的 trace 文件都随记录提交方便对照复现。从源码看后续演进值得一提的是这一思路在后来的版本中继续演进。当前仓库的优化器实现在 track_1_short/optim/anvil.pyAnvilAndAdam其文档字符串明确写道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 算法更便宜、也更立竿见影。赞分享人工智能大模型预训练分布式训练模型优化深度学习【免费下载链接】modded-nanogptNanoGPT (124M) in 90 seconds项目地址https://gitcode.com/GitHub_Trending/mo/modded-nanogpt点击查看免费下载相关推荐FlexGen 技术解读ZeRO 梯度通信优化中的 reduce-scatter 与 Partition-Aware 梯度归约FlexGen 技术解读ZeRO 梯度通信优化中的 reduce scatter 与 Partition Aware 梯度归约导读本文聚焦 DeepSpe推理引擎大模型modded-nanogpt 稀疏化 bigram 嵌入梯度通信all_to_all 稀疏 reduce-scatter 与 CPU 流水线重叠实战modded nanogpt 稀疏化 bigram 嵌入梯度通信all_to_all 稀疏 reduce scatter 与 CPU 流水线重叠实战导读本人工智能大模型预训练分布式训练模型优化深度学习PyTorch Lightning 模型训练加速完全指南从分布式加速器到梯度同步优化的实践方案PyTorch Lightning 模型训练加速完全指南从分布式加速器到梯度同步优化的实践方案本指南以 PyTorch Lightning 官方文档 spe人工智能深度学习机器学习预训练分布式训练微调上一篇PCV库与OpenCV对比为什么这款纯Python视觉库更适合初学者下一篇BlurAdmin中的测试驱动开发从单元测试到E2E测试完整流程创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

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

相关推荐

EDK II EmulatorPkg 平台 CI 构建指南：基于 Pytools 的本地构建与模拟器自动化运行

himalaya 邮件发送副本保存：`message.send.save-copy` 配置解析与“先发送后保存“语义

工业设备非易失存储方案：基于R7FA4E2B93CFM与MR25H40CDF MRAM的日志设计

更多资讯

接口自动化测试数据构造全攻略：从静态数据到Mock与加密

硬盘修复后文件消失？底层原理与数据恢复实战指南

钢材缺陷检测三格式标签实战：VOC/COCO/YOLO工业级打通指南

Linux系统管理实战：从用户权限到运维脚本的完整进阶指南

GeM池化：可学习幂参数的图像特征聚合方法

电影评论情感分析系统实战：BiLSTM训练与Flask部署

今日更新

CMOS设计三大物理效应：体效应、沟道长度调制与亚阈值导电

微信小程序在线阅读毕设全解析：Java后端+MySQL+避坑实战

Android音频问题五层分析法：从App到Codec的系统化诊断

本周精选

CMOS设计三大物理效应：体效应、沟道长度调制与亚阈值导电

微信小程序在线阅读毕设全解析：Java后端+MySQL+避坑实战

Android音频问题五层分析法：从App到Codec的系统化诊断

本月热门

我发现了一个新思路：用 Remotion + Claude Code 像写代码一样自动化生成短视频

Windows下 Codex 中 Chrome 和 Computer Use 插件不可用问题排查及解决参考方式：TaoToken 统一 Key 配置与验证

2026 大模型集体涨价：用 Python 做企业 Token 成本测算与选型避坑（附配置）

读完文章，想聊聊您的网站？