企业官网建设流程全解析

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

导读

本篇文章围绕 modded-nanogpt 在 2025-12-14 创造的 128.8 秒（8×H100、124M 参数、fineweb10B）训练纪录展开，核心剖析该纪录中包含的五项更新，重点深挖其中的Partial Key Offset（部分键偏移）注意力机制改造：它如何通过"只平移 Key 的静止维"让一层注意力即可完成归纳（induction），从而在不显著增加计算的前提下压低验证损失。文章同时结合仓库源码（track_1_short/model/gpt.py、track_1_short/model/attention.py、track_1_short/perf/kernels/qkv_rope.py）给出可复现的实现细节、统计验证方法与完整训练脚本解读。读完本文，你将理解：为何 key offset 只作用在长滑动窗口与静止头维、x0_lambda合并的代数化简、batch size 调度与窗口调度如何精确对齐，以及这类"纪录型"训练脚本如何用严格的统计检验来验收每一次提速。

纪录背景：128.8s 是怎么来的

原纪录 README 的第一行即点明主题："New WR 128.8s: Partial Key Offset"（新的世界纪录 128.8 秒，靠的是部分键偏移）。该纪录由五项更新叠加而成，合计带来约 2.4 秒的提速，其分解如下：

1. Partial Key Offset：为长滑动窗口实现"部分键偏移"（详见下一节）；
2. 合并 layer 0 的残差缩放：把第 0 层的x_lambda*x + x0_lambda*x0化简为(x_lambda+x0_lambda)*x，并顺带清理代码，使 11 层模型的结构表达得更清晰；
3. 减少 50 步训练：每步约 60ms，直接节省约 3 秒墙钟时间（这是"更快"与"更低损失"之间的权衡取舍）；
4. 对齐 batch size 调度与窗口调度：让 batch size 的切换时刻与滑动窗口尺寸的更新时刻严格同步（因为 0.33 ≠ 1/3，原本两者存在错位）；
5. 初始 value embeddings 清零：对最终指标影响很小，但可能降低训练方差，同时符合"零初始化作为最低假设配置"的原则。

这些更新对应的完整可运行脚本保存在 records/track_1_short/2025-12-14_PartialKeyOffset/ 目录下（同名训练脚本有 11 份，为同一配置的不同运行副本），README 位于 records/track_1_short/2025-12-14_PartialKeyOffset/README.md。

注：以下验证损失、训练时间等数字均直接引自该纪录 README 与同目录训练日志，反映的是 2025-12-14 当日 8×H100 环境下的观测结果。

Partial Key Offset：一层注意力也能完成归纳

从"键偏移"说起

在仅含因果掩码的标准注意力里，token 位置 t 只能看到 ≤ t 的键。若想让注意力"抄写"前一个 token（例如序列abc...中预测b时要能直接关注到a的表示），通常需要多层叠加：第一层把相邻信息"搬运"进表示，后续层再利用位置编码实现归纳。而key offset的做法更直接：把 Key 张量在时间维上整体前移一个位置，使位置 t 的 query 恰好能对齐位置 t-1 的 key（对 stationary 维而言，k[t] = k[t-1]），从而在单层注意力内就建立了"相邻 token 键值对齐"的捷径，让模型仅凭一层即可完成 1-layer induction。

为什么只偏移"静止维"

本纪录的关键洞察是：key offset 只应用于静止的 head 维（stationary head dims），而不是所有维度。在 head_dim = 128 的配置下，前 64 维承载 RoPE 旋转（其中前 32 维为完整旋转频率、后 32 维为半截断旋转的零频率维），而后 64 维是"静止"的。README 中明确说明：

The partial key offset is only applied to the stationary head dims (32-64 and 96-128). This was found to perform better than applying it to all dims. This approach gives the queries more freedom to attend to multiple positions at once through a single key.

也就是说，只偏移静止维后，同一个 key 的旋转维仍携带自身位置的精确信息，而静止维则携带"上一个 token 的语义信息"。这样一来，query 可以同时通过旋转维关注当前位置、通过静止维关注上一位置，一个 key 同时服务于两个相邻位置，给 query 更大的注意力自由度；若把所有维度都偏移，则会破坏位置唯一性，效果反而更差。

实现代码（记录脚本内嵌版）

原 README 给出了最核心的两行实现，位于注意力 forward 中、RoPE 之后（见记录脚本150d40bf...txt第 943-946 行）：

