modded-nanogpt Recursive From-Best 实践：FP8 注意力投影、NorMuon 退火探索噪声与精简 ReLU² MLP 内核的 124M 训练优化

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

本篇文章以仓库records/track_1_short/2026-06-11_RecursiveFromBest/README.md为核心，完整复现并剖析本次 Track 1 提交（以下简称this_pr）相对世界纪录 #83 在相同 8×H100 硬件上的对照实验结果，以及驱动提速的五项关键技术改动：FP8 注意力 QKV/O 投影、NorMuon 中退火的行 RMS 缩放 Langevin/SGLD 探索噪声、bigram 与 value-embedding 银行上的谨慎（cautious）Adam、更精简的融合 ReLU² MLP Triton 内核，以及调度/架构重调。读完本文，你将理解这类"从最优基线递归改进"的对照实验方法论，掌握各项改动的源码级依据，以及如何在自己的 124M NanoGPT 训练任务中复现同样的实验设置与统计判读。

一、实验背景：为什么需要"Recursive From-Best"

modded-nanogpt 的 Track 1 以固定步数内把 124M 参数量级模型的验证损失压到最低为目标。当优化已经推进到世界纪录级别时，单次改进往往只有千分之一的验证损失波动，因此判断一次提交是否真正优于上一代纪录，必须把新旧两个版本放在同一套硬件、同一份数据流上重跑多轮，再借助统计检验来判断均值差异是否显著。

本次提交的做法正是"递归从最优基线出发"（Recursive From-Best）：以世界纪录 #83（即记录目录中的baseline_pr299，README 描述其为 "Sign Trick on Bigram Embed" 的官方提交）为起点，将其在同一个 8×H100 Modal 环境上重跑 10 轮作为对照组，再在完全相同环境下运行新提交 13 轮作为实验组，最后用单侧单样本 t 检验比较两者相对 3.28 阈值损失的显著性。README 明确说明：PrimeIntellect 官方验证仍在等待中，表中的数字是同一硬件下 Modal 上的实测结果。

二、对照实验结果与统计判读

2.1 核心对照表

README 中的核心对照数据如下（n为有效 run 数，steps为训练步数，train_time为单次完整训练耗时）：

两个关键事实：

1. 速度显著提升：同一套 8×H100 硬件上，this_pr平均训练耗时 77.34s，比 #83 重跑基线的 80.61s 快约 4%，且耗时波动更小（std 0.20s vs 0.22s），甚至低于官方提交记录的 79.7s。
2. 验证损失基本持平：this_pr的平均验证损失 3.27893 比基线 3.27850 略高约 4e-4（约 0.3 个标准差），但其标准差更小（0.00137 vs 0.00144）。两组相对 3.28 阈值的单侧 p 值分别为 0.0078 与 0.0047，均显著小于 3.28——也就是说两个版本都稳定地跑进 3.28 以内，差异主要落在同一起跑线上，而本次提交的核心增量是同等质量下把耗时压低约 3.3 秒。

2.2 逐 run 明细

this_pr 的统计文件列出了 13 轮各自的 val_loss 与训练耗时：val_loss 分布在 3.2766～3.2819 之间，训练耗时全部落在 77.1～77.9 秒区间；基线统计文件则记录了 10 轮 val_loss 3.2762～3.2805、耗时 80.3～81.1 秒的分布。两份文件均给出了验证器输出与检验公式：scipy.stats.ttest_1samp(accs, 3.28, alternative="less").pvalue。

2.3 统计方法解读

两份统计文件明确标注了判读标准：对每轮得到的 val_loss 序列做单侧单样本 t 检验，原假设是均值不小于 3.28，备择假设是均值小于 3.28。由于 3.28 通常是该 Track 的目标线，p(mean<3.28)越小，说明该提交稳定突破目标线的置信度越高。这里两组 p 值都处于 0.005～0.008 量级，足以支撑"稳定跑进 3.28"的结论，但不足以支撑"新提交损失严格优于旧提交"——后者需要看置信区间重叠情况，而两者均值的差距远小于各自标准差，因此本文的合理解读是：新提交在保持同等验证损失水平的前提下实现了可重复的约 4% 提速。

三、五项主要改动逐一拆解

README 将本次提交的改动概括为五项，下面逐项结合源码展开。

3.1 FP8 注意力投影：QKV 与 O 的训练前向量化

Main changes 第 1 条：FP8 attention projections for QKV and O in the training forward pass.

这是本次提速的直接来源之一。注意力层的 QKV 投影与 O 投影在训练前向中使用 FP8（e4m3）量化，将原本占前向相当比例的矩阵乘法换成 FP8 张量核心路径。

源码侧的直接证据：

