modded-nanogpt 短赛道 MUDD 门控 + 轻量 DC 校正：把 124M 训练时间压到 75.5 秒

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

本篇技术指南完整梳理 modded-nanogpt 短赛道（track_1_short）中2026-05-27_MuddGatedAndDC这一条实验记录的三个迭代（v1/v2/v3）：在 XSA 门控基线之上引入MUDD 门控（Multiway Dynamic Dense Connections gate）统一生成各注入路径的逐 token 系数，并用轻量 DC（Differential Correction）对最后一个非配对注意力层做局部窗口 QK softmax 修正。文章结合仓库源码与记录日志，详细解释门控布局、DC 张量的生成方式、bigram 兼容性调整与可复现的统计口径；读完你可以掌握「如何用一个共享 MLP 门控取代多条静态门控路径、同时砍掉 95 步训练调度并提升验证损失」的完整实验方法。

1. 记录概览：同一想法上的三次迭代

该记录位于 records/track_1_short/2026-05-27-MuddGatedAndDC/README.md，跟踪了同一个技术思路的三轮实验：

最终提交版本为this_pr_v3。实验环境为 8× NVIDIA H100 80GB HBM3，Python 3.12.3，PyTorch2.10.0+cu128，Triton3.6.0，驱动580.159.03。

2. 结果数据与统计口径

2.1 v1/v2 与 XSA 基线对比

v1/v2 均基于2026-05-07_XSAGatedLayers短赛道基线，结果如下（Time delta相对xsa-baseline-s1385，单位为秒）：

Runs Steps Time mean Time sd Time delta Loss mean Loss sd Loss p xsa-baseline-s1385 10 1385 81.2187 0.0681 0.0000 3.27841 0.00133 2.17e-03 xsa-baseline-s1335 9 1335 78.5576 0.0401 -2.6611 3.28477 0.00168 1.00e+00 this_pr_v1 11 1335 78.9669 0.0425 -2.2518 3.27714 0.00123 8.15e-06 this_pr_v2 10 1325 77.7050 0.0297 -3.5137 3.27808 0.00098 7.76e-05

关键解读：

• 把调度从 1385 步砍到 1335 步后，纯 XSA 基线（xsa-baseline-s1335）的损失明显掉到目标线之上（Loss mean 3.28477，p = 1.00e+00），证明步骤削减本身是有代价的；
• v1 在同样的 1335 步下把损失拉回3.27714，不仅优于匹配步数的基线，还略微优于 1385 步的原始基线；
• v2 再砍 10 步（1325 步），时间进一步降到77.7050s，损失 3.27808 仍远低于 3.28 目标。

2.2 v3 与上游新基线的对比

v3 基于带有bigram 符号技巧（bigram sign trick）+ FP8 MLP的更新上游基线，Time delta在该表中以正数表示加速：

Runs Steps Time mean Time sd Time delta Loss mean Loss sd Loss p upstream-sign-fp8 - 1390 79.2000 - 0 - - - this_pr_v3 10 1290 75.5199 0.0487 3.6801 3.27851 0.00097 4.53e-04

• upstream-sign-fp8行是复制自上游的最新基线结果，未在本文件夹内重跑，仅作为时间参考（因此没有重复损失日志）；
• this_pr_v3在1290 步（比基线的 1390 步少 100 步）达到75.5199s，比基线快3.68s，平均损失 3.27851，通过了p < 0.01的目标规则（p = 4.53e-04）。

2.3 统计口径

README 中定义的Loss p是相对 3.28 目标线的一样本 t 检验（越小越代表「有把握低于 3.28」）：

scipy.stats.ttest_1samp(losses, 3.28, alternative="less").pvalue

v2 相对 v1 的损失差异（+0.00094）与 v3 相对 v2 的损失差异（+0.00043）都使用双侧 Welch t 检验评估，分别得到p = 0.066与p = 0.336，在 5% 显著性水平下均不显著。也就是说：v2/v3 的时间收益并不以统计显著的损失退化为代价。

2.4 v2 → v3 的加速构成

v2 与 v3 并非严格同基线的比较——v3 先于上游 bigram-sign / FP8-MLP 基线更新。加速量大致与重基线化收益吻合：FP8 MLP 约贡献 0.5s，bigram 符号技巧约贡献 1.5s，因此 v2 → v3 的 2.19s 时间收益应主要归因于上游基线自身的提速，而非 MUDD/DC 的新增收益。

3. v1：MUDD 门控 + 轻量 DC

3.1 设计动机

v1 的两项模型改动与路径清理基于同一观察：DC、XSA 与attn_gate_w都在调制注意力流。一旦第 10 层获得 DC 校正，旧的 XSA / 注意力门控调制就显得冗余，因此 v1 只保留一个更小的门控子集，并让 MUDD 门控来生成所有激活系数。

3.2 MUDD 门控布局（v1）

门控采用「pre / post」两段式结构，x0指第 0 层之前的归一化隐藏状态：

• pre gate（由x0计算）：XSA 作用于层1, 3, 4，注意力门控作用于层3，x0/bigram 注入门控作用于层0..5；
• post gate（在层6生成）：XSA 作用于层7，注意力门控作用于层10，x0/bigram 注入门控作用于层7..10，以及层3→ 层6的跳跃连接系数。