if key_shift: # shift keys forward for the stationary head dims. Enables 1-layer induction. k[:, 1:, :, self.head_dim//4:self.head_dim//2] = k[:, :-1, :, self.head_dim//4:self.head_dim//2] k[:, 1:, :, self.head_dim//4+self.head_dim//2:] = k[:, :-1, :, self.head_dim//4+self.head_dim//2:]

以 head_dim=128 代入：head_dim//4 = 32、head_dim//2 = 64，所以这两行分别平移了[32:64)与[96:128)两个区间——恰好是两个静止的 32 维块，合计 64 维。k[:, 1:, ...] = k[:, :-1, ...]表示"位置 t 取位置 t-1 的值"，首位置（t=0）保持自身值不变。

何时启用：只给长滑动窗口

关键设计决策是并非所有层都启用 key shift，只作用于承担归纳任务的长窗口层。在记录脚本的GPT.forward中（150d40bf...txt第 1079-1081 行）：

bm_sizes = [short_bm, short_bm, short_bm, long_bm, short_bm, short_bm, None, short_bm, short_bm, short_bm, long_bm] assert len(bm_sizes) == self.num_layers key_shift = [b==long_bm for b in bm_sizes] # apply key shift to long windows

即 11 层中第 3 层与第 10 层为长窗口层（其余为短窗口，第 6 层无注意力），key shift 仅在这两个长窗口层开启。README 给出了背后的实验依据：

The lowest loss was achieved when applying it to every layer, but the cost:speed ratio seems best when only applying to the long windows, which are the ones primarily responsible for induction.

也就是说：全层应用能得到最低损失，但会引入额外开销；而长窗口正是归纳（induction）的主力，把 key offset 限定在长窗口层，能在"每步约 60ms"的极致速度预算下取得最好的性价比。README 还补充了消融结论：只对部分 head 子集应用 key offset 效果更差。

从纪录脚本到正式仓库：Triton 内核化

纪录时的实现是 PyTorch 的切片赋值（由torch.compile融合编译），而这一机制在后来的正式仓库中被下沉为 Triton kernel，并演进出更细的规则。这为我们理解其工程演化提供了绝佳参照。

正式模型中的开关与分层规则

在 track_1_short/model/gpt.py 中，模型现在用LONG_WINDOW_LAYERS显式声明长窗口层，并把 key offset 与窗口一一绑定（第 560-562 行）：

# sliding-window sizes and key shift: the long windows get the partial key offset bm_sizes = [ws_long if i in LONG_WINDOW_LAYERS else ws_short for i in range(self.num_layers)] key_offset = [i in LONG_WINDOW_LAYERS for i in range(self.num_layers)]

而 track_1_short/model/attention.py 第 37-38 行给出LONG_WINDOW_LAYERS = (3, 10)——与纪录脚本中的两个长窗口层完全一致；该文件中AttnArgs携带key_offset: bool字段，CausalSelfAttention.forward中还有一条重要前置断言（第 174 行）：

# The partial key offset shifts a key's non-rotating dims, so only a head with some can carry it. assert not key_offset or self.qk_dim > Yarn.ROTARY_DIM

即：只有 qk_dim 大于旋转维（64）的 head 才允许做 key offset——这与 README 中"只偏移静止维"的语义严格对应。

Triton 内核中的逐元素实现

正式实现位于 track_1_short/perf/kernels/qkv_rope.py 的_qk_norm_rope_forward_kernel（第 72-82 行）：

if KEY_OFFSET: # Partial key offset: a key's stationary (non-rotating) dims come from the previous token. shift_row = (input_head >= num_heads) & (token > 0) previous_token = tl.maximum(token - 1, 0) x_previous = tl.load( qk + previous_token[:, None] * stride_qkt + input_head[:, None] * stride_qkh + offs_d[None, :], mask=mask & shift_row[:, None], other=0.0, ).to(tl.float32) previous_rstd = tl.rsqrt(tl.sum(x_previous * x_previous, axis=1) / qk_dim + 1.1920928955078125e-7) shift = shift_row[:, None] & (offs_d[None, :] >= rotary_dim) y = tl.where(shift, x_previous * previous_rstd[:, None], y)

这里 QK 是打包成[tokens, 2*num_heads, qk_dim]的单个张量（前半是 Q、后半是 K，因此input_head >= num_heads即 K 行），内核在归一化与 RoPE 的同一个 epilogue 里完成 key offset：