• model/gpt.py 中的quantize_attn_fp8()负责维护注意力权重与激活的 FP8 缓存与缩放因子；
• model/attention.py 的forward()接受qkv_fp8元组（训练专用，验证时传None），其中包含weight_f8, weight_f8_t, weight_scale, x_scale, grad_scale, x_f8, x_f8_t，即权重双布局 FP8 表示、激活 FP8 与对应的缩放因子；
• perf/kernels/fp8_attention_quant.py 提供quantize_dual_layout_packed_batched等 Triton 内核，为 QK、V 投影权重同时产出行/列两种布局的量化表示，供前向 GEMM 与反向 GEMM 复用；
• perf/cuda_graphs/fp8_refresh_graphs.py 将quantize_attn_fp8()与 MLP 量化封装进 CUDA Graph，避免每步在 eager 模式下反复启动量化内核。

值得注意的限制条件：FP8 量化只应用于训练前向的注意力投影；验证阶段仍走完整 bf16 路径（qkv_fp8=None），保证评估结果的精度不受影响。README 也特别注明"没有额外的torch._inductor.config或 compile 标志"——即本次提速不依赖额外编译优化开关。

3.2 NorMuon 中的退火行 RMS 缩放 Langevin/SGLD 探索噪声

Main changes 第 2 条：Annealed row-RMS-scaled Langevin/SGLD exploration noise in NorMuon, zero by cooldown.

NorMuon 是仓库优化器的核心家族之一。在 optim/anvil.py 的类文档中可以看到 ANVIL（纪录 #360 的后续）与标准 Muon 的区别：双轨道（fast/slow）动量 EMA、以 ANVIL cascade（六次重推导的五次映射 + Frobenius 归一化）取代 Newton–Schulz、以及来自 NorMuon 的逐 lane 能量均衡器（低秩、Adafactor 风格的方差估计，见_rail_equalizer，optim/anvil.py）。

本次提交在此基础上引入了探索噪声：按行 RMS 缩放、形态类似 Langevin 动力学 / SGLD（随机梯度 Langevin 动力学）的高斯扰动，叠加到 NorMuon 的更新量上。其核心设计是"退火"——噪声幅度随训练进度衰减，到 cooldown 阶段（学习率冷却期）必须归零，从而避免噪声干扰最后阶段的收敛。这本质上是训练早期用噪声换取对损失曲面更宽的探索、后期确定性收敛的标准思路在正交化优化器（Muon 系）上的应用。