清理之后：

• 激活的 XSA 层为1, 3, 4, 7，层8与10上旧的 XSA 使用被移出激活路径；
• 注意力门控仅来自 MUDD 门控的层3与10；
• 固定的x0_lambdas/bigram_lambdas、静态xsa_alphas与独立的skip_gate都不再是激活的控制路径。

3.3 轻量 DC 路径

DC 路径被刻意设计得极窄：

• 仅 post-only：无 pre 组合、无 dense-dense 项；
• 只作用于最后一个非配对注意力层（layer 10）；
• 基础注意力仍由FA3提供，自定义 Triton 内核只负责重算局部窗口 QK softmax 并叠加 DC 校正。

v1 把调度从1375 + 10 = 1385砍到1315 + 20 = 1335，同时提升平均验证损失。相对匹配步数的xsa-baseline-s1335，v1 的平均损失改善约0.00763（双侧 Welchp = 1.44e-08）。已知的孤立实验表明：DC 单独大约贡献0.0046的改善，其余来自 MUDD 门控与路径清理。

4. v2：由门控生成 DC（Gate-Generated DC）

v2 保留同一基本思路，但移除了独立的 DC 参数路径——layer 10 的 DC 张量现在直接取自 MUDD 门控的 post gate：

dc_weights[10] = (post_gate[..., 29:35], post_gate[..., 35:41])

即 post gate 的第 29..34 维作为w1、第 35..40 维作为w2（各 6 维，对应 6 个注意力头）。dc_gate函数此时只做三件事：校验形状、对w1按头做 RMS 归一化、返回连续的(post_w1, post_w2)——如 this_pr_v2 的提交代码 中dc_gate的实现所示，post_w1经F.rms_norm(..., (num_heads,), eps=1.0e-6)后与原张量同 dtype 返回。

v2 的 MUDD 门控布局被收窄并显式化：

• pre gate（由x0计算）：XSA 作用于层1, 3，注意力门控作用于层3，x0/bigram 注入门控作用于层0..3；
• post gate（在层4起始处生成）：XSA 作用于层4, 7，注意力门控作用于层10，x0/bigram 注入门控作用于层4, 5, 7, 8, 9，外加层3的跳跃系数与 layer-10 的 DC 张量；
• 最后一层的 x0/bigram 注入保留在既有 layer-10 MUDD 块中，通过mu[10]与mu[11]实现。

对应提交代码中的门控通道排布为：pre 共 26 维（xsa[1,3] 12 + attn[3] 6 + inject[0..3] 8），post 共 41 维（xsa[4,7] 12 + attn[10] 6 + inject[4,5,7,8,9] 10 + skip 1 + dc[10] 12）。初始化时 DC 的w1通道 bias 设为1.0 / _mudd_gate_scale，即初始输出从1.0起步，保证模型加载后 DC 以接近单位增益、可被学习的方式介入。

v2 使用1315 + 10 = 1325步并以p = 7.76e-05通过目标。相对 v1，它以+0.00094平均损失换来1.26s更少的墙钟时间；该损失差距在本样本中不显著（双侧 Welchp = 0.066）。

5. v3：bigram 兼容重基线化（提交版本）

v3 把「MUDD 门控 + 门控生成 DC」重基线化到带bigram 符号技巧与FP8 MLP的最新上游基线。重基线后遇到的主要问题是 bigram 路径：上游默认bigram_dim = 192、bigram_vocab_size = 50304 * 15，在该设置下 MUDD 门控无法恢复 v2 观察到的收益——推测是 bigram 维度 / 词表大小的权衡没有让 MUDD 生成的 bigram 门控充分利用 bigram 嵌入路径。

v3 改为：

bigram_dim = 768 bigram_vocab_size = 50304 * 15 // 2

这是一个面向门控 bigram 路径的损失/速度权衡：维度足够宽，MUDD 门控能发挥作用；词表相对原始768维 bigram 嵌入会使用的规模减半。v3 提交代码中对应实现可见于 this_pr_v3 的提交代码：self.bigram_embed = nn.Embedding(args.bigram_vocab_size, args.bigram_dim)并零初始化，配合符号表bigram_sign_table实现 PR299 的压缩技巧。

v3 的另一项改动是让_mudd_gate_scale成为可学习参数，初始化为0.1：

self._mudd_gate_scale = nn.Parameter(torch.tensor(0.1))

原因是训练早期 MUDD 门控表现出一些不稳定性；把全局门控缩放交给模型自己调节后，门控路径更容易稳定下来。提交的 v3 调度为1275 + 15 = 1290步。

6. 源码级实现原理

6.1 MUDD 门控的底层实现

短赛道当前源码 track_1_short/model/gpt.py 中保留并持续演进了这套 MUDD 门控架构：

MUDD_GATE_HEAD_LANES = 6 # 每个注意力头的门控通道 MUDD_GATE_SCALE = 0.1 # 门控输出初始缩放；bias 按它预除以存储 POST_GATE_LAYER = 4 # post gate 在层 4 起始处生成 XSA_LAYERS = (1, 3) # XSA 强度来自 pre gate ATTN_GATE_LAYERS = (3, 10) # 注意力门控层