• shift_row限定"只处理 K 行且 token > 0"；
• 从token-1行重新加载原始 QK 值，并按 qk_dim 重新计算 rstd（因为被平移的维同样经过 QK norm）；
• shift掩码offs_d >= rotary_dim精确限定只有非旋转（静止）维才被替换为上一 token 的值；
• 最后把旋转维与静止维拼回同一个 y 输出。

这一实现与 README 中"仅静止维、时间前移一位、首位置保持不动"的语义完全一致。此外，qk_norm_rope_forward外层（第 97-99 行）还有两个工程细节：assert not (paired and key_offset)（paired head 不兼容 key offset），以及block_m = 4 if key_offset else 8（偏移需要多读一行 token，因此减小 tile 尺寸）。向后兼容的PackedFP8QKVFunction也把key_offset作为参数贯穿 FP8 训练路径。

其余四项更新逐条拆解

1. 第 0 层残差缩放合并

原实现中每一层都执行x = resid_lambdas[i] * x + x0_lambdas[i] * x0，其中x0是归一化前的首层输入。由于第 0 层的x0_lambdas[0]恰好作用于x0自身，可以把两项合并为一次乘法。纪录脚本第 1110-1113 行：

if i == 0: x = (resid_lambdas[0] + x0_lambdas[0]) * x else: x = resid_lambdas[i] * x + x0_lambdas[i] * x0

这里resid_lambdas与x0_lambdas均来自同一个可学习scalars参数（初始化1.1 * ones(num_layers)与0 * ones(num_layers)，见150d40bf...txt第 1035-1047 行）。合并后第 0 层少一次乘加，且由于二者是同一个张量的相邻切片，(resid_lambdas[0] + x0_lambdas[0])可以提前融合计算；对 11 层模型的代码结构也更清晰。

2. 减少 50 步训练

训练步数由原来的 ~2160 步缩减为 2110 步（num_scheduled_iterations = 2070+num_extension_iterations = 40，见150d40bf...txt第 1329-1331 行）。按 README 所述"每步 60ms"计算，50 步约合 3 秒；考虑到 key offset 本身带来的额外计算，最终总时长仍达到 128.8s。这体现了该类纪录脚本的典型权衡：在不破坏最终验证损失的前提下，通过裁剪尾部低收益步数换取墙钟时间（验证部分见第四节）。

3. 对齐 batch size 与窗口调度

调度代码的核心是get_ws与get_bs均按同一比例x = step / num_scheduled_iterations分段（150d40bf...txt第 1442-1458 行）：

def get_ws(step: int): if step >= args.num_scheduled_iterations: return args.ws_final // 2, args.ws_final x = step / args.num_scheduled_iterations ws_idx = int(len(args.ws_schedule) * x) return args.ws_schedule[ws_idx] // 2, args.ws_schedule[ws_idx] def get_bs(step: int): if step >= args.num_scheduled_iterations: return args.train_bs_extension x = step / args.num_scheduled_iterations bs_idx = int(len(args.train_bs_schedule) * x) return args.train_bs_schedule[bs_idx]

三项调度（ws_schedule = (3, 7, 11)、train_bs_schedule = (8*2048*8, 16*2048*8, 24*2048*8)、LR 调度）都使用同一个分段边界。README 特别指出更新前的 bug：0.33 != 1/3——旧代码用0.33之类的近似常量做分段比例，导致 batch size 切换时刻与窗口尺寸更新时刻错位；修正为按int(len(schedule) * x)统一计算后，两个调度的切换点精确对齐，避免在窗口/批大小不匹配的过渡期浪费步骤。

4. 初始 value embeddings 清零

模型包含 3 个 token value embeddings（self.value_embeds，按012...012结构分布于各层），纪录脚本在第 1022-1023 行对它们做零初始化：

for embed in self.value_embeds: nn.init.zeros_(embed.weight)

README 的评价是：影响非常轻微，但可能降低训练方差，且零初始化符合"最低假设配置"（lowest assumption config）原则——初始时不向 V 注入任何先验信息，让网络自行决定何时、如何利用 value embedding。

计时与验证：用统计检验验收提速

纪录型训练脚本的验收必须同时回答两个问题：损失是否真的更低？计时是否可信？原 README 给出了可复制的验证脚本（scipy.stats+torch）：

