modded-nanogpt GPT-2 Medium 赛道 25.29 分钟纪录解析：ISRU 软上限、分片混合精度 Muon 与权重合并

本篇文章以 records/track_2_medium/2025-04-16_Record7/README.md 为骨架，深入拆解 modded-nanogpt「GPT-2 Medium 极速训练」（Track 2，目标验证损失 2.92）第 7 号纪录：在 8 块 NVIDIA H100 上把训练时间从 25.95 分钟压到 25.29 分钟的五项改动。文章逐条还原每项优化的动机、源码实现与实测收益（步数、单步耗时、验证损失），并对照前一纪录（2025-03-25_ArchOptTweaks）与完整运行日志（223_3310d0b1-b24d-48ee-899f-d5c2a254a195.txt），帮助读者掌握：何时该牺牲 FP8 换精度、logits 软上限的轻量化替代、如何用 uint16 尾数拆分实现混合精度 Muon、以及合并权重张量这一通用显存/时间优化手法。

一、赛道背景：从 3.28 到 2.92，从 GPT-2 Small 到 GPT-2 Medium

仓库 README.md 说明，modded-nanogpt 的核心玩法是「NanoGPT 极速跑」（NanoGPT speedrun）：用 8 块 NVIDIA H100，以尽可能少的墙钟时间把语言模型在 FineWeb 验证集上的交叉熵损失训练到目标值。主赛道（Track 1）的目标是 3.28，对标 Karpathy 的 llm.c GPT-2（small）复现；Track 2（GPT-2 Medium）则把目标降到 2.92，对应 350M 参数量级的 llm.c 基线，性能大致与原始 GPT-2 Medium 相当。除目标损失外，其余规则（数据管线不可改动、只能使用固定 token 流、禁止额外 torch._inductor.config 标志等）与主赛道一致；Track 2 在 Record 6 之后被允许开启 coordinate_descent_tuning（编译会变慢，但可换来运行时收益，本纪录的代码中确实保留了该标志）。

Track 2 的纪录演进非常密集：llm.c 基线 5.8 小时 → Record 2 的 29.3 分钟 → Record 6 的 25.95 分钟（2x MLP wd、qkv norm、all_reduce/opt.step 重叠、优化 skip 模式）→ 本文的主角 Record 7：25.29 分钟（2025-04-16，贡献者 @YouJiacheng），随后被 Record 8（24.50 分钟，三次方滑动窗口调度）刷新。

Record 7 的 README 只有五条清单，但每一条都对应日志文件中可验证的源码改动与实测数据：

1. 移除 FP8 LM head：训练步数 6710 → 6450，墙钟时间几乎不变；
2. logits softcap 由 tanh 换成 ISRU：每步快约 1.5ms，验证损失略降（2.91970 → 2.91944）；
3. 新的分片混合精度 Muon，移除 CastedLinear：每步快约 3.6ms；
4. 合并 QKV 与 O 权重：每步快约 0.8ms；
5. 合并标量权重：每步快约 0.6ms。

二、优化一：移除 FP8 LM head，用步数换精度

前一纪录（Record 6）在 lm_head 上启用了 FP8 量化线性层，源码见 train_gpt-20250329.txt：

self.lm_head = CastedLinear(model_dim, next_multiple_of_n(vocab_size, n=128), use_fp8=True, x_s=0.5, w_s=2**-9, grad_s=2**-19)

其底层是 torch.library.custom_op 注册的 nanogpt::mm / nanogpt::mm_backward FP8 矩阵乘算子（torch._scaled_mm，前向 e4m3、反向 e5m2），把激活、权重与梯度分别除以缩放因子后量化为 float8 再计算。

Record 7 把这套机制整体移除：lm_head_w 重新变为一个普通 bf16 参数（223_3310d0b1-b24d-48ee-899f-d5c2a254a195.txt）：

self.lm_head_w = nn.Parameter(torch.zeros(next_multiple_of_n(vocab_size, n=128), model_dim))

损失计算也回到直接 F.linear(x, self.lm_head_w.bfloat16())。动机在于：LM head 是训练后期决定最终验证损失的瓶颈，FP8 量化带来的舍入噪声会抬高损失下限，迫使训练多跑约 260 步（6710 → 6450）才能回到同一损失水平；而 head 的矩阵乘在整段训练中的计算占比很小，换回 bf16 后墙钟时间几乎无变化。这是一次典型的「用算力换精度、再折算成更少步数」的取舍，从日志末尾的最终结果（val_loss:2.919352，见第 7125 行）看，这个取舍是划算的。