init_mudd_gate创建两组参数：mudd_gate_w1（形状(2, 64, model_dim)，kaiming 初始化）与mudd_gate_w2（形状(2, max_num_coef, 64)，零初始化），以及按层位排布 bias 的mudd_gate_b2。核心前向forward_mudd_gate是一个共享 MLP：

x = F.gelu(F.linear(x, self.mudd_gate_w1[id])) return (F.linear(x, self.mudd_gate_w2[id, :num_coef]) + self.mudd_gate_b2[id, :num_coef]) * self._mudd_gate_scale.type_as(x)

随后unpack_pre_mudd_gate/unpack_post_mudd_gate按固定通道布局把门控向量切分给xsa_alphas、attn_gates、x0_gates、bigram_gates，post gate 还会额外返回 layer-6 的跳跃门控。这与 README 中 v2/v3 描述的布局一一对应：pre 门控服务POST_GATE_LAYER之前的层，post 门控服务它自身及之后的层。

6.2 训练超参

track_1_short/training.py 中为 MUDD 门控参数配置了独立的 Adam 优化段：

"mudd_gate_w1": {"optim": "adam", "comms": "replicated", "adam_betas": [0.9, 0.99], "lr_mul": 0.1}, "mudd_gate_w2": {"optim": "adam", "comms": "replicated", "adam_betas": [0.9, 0.99], "lr_mul": 0.1}, "mudd_gate_b2": {"optim": "adam", "comms": "replicated", "adam_betas": [0.9, 0.99], "lr_mul": 0.1, "wd_mul": 0.0}, "_mudd_gate_scale": {"optim": "adam", "comms": "replicated", "adam_betas": [0.9, 0.99], "lr_mul": 0.1, "wd_mul": 0.0},

要点：门控参数走 Adam（而非投影矩阵的 ANVIL/Muon 路径）、全副本通信、lr_mul = 0.1的低学习率缩放；b2与_mudd_gate_scale关闭权重衰减。这与 README 记录的「MUDD 门控需低学习率稳定训练」一致，也与前序 MuddFormer 记录 中 dense 参数lr_mul = 0.25的低 LR 传统一脉相承。

6.3 DC 校正内核

在 v1/v2/v3 的提交日志（train_gpt.py+dc_triton_kernels.py）中可以看到完整链路：

y = flash_attn_interface.flash_attn_varlen_func(...) # 基础 FA3 注意力 if dc_w is not None: dc_weights = dc_gate(x, dc_w, self.num_heads) y = dc_attention_postonly_nodd_correction_add_base_triton( y, q, k, v, dc_weights, None, scaling=yarn.attn_scale, window=112, seq_lens=seqlens, )

即先跑 FA3 得到基础输出y，再由自定义 Triton 内核在112 的局部窗口内重算 QK softmax（使用与 FA3 相同的yarn.attn_scale缩放），叠加 DC 校正后写回。注意在 v1 中 DC 权重来自独立参数（dc_w1_bank/dc_w2_bank，且仅在dc_layers = [10]且非配对层上生效），到 v2 起 DC 权重完全改由 MUDD 门控 post gate 切片生成，dc_gate退化为纯形状校验 + 按头 RMS 归一化 + 连续性保证。

7. 实验文件与复现说明

记录目录 records/track_1_short/2026-05-27-MuddGatedAndDC 下的组织方式：

每个.txt日志本身就是该次运行的完整提交快照（包含train_gpt.py、triton_kernels.py与dc_triton_kernels.py的全文），因此任意一条日志都可以独立回溯当时的代码状态。统计结果可直接用 README 给出的ttest_1samp/ Welch t 检验口径复算。

8. 小结：该记录沉淀的方法论

• 用「门控统一」替代「多条静态路径」：DC、XSA、注意力门控、x0/bigram 注入、层间跳跃全部收敛到同一个 64 维 MLP 门控输出的通道切片上，单步开销几乎为零，却提供了足够的「每步损失容量」去削减训练调度；
• DC 要「窄」才有性价比：post-only、无 pre 组合、无 dense-dense 项、只作用于最后一个非配对注意力层，配合 FA3 + 112 窗口 Triton 修正，是「低成本高质量」的关键取舍；
• bigram 配置与门控协同：v3 的768维 bigram + 减半词表是让 MUDD 门控在新基线上恢复收益的必要条件，说明门控类设计对底层嵌入路径的形状高度敏感；
• 可学习全局缩放是稳定器：把_mudd_gate_scale从常量改为初始值 0.1 的可学习参数，缓解了门控路径早期的训练不稳定。

最终，this_pr_v3以1290 步 / 75.5199s / 平均损失 3.27851（p = 4.53e-04）提交，比上游 bigram-sign + FP8-MLP 基线的 1390 步参考时间快 3.68 秒，同时严格通过 3.28 损失目标线——这是短赛道「每步更高质量 + 更少步骤」路线又一次可复现的实证。

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