modded-nanogpt 记录解读：统一零拷贝优化器与转置 LM Head —— 从梯度堆叠到显式通信调度

本指南基于 modded-nanogpt 仓库中的实验记录 2026-01-18_UnifiedOptimizers/README.md，系统讲解三项核心技术：面向单参数配置的统一优化器 NorMuonAndAdam、基于"参数银行"消除 memcpy 的零拷贝梯度通信，以及通过转置内存布局消除 LM Head 慢速梯度累加核函数。读完本文，你将理解这些改动如何在不改变数学与精度的前提下压缩 8×GPU 训练的单步耗时，并看到它们在后继源码（AnvilAndAdam、qk_bank/vo_bank/mlp_bank 银行、Triton 转置内核）中的延续与演进。

背景：一项"不改变数学"却减掉 10 步的优化

该记录（record 2026-01-18）以 record 59 为开发基线，在合入 record 60（logit softcap kernel）并重新计时后，总训练步数减少了 10 步（单独对比 record 59 基线时减少 5 步）。需要特别强调的是：本次 PR 的改动不应改变任何数学或精度，但实验显示它在多台机器上带来一致的 loss 下降，从而允许删除若干训练步。作者自己也无法完全解释这一现象，推测与"用显式通信/优化工作顺序替代梯度累加 hook"有关。

与 record 60 合入后的对比（8 次运行，Time 单位秒，Loss 为验证损失）：

                   Runs   Time μ  Time σ  Time +/-  Time p  Loss μ  Loss σ  Loss +/-      p
baseline              8 105.5361  0.0923    0.0000     NaN  3.2783  0.0008    0.0000 0.0003 # At 1775 steps
unified-optimizer     8 104.4579  0.1430   -1.0782  0.0000  3.2782  0.0015   -0.0001 0.0060 # At 1765 steps

baseline:
  losses = [3.2777, 3.2785, 3.2790, 3.2785, 3.2788, 3.2791, 3.2766, 3.2784]
  times  = [105.3690, 105.5950, 105.4140, 105.5520, 105.5610, 105.5820, 105.6090, 105.6070]
  (n = 8)

unified-optimizer:
  losses = [3.2781, 3.2803, 3.2767, 3.2799, 3.2794, 3.2767, 3.2767, 3.2780]
  times  = [104.6490, 104.4420, 104.2430, 104.4540, 104.4860, 104.3220, 104.6550, 104.4120]
  (n = 8)

以 record 59（MLP kernel）为基线、减少 5 步的对比（10 次运行）：

                 Runs   Time μ  Time σ  Time +/-  Loss μ  Loss σ  Loss +/-      p
Baseline            8 107.0998  0.0763    0.0000  3.2792  0.0017    0.0000 0.1039  # At 1775 steps
unified-optimizer  10 106.4940  0.0245   -0.6058  3.2790  0.0006   -0.0001 0.0002  # At 1770 steps

记录作者还额外把 triton_kernels.py 的改动一并纳入，并在 train_gpt.py 顶部加了一行代码，将 kernel 源码直接拉进日志文件——这一点可以在记录目录下的 unified-optimizer/2fc79469-a527-4bde-8540-8426ed3352d1.txt（即该记录版本的 train_gpt.py）前 9 行得到印证。

1. 统一零拷贝优化器

1.1 堆叠（Stacking）vs. 分片（Sharding）：参数银行

旧方案的问题

此前 Muon 参数的通信流程是：

1. 将不同的 attention / MLP 矩阵分配到不同 GPU；
2. 把每块 GPU 上的梯度**堆叠（stack）**到一起；
3. **散射（scatter）堆叠结果，做 Muon 数学，再收集（gather）**更新；
4. 把收到的更新**解堆叠（unstack）**并应用到本地参数。

这套流程的缺点很直接：

• 堆叠/解堆叠需要耗时的 memcpy，使 comms 与 compute 难以完全重叠（尽管当时已有一个 PR 做到了）；
• 它强依赖 attention 权重矩阵与 MLP 权重具有相同的大小、形状和朝向；
• 它跑两趟：第一趟处理全部 attention 权重 + 部分 MLP 权重，第二趟只处理 MLP。attention 矩阵处理更快，但两趟都被 MLP 的工作长度拖住。

参数银行（Parameter Banks）