import scipy.stats import torch losses = [3.2788,3.2774,3.2786,3.2792,3.2762,3.2769,3.2781,3.2778,3.2761,3.2783,3.2809] times = [128.892,128.907,128.912,128.844,128.822,128.869,128.818,128.882,128.886,128.95,128.946] print("p=%.4f" % scipy.stats.ttest_1samp(losses, 3.28, alternative="less").pvalue) # p=0.0004 print("losses:", torch.std_mean(torch.tensor(losses))) # losses: (tensor(0.0014), tensor(3.2780)) print("time:", torch.std_mean(torch.tensor(times))) # time: (tensor(0.0443), tensor(128.8844))

解读如下：

• 11 次独立运行的最终验证损失均值3.2780 ± 0.0014，全部落在 3.28 下方；
• 对 H0（真实均值 ≥ 3.28）做单侧 t 检验，p = 0.0004，远小于常规显著性阈值，说明"该配置显著低于 3.28"在统计上是稳健的；
• 11 次运行墙钟时间均值128.88 ± 0.044 秒，方差极小（约 0.3%），表明计时高度可复现。

README 还补充了旧纪录的复测数据：131.2: [131.270, 131.241, 131.213]，并注明"这次似乎拿到了一台略慢的机器"（appears I got a slightly slower machine this time）——即 128.8s 相对 131.2s 的约 2.4s 提升中，部分窗口是在不同批次硬件上测得的，这也是所有"秒级纪录"对比都必须声明的环境前提。

如何在当前仓库中查看与复现

该纪录的完整训练脚本保存在 records/track_1_short/2025-12-14_PartialKeyOffset/ 目录（含 README 与 11 份脚本运行副本，脚本开头会自动把自身源码写入日志以便审计）。复现要点与环境要求如下：

• 硬件：8×NVIDIA H100 80GB（记录日志中的nvidia-smi快照显示为 H100 80GB HBM3，CUDA 12.6 驱动 560.35.03）；
• 软件栈：Python 3.10.12、PyTorch 2.10.0.dev（CUDA 12.6 编译）、Triton 3.6.0；数据为 fineweb10B 的.bin分片（data/fineweb10B/fineweb_train_*.bin）；
• 分布式启动：脚本依赖torchrun设置RANK/WORLD_SIZE/LOCAL_RANK，要求world_size为 8 的约数（grad_accum_steps = 8 // world_size），PYTORCH_ALLOC_CONF="expandable_segments:True"已内置；
• 关键超参：num_iterations = 2110（2070 调度步 + 40 扩展步）、cooldown_frac = 0.55、block_size = 128、ws_schedule = (3, 7, 11)、ws_final = 13、train_bs_schedule = (8*2048*8, 16*2048*8, 24*2048*8)、val_tokens = 10485760；
• 模型：50257 词表（向上取整到 128 的倍数）、11 层、6 头、head_dim=128、model_dim=768；优化器为DistAdam（embed/scalars/lm_head，lr=0.008）与NorMuon（attn/mlp/gate，lr=0.023，weight_decay=1.2）双优化器，Muon 动量按步数做 warmup/cooldown，且 LR 与动量调度全部与步数联动。

若希望直接体验 key offset 在当前主干模型中的实现，可阅读 track_1_short/model/gpt.py（分层窗口与key_offset的绑定）、track_1_short/model/attention.py（AttnArgs.key_offset与CausalSelfAttention的消费），以及 track_1_short/perf/kernels/qkv_rope.py（Triton 内核中的逐元素实现与PackedFP8QKVFunction的 FP8 路径）。

小结

Partial Key Offset 是 modded-nanogpt 速度纪录迭代中的一个典型样本：它用"只偏移 Key 的静止维、且只作用于长窗口层"两个约束，把"单层归纳"这一能力以近乎零额外成本的代价嵌入注意力，换取更低的验证损失；再配合第 0 层残差缩放合并、步数裁剪、调度对齐与 value embedding 零初始化，四项小改进合计约 2.4s，最终达成 128.8s 的 124M 训练纪录。更重要的是，README 给出的统计验收脚本提醒我们：秒级纪录的每一个数字，都应该用多次运行的标准差与假设检验来背书。这份从实验到工程、再到统计验收的完整闭环，正是 modded-nanogpt 持续刷新纪录的方法论所在。

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