需要说明：由于本提交的记录目录只归档了日志与统计文件（this_pr/*.txt），具体噪声幅度调度、行 RMS 计算与退火曲线系数以 README 描述为准；可以推断该噪声作用于 NorMuon 更新路径（对应 optim/anvil.py 的 ANVIL bank 更新体），且其"零到 cooldown"属性与仓库调度器中的cooldown_frac机制（schedule.py）天然衔接。

3.3 bigram 与 value-embedding 银行上的 C-Optim / 谨慎 Adam

Main changes 第 3 条：C-Optim / cautious Adam on the bigram and value-embedding banks.

"谨慎（cautious）"在这里指符号对齐的权重衰减：只有当"更新方向"与"参数当前值"符号一致时，才对参数施加权重衰减项；符号相反时跳过衰减。这样既保留权重衰减的收缩作用，又避免它在梯度方向与参数符号相反时产生反方向拉扯。

源码证据非常清晰：

• optim/anvil.py 的_adam_update_step：mask = (update * p_slice) > 0，随后update.addcmul_(p_slice, mask, value=eff_wd_t)——只有乘积为正的位置才被计入衰减；
• perf/kernels/ngram_adam.py 的行稀疏 Adam 内核同样实现update += where(update * p > 0, p * decay, 0)，注释标注其源自纪录 #360（ANVIL2）的 bigram 内核；
• ngram_table.py 说明 bigram 表并非模型参数，而是拥有独立行稀疏 Adam（自己的更新节拍is_update_step），共享基础 lr / weight decay / eps，在优化器通信调度内运行；
• value-embedding 银行同样是 Adam 参数，但跟随 n-gram 表的节拍更新，采用行稀疏梯度交换 + 副本拉取（perf/value_embed_pull.py），其 betas 与 wd_mul 由value_embed_betas_and_wd_mul随节拍调整（见 training.py）。

配置层面，training.py 的param_table展示了各类参数的优化器归属：qk_bank/vo_bank/mlp_bank走 ANVIL（sharded 通信），value_embeds走 Adam（sharded），embed与lm_head走 Adam 且wd_mul: 150.，而scalars的wd_mul: 0.0。这意味着"谨慎 Adam"精确作用于 bigram 表与 value-embedding 这些稀疏、高频更新的嵌入型参数上，与它们的稀疏梯度交换机制（行稀疏更新）相配合。

3.4 精简的融合 ReLU² MLP Triton 内核

Main changes 第 4 条：Leaner fused ReLU^2 MLP Triton kernel, storingpost = relu(pre)^2and reconstructingsqrt(post)in backward.

MLP 前向在 Triton 内核内部融合了线性层与 ReLU 平方激活，并且只存储激活后值post = relu(pre)^2（FP8），而不是同时保存 pre 与 post；反向时通过sqrt(post)重建 pre 的符号/数值信息，从而省下一份中间张量的写入与显存占用。

证据位于 perf/kernels/mlp.py：

• linear_relu_square_kernel的STORE_PRE、STORE_POST_BF、EMIT_F8、EMIT_T等编译期开关精确控制"是否写 pre / 是否写 post"（第 37～53 行）；
• 前向路径把post量化进 FP8（inv_post_scale、post_scale_ptr、post_amax_ptr，第 68～119 行）；
• 反向路径正是 README 所说的"重建"：c0 = 2.0 * acc0 * sqrt(a0_raw * post_scale)（第 132～133 行），即从存储的 post 反解出 pre，乘回梯度得到输入梯度。

结合上一节的 FP8 量化，可以看到本次提交的总体思路是让前向/反向尽可能只产生 FP8 中间结果、并压缩中间张量，从而在不牺牲验证损失的前提下缩短单步时间。

3.5 调度与架构重调

Main changes 第 5 条：Schedule/architecture retunes: fewer paired-head layers, tied embed throughout, and step-count/window changes.

• 更少的 paired-head 层：paired-head 是注意力头配对设计的架构选项（model/attention.py 的paired参数），本次减少了使用该结构的层数；
• 全程 tied embed：embedding 与 lm_head 权重在整个训练过程中保持共享（tied），不再在扩展阶段解绑——而仓库默认路径里存在"split 步骤后解绑"的设计（training.py 注释：embed 与 lm_head 在 split step 前 tied，之后 untied）。本次提交选择全程 tied，减少了一个参数银行及其更新开销；
• 步数与窗口调整：this_pr训练 1398 步（基线 1385 步），同时涉及窗口长度（context window / 滑动窗口）等调度参数的变化。

这些重调与 2.1 节的速度数字直接相关：步数略有增加但总耗时反而缩短，说明每步的延迟收益（FP8 注意力、精简 MLP 内核、更少 paired-head 层）大于步数增量。

四、实验设置与可复现性

README 的 Notes 部分给出了本次对照实验的完整环境前提，写文章时一并保留：

• 数据流：训练与验证 token 流保持 FineWeb 的 train / val 划分不变（数据加载实现在 data.py 中，训练入口见 training.py）；
• 编译配置：没有添加额外的torch._inductor.config或 compile 标志，所有收益来自模型/内核/调度层面的改动；
• 硬件与软件栈：日志采集自 Modal 平台、8× NVIDIA H100 80GB，PyTorch 2.10.0+cu128，Triton 3.6.0。

复现步骤建议：在具备 8×H100 的环境中，先按 baseline_pr299 的 10 轮设置重跑 #83，再运行 this_pr 的 13 轮设置，用仓库统计文件给出的scipy.stats.ttest_1samp(accs, 3.28, alternative="less")复算 p 值，对照 2.1 节表格即可验证本次结论。单轮耗时约 77～81 秒，整套对照实验在一台 8×H100 上可以在半小时内完成。

五、局限与判读边界

1. 官方验证未完成：README 明确 PrimeIntellect 官方验证仍在等待中，表格内均为同硬件 Modal 实测，不应视为最终官方成绩；
2. 损失水平持平而非提升：从数据看，本次提交的价值主要是同等 val_loss 下约 4% 的耗时缩短（77.34s vs 80.61s），val_loss 均值差（3.27893 vs 3.27850）远小于组内波动，不宜解读为损失改善；
3. 探索噪声的精确实现未归档：退火噪声的具体调度系数需要从该 PR 的运行配置获取，本仓库只记录到 README 描述层面，阅读时请以"存在性 + 设计意图"理解，而非精确超参。

总而言之，2026-06-11_RecursiveFromBest是 modded-nanogpt Track 1 优化链中一次典型的"从最优基线递归改进"：通过 FP8 注意力投影、NorMuon 退火探索噪声、谨慎 Adam 化稀疏嵌入银行、精简融合 MLP 内核与调度重调，在守住 3.28 验证损失线的同时把 124M 训练压进 77.3 秒均值，并以同硬件重跑 + 单侧 t 检验的严谨方式呈报结果——这套"同基线、多轮、统计检验"的对照方法论本身，值得所有追求极致训练效率的团队借鉴。

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