该 PR 采用"参数银行"策略——该策略此前已用于 attention gates 与 ve gates，此处扩展到 attention 与 MLP 投影矩阵：

• 全部 10 个 attention 模块的 QKVO 权重矩阵存为单个参数，形状 (10, 4×768, 768)；无需内存移动即可"reshape"成 (40, 768, 768)，然后均匀切分到 8 块 GPU（每块 5 个 attention 矩阵）。
• 11 个 MLP 模块、每模块 2 个矩阵，需要 padding 才能均匀切分。参数银行形状为 (12, 2, 3072)，reshape 为 (24, 3072) 后每块 GPU 处理 3 个矩阵（注意：这同时意味着要为多出的整整一个 MLP 支付通信与优化器开销，记录作者在 Ideas 一节探讨了如何利用这批空闲权重）。

改组后两趟负载均衡：第一趟每块 GPU 只处理 (5, 768, 768) 的 attention 权重，明显更快；第二趟每块 GPU 处理 3 个 MLP 矩阵（此前为 2 个），会慢一些。

当前主线的实现延续了该思想并有更完整的形状定义：track_1_short/model/gpt.py 中 qk_bank 为 (num_qk_padded, 2*head_dim, model_dim)、vo_bank 为 (num_vo_padded, hdim, model_dim)，而 mlp_bank 明确为 (12, 2, 2816, 768)、reshape 为 (24, 2816, 768)——slot 11 被并行 MLP 使用，slot 7 保持死亡（NO_MLP_LAYERS）但使银行保持偶数，这与记录中"padding 以保证均匀切分"的描述一一对应。

1.2 独立参数：用 param_table 取代参数组

参数银行的另一个收益是消除"参数组（parameter groups）"概念。组存在的意义是处理"同一功能矩阵的多个实例"（如 10 层的 attention 权重共享学习率），但它给优化器代码增加了一层抽象与复杂度。有了参数银行，模型不再有逐层参数，只剩下表中 13 个截然不同的参数（定义于 TrainingManager 的 init 函数）：

# - Ordering dictates when to launch reduce/reduce_scatter operations
# - "sharded" parameters use reduce_scatter/all_gather and "replicated" ones use all_reduce
self.param_table = {
    "attn":           {"optim": "normuon", "comms": "sharded",     "adam_betas": None},
    "mlp":            {"optim": "normuon", "comms": "sharded",     "adam_betas": None},
    "scalars":        {"optim": "adam",    "comms": "replicated",  "adam_betas": [0.9,  0.99]},
    "ve0":            {"optim": "adam",    "comms": "sharded",     "adam_betas": [0.75, 0.95]},
    "ve1":            {"optim": "adam",    "comms": "sharded",     "adam_betas": [0.75, 0.95]},
    "ve2":            {"optim": "adam",    "comms": "sharded",     "adam_betas": [0.75, 0.95]},
    "smear_gate":     {"optim": "adam",    "comms": "replicated",  "adam_betas": [0.9,  0.99]},
    "skip_gate":      {"optim": "adam",    "comms": "replicated",  "adam_betas": [0.9,  0.99]},
    "attn_gate_bank": {"optim": "adam",    "comms": "replicated",  "adam_betas": [0.9,  0.99]},
    "ve_gate_bank":   {"optim": "adam",    "comms": "replicated",  "adam_betas": [0.9,  0.99]},
    "x0_lambdas":     {"optim": "adam",    "comms": "replicated",  "adam_betas": [0.65, 0.95]},
    "lm_head":        {"optim": "adam",    "comms": "sharded",     "adam_betas": [0.5,  0.95]}, # Before embed for grad aggregation
    "embed":          {"optim": "adam",    "comms": "sharded",     "adam_betas": [0.5,  0.95]}, # Last (skipped when tied)
}

表中每个条目还隐含可选的 lr_mul、wd_mul 等乘子（记录版本未显式写出，可在后继代码中看到默认值 1.0）。最后一组"理论上仍可分组"的是 value embedding 表——作者选择把它们分别命名为 ve0、ve1、ve2，当作三个独立参数处理。