三、优化二：logits softcap 从 tanh 换成 ISRU

Gemma 2 引入的 logits softcap 在本项目中经历了多次演进：最初是 tanh 软上限（Record 6 代码注释记录为「@KoszarskyB 把上限从 30 降到 15，@YouJiacheng 平移 +15，2*sigmoid(2*x)=tanh(x)+1」）。Record 7 的 forward 中换成了 ISRU（Inverse Square Root Unit） 形式（日志第 304 行）：

loss = F.cross_entropy(15 * logits * torch.rsqrt(logits.square() + 225), target_seq)

数学上 15 * x / sqrt(x² + 225) 等价于 15 * ISRU(x / 15)，即带缩放因子的软上限：x 趋近 ±∞ 时输出饱和在 ±15，与 tanh softcap 的渐近行为一致，但计算路径完全不同——tanh（或 2*sigmoid-1）需要求指数或三角函数，而 rsqrt 只需要一个快速倒数平方根指令，GPU 上开销显著更低，同时数值稳定性更好（无需担心指数上溢）。

README 记录的收益是双向的：单步快约 1.5ms，且验证损失从 2.91970 微降到 2.91944。对验证阶段同样生效（验证分支第 310 行使用相同的 ISRU 公式，分 4 块 chunk 计算以降低峰值显存）。这提示一个可复用的经验：在保证饱和点（此处为 15）不变的前提下，softcap 的实现函数可以替换为计算更廉价的代数等价形式，既省时间又不伤精度。

四、优化三：分片混合精度 Muon，移除 CastedLinear

这是五项改动中收益最大的一项（每步约 3.6ms），也是工程上最有意思的部分。它包含两个层面：

4.1 混合精度状态：bf16 权重 + uint16 尾数拆分

Muon 的更新核心在 update 函数。为了既享受 bf16 权重的低显存与快访存，又不丢失 fp32 量级的累积精度，代码把每个 bf16 参数张量视图复用为 uint16 的两个半区：高 16 位（bf16 的指数+高位尾数）与一个独立的 uint16 mantissa 张量（低位尾数），二者拼成一个 uint32，再按 fp32 解释：

acc_m_u32 = (acc_bf16_view_u16.to(torch.uint32) << 16) | mantissa.to(torch.uint32)
acc_m_u32.view(torch.float32).mul_(1 - eff_weight_decay)
acc_m_u32.view(torch.float32).add_(other=v, alpha=-eff_lr)
acc_bf16_view_u16.copy_((acc_m_u32 >> 16).to(torch.uint16))
mantissa.copy_(acc_m_u32.to(torch.uint16))

即在 fp32 视图上执行权重衰减与更新，然后把结果拆回 bf16 主存储 + uint16 尾数，实现「bf16 存储、fp32 精度」的在线累积。相比 Record 6 的 Muon（每个尺寸分组维护一个 (world_size, size) 的 bf16 update_buffer 用于 all_gather 通信），新实现的通信量只随参数本身大小走，且把动量、正交化、更新全部收进一个 @torch.compile 的 update 核里。

4.2 分片（sharded）执行：每卡只算自己的参数

Muon.step 中，参数按 base_i + self.rank 的步长分发（第 100-113 行）：8 块卡各自负责一组参数完成动量 + Newton-Schulz 正交化（zeropower_via_newtonschulz5，5 次五次多项式迭代，系数针对 bf16 稳定性调优），随后用异步 all_gather（async_op=True + future 集合）把更新广播给所有 rank，实现与计算重叠。这正是 README 所说「New sharded mixed precision Muon」——计算分布化 + 通信异步化 + 精度混合化的组合拳。

4.3 移除 CastedLinear：回归朴素线性层

Record 6 中 FP8 算子通过 CastedLinear 封装，Record 7 把模型里所有线性层换成朴素的 F.linear + bf16 参数（如注意力里 F.linear(x, self.qkvo_w[:3]...)、MLP 里 F.linear(x, self.fc_w)），CastedLinear 类与 FP8 custom op 从代码中完全消失。原因与优化一相同：FP8 仅在 head 上有收益（训练后期大 logits 需要高精度梯度），在 MLP/QKV 上引入的量化与缩放开销不划算；移除后模型整体可以交给 torch.compile 更顺畅地融合，这也是 3.6ms/step 收益的主要来源之一。