这个"按参数而非按优化器配置"的设计被完整继承到当前主线：track_1_short/training.py 中的 param_table 扩展为 qk_bank/vo_bank/mlp_bank（ANVIL，即 Muon 系）、scalars/smear_gate/ve_gate_bank 等（Adam，replicated），以及 lm_head/embed/value_embeds（Adam，sharded），并给 scalars 等标量设置了 lr_mul: 5.0、wd_mul: 0.0，给 lm_head/embed 设置 wd_mul: 150.。对应优化器 track_1_short/optim/anvil.py 中的 ParamConfig 数据结构（label/optim/comms/adam_betas/lr_mul/wd_mul/lr/...）也完全遵循该记录确立的形态。

1.3 用显式排序取代梯度累加 Hook

此前注册梯度累加 hook 的最大收益是让通信与耗时的 lm_head 梯度累加核函数重叠。由于 LM Head 已重写（见第 2 节），重叠小核函数的收益变小，作者干脆移除了 hook 特性，以换取更简洁的实现和对通信/计算顺序的灵活控制。

param_table 中的参数顺序决定了散射 / all_reduce 的发起顺序；另有一张独立列表规定处理顺序：

# - Process smaller/faster params first while large reduces complete
# - lm_head must complete before embed sync (when tied)
self.work_order = [
    "scalars", "smear_gate", "skip_gate", "attn_gate_bank", "ve_gate_bank", "x0_lambdas",  # Small, fast
    "ve0", "ve1", "ve2",  # Medium
    "lm_head", "embed",   # lm_head must complete before embed sync (when tied)
    "attn", "mlp",        # Large, polar express - process last to maximize overlap
]

这段顺序由 Claude Opus 给出，注释即其理由。作者认为该策略优于之前的 hook 方案，且很可能正是验证损失下降、从而允许删掉 5 步的原因。但作者也诚实指出一个"反直觉"的观察：work_order 与通信顺序（scatter_order）大致相反。

当前主线的 AnvilAndAdam.step 以三阶段实现同样的思想：Phase 0 对 replicated 参数做一次扁平 all_reduce（track_1_short/perf/replicated_adam.py）；Phase 1 按 scatter_order 发起 reduce/reduce_scatter；Phase 2 按 work_order 等待 reduce、计算更新、发起 all_gather；Phase 3 先等 lm_head gather（以便同步 embed），再等其余 gather。其流程注释与记录文档完全同源。

1.4 调试通信：profiler 的困境与 hook 触发顺序

记录作者坦言近期难以拿到好的 trace 数据——profiler 造成的减速比往常严重得多，通信数据也不对劲，导致难以制定并验证一个深思熟虑的通信策略。他甚至认为存在比当前显式排序更优的 hook 方案，"这个 PR 很大程度上是'无论如何，它就是比我们之前的方案好'"。

作为未来调试参考，作者打印了 hook 触发顺序，它揭示了梯度累加的实际次序：

1. scalars
2. embed
3. smear_gate
4. lm_head
5. skip_gate
6. mlp
7. attn
8. ve_gate_bank
9. ve2
10. attn_gate_bank
11. ve1
12. x0_lambdas
13. ve0

"并不直观！"——作者原话。为了把 NorMuon 参数也纳入观察，他给它们也加了 hook，并手动把梯度累加步数设为 2 以触发。

2. 转置 LM Head

另一个长期存在的通信问题：LM Head 的梯度累加核函数很慢。在 Adam+Muon 步与 Muon-only 步中的表现如下图所示：

解决办法是把 LM Head 的内存布局转置为 (768, 50304)，并同步更新 FP8 实现以支持该布局。从测试运行和单 GPU trace 看，其前向/反向性能与旧实现几乎一致。

2.1 内存布局不匹配的根因

旧实现中，反向传播为 LM Head 产生 (768, 50304) 的梯度，而权重按 (50304, 768) 存储。作者的经验是：矩阵乘法中 CUDA 往往能隐藏布局不匹配的代价（值可被读取一次、多次复用），但逐元素操作无法隐藏——逐元素操作每个值只被读取并精确使用一次，布局不匹配时不得不在某个矩阵上跨步遍历，导致核函数显著变慢。这正是当时梯度累加核函数比同尺寸同形状的其他 embedding 表慢约 4 倍的原因。

转置 head 使权重与梯度对齐，对 NorMuon-only 步收益最大——因为此前在这些步里，我们从未能把任何通信与这个慢核函数重叠。