五、优化四：合并 QKV 与 O 权重（+0.8ms/step）

注意力模块把原来分散的 Q、K、V、O 四个线性层合并为一个形状为 (4, hdim, dim) 的单一参数 qkvo_w：

self.qkvo_w = nn.Parameter(init_linear(torch.empty(4, hdim, dim)).bfloat16())
self.qkvo_w.detach()[3].zero_()  # out zero init（muP 风格的投影零初始化）

前向一次 F.linear(x, self.qkvo_w[:3].flatten(end_dim=1)) 就同时产出 Q、K、V 三个投影，O 投影走 F.linear(y, self.qkvo_w[3])。合并带来的收益是多重的：

• 把 4 次 kernel launch 压成 2 次（QKV 一次、O 一次），减少调度开销；
• 一次大 GEMM 比多个小 GEMM 更利于 tensor core 利用率与 torch.compile 的图融合；
• 参数集中后 Muon 的按尺寸分组与 sharded all_gather 通信模式更规整。

README 记录的收益为每步约 0.8ms。代码注释还标注了该思路的来龙去脉：「merged QKV weights」由 @fernbear.bsky.social 提出、@YouJiacheng 改进；O 投影零初始化则来自 @Grad62304977（muP 风格，让浅层初始时保持恒等映射）。注意力其余细节一并保留：QK-norm + 半截断 RoPE（Rotary，仅前 dim/4 维有频率）、可调 attn_scale=0.12（替代默认 head_dim**-0.5）、以及 ve（token value embeddings，012…012 的 U-net 式插入结构）按 sa_lambdas 混合进 V。

六、优化五：合并标量权重（+0.6ms/step）

模型把三类可学习标量——逐层 skip 权重（skip_weights，长度 = 层数 16）、逐层 block 混合系数（lambdas，每层 2 个：残差流与 x0 的加权）、逐层注意力 SA 混合系数（sa_lambdas，每层 2 个）——合并进单个参数张量 scalars：

self.scalars = nn.Parameter(torch.cat([
    torch.ones(num_layers),            # skip_weights
    *[torch.tensor([1.0, 0.0]) for _ in range(num_layers)],   # block lambdas
    *[torch.tensor([0.5, 0.5]) for _ in range(num_layers)],   # SA lambdas
]))

前向时用切片视图解包（第 292-294 行）：

skip_weights = self.scalars[:len(self.blocks)]
lambdas      = self.scalars[1*len(self.blocks): 3*len(self.blocks)].view(-1, 2)
sa_lambdas   = self.scalars[3*len(self.blocks): 5*len(self.blocks)].view(-1, 2)

合并后这 80 个标量在 AdamW 里只占一个「参数」，optimizer 的迭代、显存管理、all_reduce 通信都从「几十个小张量」收敛为「1 个大张量」，省下大量 kernel launch 与通信原语开销，实测每步约 0.6ms。这一招的普适性很强：任何按固定索引切片使用的独立小参数，都值得先 cat 成一个大张量再进优化器，本项目后续纪录（如 VeFused、GatesToCompiledAdam）延续了同样的思路。

七、完整训练配置与调度细节

Record 7 的超参数与调度都封装在 Hyperparameters 中，要点如下：

配置项	取值	说明
train_files	data/fineweb10B/fineweb_train_*.bin	训练数据（FineWeb10B 预处理成 .bin）
val_files / val_tokens	fineweb_val_*.bin / 10485760	验证 token 数固定，保证跨纪录可比
train_seq_len / val_seq_len	64×1024 / 4×64×1024	FlexAttention 序列长度；验证用更长序列
num_iterations	6450	相比 Record 6 的 6710 减少 260 步
cooldown_frac	0.6	训练后 60% 用于学习率线性冷却
val_loss_every	125	每 125 步评估一次验证损失

模型为 16 层、8 头、1024 维、词表 50257（lm_head 扩展到 128 的倍数 50304），批量 = world_size × train_seq_len（8×64K token）。参数分四个集合交给两个优化器（第 423-438 行）：

• hidden_matrix_params（各层 2D 矩阵，按 numel 降序）→ Muon，lr=0.025，momentum=0.95；
• embed_params（embed + 3 个 value_embeds）→ AdamW，lr=0.3；
• scalar_params（合并后的 scalars）→ AdamW，lr=0.015；
• head_params（lm_head_w）→ AdamW，lr=1/320。