2.2 权重绑定（Tied Embeddings）下的"解绑"

重新绑定 embedding 之前，这个改动收益更明显。绑定带来一个矛盾：输入 embedding 仍必须保持 (50304, 768) 的当前形状，这样 embedding 才能在内存中连续排布、被高效选取。

解决方式是"名义上解绑"：始终维护两个独立矩阵，但通过手动管理梯度与更新来保持等效绑定：

• 在散射之前，手动把 LM Head 与输入 embedding 的梯度相加；
• gather 完成后，把收到的更新分别复制进 embed 与 lm_head 权重。

这实际上又回到了起点——对两个未对齐矩阵做昂贵的逐元素操作——但这次它发生在优化器 step 代码内部：那里有充足的算力，且可以与其他通信重叠。

该逻辑完整保留在 track_1_short/optim/anvil.py 的 step 流程中（lm_head 标签处 transpose_add(embed_param.grad, param.grad)，Phase 3 中 transpose_copy(lm_param.data, embed_param.data)），并有 split_embed 标志控制绑定/解绑两种模式；track_1_short/training.py 注释明确"embed 在 split step（扩展阶段）前与 lm_head 绑定，之后解绑"。记录版本中这两步还是朴素的 .add_/.copy_，当前主线已将其替换为分块 Triton 转置内核 track_1_short/perf/kernels/transpose.py：每个 program 搬运一个 2D tile，从 src 合并读取、经寄存器 tl.trans 后合并写入 dst，双端都发出完整的 128 字节段，接近内存带宽上限，替代了 PyTorch 约 75k 块的朴素逐元素转置核。

2.3 解绑后的优化器状态

解绑还带来一个状态同步难题：由于朝向不同，每块 GPU 持有 lm_head 的 (96, 50304) 分片和 embed 的 (6288, 768) 分片，无法简单地把 Adam 动量缓冲从一方拷贝到另一方。感谢 @ClassicLarry 的建议，方案是执行一次显式通信：先把 lm_head 动量 all-gather 成完整 (768, 50304)，转置后每个 rank 截取属于自己的 embed 分片，从而保证每块 GPU 都有正确的 embed 动量状态，行为与精度和既有实现完全一致。

这一方案在 track_1_short/optim/anvil.py 的 copy_lm_state_to_embed 中得以保留：对 exp_avg 与 exp_avg_sq 分别执行 all_gather_into_tensor 得到完整 lm_head 动量，full_lm.T[rank * embed_chunk_size:(rank + 1) * embed_chunk_size] 即本 rank 的 embed 分片，并保留 embed_state['step'] = lm_state['step'] 以维持偏差校正的一致性。

3. Ideas：释放出的算力还能做什么

memcpy 消除后，优化器窗口内腾出了大量算力。记录作者尝试过若干方向，未必都能用，但值得继续探索或启发他人：

3.1 插值式 Embedding 更新

既然现在手动合并输入 embedding 与 LM Head 的梯度更新来维持绑定，就有机会"以不同比例混合"：例如在合并前把输入 embedding 梯度相对 lm_head 放大或缩小。

3.2 混合 / 切换 Muon 与 Adam

Adam 收敛极快，作者认为完全可以在 MLP 与 attention 权重上同时运行 NorMuon 与 Adam，并尝试了训练全程混合两者更新的实验，但没有出现有前景的结果。更感兴趣的是在训练末期把投影矩阵从 Muon 切换到 Adam：相比 Muon 的正交化约束，Adam 更新的逐参数自由度或许能让模型在结尾收敛到更好的局部最优（或更快找到）。(@Larry 在 PR 评论里有相关笔记。)

3.3 在 step 内提前开始 forward

优化器窗口内等待 attention/MLP gather 完成时，可以提前开始下一步的输入 embedding 选取、归一化、应用 smear gate，甚至跑第一个 attention 层（其权重比 MLP 先到）。作者已经开始实现，但在意识到梯度累加步下这几乎不可能实现后搁置；他设想应定义一个仅适用于 8×GPU（无梯度累加支持）的独立训练循环来落地该想法。

3.4 利用多余的 MLP 权重