AdamW 使用 betas=(0.8, 0.95)、eps=1e-10（小 epsilon 是 @YouJiacheng 提出的修正 world_size 依赖的替代方案）、无权重衰减、fused=True。学习率调度为「平台 + 冷却」两段式（前 40% 恒为 1.0，后 60% 线性降到 0）；Muon 的 momentum 在前 300 步从 0.85 线性预热到 0.95（第 548-550 行）。注意力窗口按块（128 token/块）从 128 线性增长到 1792（get_window_size_blocks，随训练进度 1728×x 取 128 的倍数）；长-短滑动窗口（long/short 两个 BlockMask，short 为 long 的一半）按 [long, short×11, long, short×3, long] 的 16 层模式分布，mask 由 FlexAttention 的 BlockMask.from_kv_blocks + 文档因果 mask 手工构造。此外保留特征：第 8 层（layer_idx 7）跳过注意力、第 9/10/11 层额外接入第 6/4/2 层的 skip 连接（skip_map）、中间层跳过 value embeddings（ve 列表中间为 None）、bf16 全模型 + torch.compile(dynamic=False)。

分布式侧：训练前先跑 10 步随机数据的 warmup 再恢复初始状态（保证计时公平）；正式训练里梯度用异步 all_reduce（AVG）与 opt.step() 重叠；数据加载使用自研 distributed_data_generator（每个 rank 独立 mmap 读取 shard，避免 host 侧同步）。

八、运行结果与验证

完整运行日志 223_3310d0b1-b24d-48ee-899f-d5c2a254a195.txt 记录了 8×H100、PyTorch 2.7.0.dev20250208+cu126（CUDA 12.6）环境下的全过程，关键数据：

• 单步均耗稳定在 ~222ms（早期）→ ~235ms（后期窗口变大），最终 step_avg:234.94ms；
• 验证损失轨迹：step 0 为 10.826 → step 125 为 4.504 → step 250 为 4.180 → step 6375 为 2.920470 → 最终 step 6450 val_loss:2.919352，低于 2.92 目标；
• 总训练时间 train_time:1515353ms，即约 25.26 分钟（与 README 记录的 25.29 分钟一致，差异来自计时口径）；
• 峰值显存 allocated 50380 MiB / reserved 55354 MiB（单卡），说明混合精度与权重合并在 80GB H100 上留有充足余量。

可以横向对照：Record 6 的基线（25.95 分钟）与 Record 8（25.29 分钟后下一纪录，075_640429f2-e726-4e83-aa27-684626239ffc.txt）都收录在 README.md 的 Track 2 纪录表 中，便于读者按序追踪每项改动的增量贡献。

九、如何查看与复现

• 阅读记录：本纪录的全部源码即日志文件 223_3310d0b1-b24d-48ee-899f-d5c2a254a195.txt 开头的 Python 代码（运行日志自嵌完整源码），README 目录 records/track_2_medium 下按日期组织全部 Track 2 纪录，可对照 2025-04-22_Record8 查看后续演进。
• 运行：仓库根目录 run.sh 与 train_gpt_medium.py 提供 Track 2 的当前训练入口（后者为后续纪录的集大成版本，已包含更多后续优化，注意与本纪录的差异）；数据侧可参考 data/cached_fineweb10B.py 生成 FineWeb10B 的 .bin 文件。环境要求 8×H100（代码断言 world_size == 8）与配套 PyTorch 版本，首次运行 torch.compile 会有额外编译耗时。
• 延伸阅读：README 的「World record history」「Speedrun track 2」两节给出完整纪录谱系与规则（含 Track 2 允许 coordinate_descent_tuning 的说明），是理解每一项改动上下文的最佳索引。

小结

Record 7 的五项改动勾勒出极速训练中「精度—速度—步数」三角平衡的典型解法：FP8 在 head 上是负资产（换回 bf16 反而省步数）、softcap 可以用代数等价形式白赚 1.5ms/step、Muon 通过「bf16 存储 + uint16 尾数 + 分片计算 + 异步通信」获得混合精度与分布式双重收益、而权重合并（QKV&O、标量）则是一类零损失的系统性提速。25.29 分钟、2.9194 验证损失的最终成绩，全部来自日志中可逐行核对的源码与数据——这正是 modded-nanogpt 社区「每条纪录都可复现、可比对」风格的直接体现。