现在要为额外一个 MLP 的权重支付通信与优化器开销。作者尝试用其中一块空闲矩阵充当 QKVO 权重、恢复第 6 个 attention 层——代价是这块矩阵随 MLP 权重一起优化，Muon 会把它当作一个大矩阵而非四个独立矩阵。结果 loss 很低（[3.2740, 3.2722, 3.2733, 3.2758]），但需要约 50 步才能回本，且 loss 太高（下表第 4 行）：

                            Runs   Time μ  Time σ  Time +/-  Loss μ  Loss σ  Loss +/-      p
This Record                   10 106.4940  0.0245    0.0000  3.2790  0.0006    0.0000 0.0002
Restore Attn 6, remove skip    4 109.3678  0.0814    2.8738  3.2754  0.0011   -0.0036 0.0017
Restore Attn 6, keep skip      4 109.5280  0.0281    3.0340  3.2738  0.0015   -0.0052 0.0019
Above, minus 50 steps          8 106.4315  0.0485   -0.0625  3.2816  0.0012    0.0026 0.9964

作者还分别尝试了移除 / 保留 skip connection（skip 似乎承担了部分缺失 attention 层的角色），并猜想若对 skip connection 或其他 attention 属性做恰当调整，这个方向或许可行。有趣的是，这一"用空闲 slot"的思路在后继架构中确实演化为"slot 11 并行 MLP"的实际用法（见 track_1_short/model/gpt.py）。

4. 记录成果在后继源码中的延续

从记录到当前主线，可以清晰看到三项核心贡献的延续与命名演进：

记录（2026-01-18）	当前主线（track_1_short）	证据位置
NorMuonAndAdam（Polar Express 正交化）	AnvilAndAdam（ANVIL 双轨动量 + 白化级联）	track_1_short/optim/anvil.py
attn / mlp 参数银行	qk_bank / vo_bank / mlp_bank	track_1_short/model/gpt.py
ve0/ve1/ve2 独立参数	单一 value_embeds（行稀疏 Adam）	track_1_short/training.py
显式 scatter_order / work_order	三阶段 step()，另加 deferred_labels 延迟 gather	track_1_short/optim/anvil.py
绑定态 embed.copy_(lm_head.T)	分块 Triton transpose_copy / transpose_add	track_1_short/perf/kernels/transpose.py
单次 H2D 批量上传标量（记录中未展开）	BankScalarStaging：一次 pinned H2D 上传全部银行标量	track_1_short/perf/bank_scalars.py

其中 BankScalarStaging 进一步解决了"每个标量一次常量重编译"的问题：所有银行每步读取的动量、衰减、双轨 beta/权重与逐矩阵学习率集中在一个 fp32 设备缓冲里，优化器在 step 顶部用一次非阻塞 pinned 拷贝完成上传，CUDA 图重放直接从内存读取当前值，避免为每步新常量触发重新编译。

5. 复现与验证

• 实验日志：本次记录的 baseline 与 unified-optimizer 两组合计 16 份运行日志分别存放在 baseline/ 与 unified-optimizer/ 目录，可对照文中的统计表核验 Time μ、Loss μ 与 p 值。
• 实现快照：记录版本的训练脚本即 unified-optimizer/2fc79469-a527-4bde-8540-8426ed3352d1.txt，其中可看到 NorMuonAndAdam 类、CastedLinearT（转置权重 (in_features, out_features) 存储，FP8 分支走 nanogpt::mm_t 自定义算子）、mm_t_backward（grad_w 天然与权重存储对齐，无需最终 .T）以及 copy_lm_state_to_embed 的完整实现。
• 复现环境：项目根目录提供 requirements.txt、run.sh 与 train_gpt.py，配合 Dockerfile 即可构建 8×GPU 训练环境；若想对照当前主线的演进实现，直接阅读 track_1_short/optim/anvil.py 与 track_1_short/training.py 即可。

小结

这条记录展示了大规模分布式训练优化中一条可复制的路径：先找到"慢在何处"（memcpy、布局不匹配导致的慢速逐元素核），再用不改变数学的方式消除瓶颈（参数银行、显式通信排序、转置布局），最后用基准测试严谨量化收益。其核心结论——"优化器窗口内的一次显式排序胜过 hook 机制"、"逐元素核无法隐藏布局错配"——对任何在多 GPU 上压榨训练吞吐的工程实践都具有直接的参考价值；而参数银行 + 逐参数 param_table 的设计，也在后继的 ANVIL 优化器中成为默认架构。