在 modded-nanogpt 中把 YaRN 融入滑动窗口训练调度：原理、实现与 159.3s 跑分的验证

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

本篇技术指南以 2025-09-10_Yarn 实验记录 为主体，系统讲解 modded-nanogpt 如何在 124M 参数 NanoGPT 的百秒级训练竞速中引入 YaRN（Yet another RoPE extensioN，论文 arXiv:2309.00071）来平滑注意力窗口的逐级放大，并结合旋转频率缩放与注意力 logit 缩放两套机制消除窗口切换时的 loss 尖峰。读完你将掌握：窗口调度为什么会在 RoPE 上引发分布外问题、YaRN 与 attn_scale 的数学公式与常数调参、单份旋转嵌入表的存储优化，以及最终验证窗口（ws_validate）的实验设计方法。

背景：更长的注意力窗口与 RoPE 的"分布外"问题

在 modded-nanogpt 的训练竞速中，一个朴素的观察是：更长的注意力窗口训练更慢，但能产出更低的 loss。因此理想方案不是全程用长窗口，而是在训练早期用短窗口加速、后期逐步放大。但这引出了 RoPE（Rotary Position Embedding）的两个真实痛点：

1. 低频率维度遭遇陌生旋转角。RoPE 各维度旋转频率不同，训练时模型只见过某个旋转角度范围内的位置编码。例如一个每位置旋转 0.1 度的维度，在窗口大小 384（ws 384）的训练中累计只会旋转0.1 * 384 = 38.4度；当滑动窗口一次性放大到 896 时，同一维度可能旋转到89.6度——这是从未见过的"分布外"数值，会带来临时的 loss 尖峰。
2. K/Q 归一化下的注意力困惑度随 key 数量增长。当 K、Q 向量做了 norm（QK norm）后，注意力窗口内的 key 数量越多，注意力机制的平均困惑度越高。此时对softmax(d * QK)施加一个缩放因子d，可以随窗口内 key 数量增长重新控制注意力分布的困惑度。

这两点正是 YaRN 要解决的核心矛盾：既要通过放大窗口拿到更低 loss，又要让 RoPE 频率分布"跟得上"窗口变化，避免切换时刻的验证损失崩坏。

单份旋转嵌入表：减少更新开销与内存占用

该记录里一个重要的工程优化是：只在模型根部保存一份旋转嵌入（rotary embeddings）的拷贝，而不是在每个注意力层各存一份。收益有三：减少窗口切换时的更新耗时、降低显存占用、潜在改善缓存局部性。实现代码如下：

# store single copy of rotary tensors angular_freq = (1 / 1024) ** torch.linspace(0, 1, steps=head_dim//4, dtype=torch.float32) # half-truncate RoPE by @YouJiacheng (w/ base freq tuning) angular_freq = torch.cat([angular_freq, angular_freq.new_zeros(head_dim//4)]) t = torch.arange(self.max_seq_len, dtype=torch.float32) theta = torch.outer(t, angular_freq) self.rotary_cos = nn.Buffer(theta.cos(), persistent=False) self.rotary_sin = nn.Buffer(theta.sin(), persistent=False)

这里有两个细节值得展开：

• 频率基线：angular_freq从(1/1024)^0 = 1线性衰减到(1/1024)^1，即最大角频率为每位置 1 弧度（约 57 度），这是对"1024^(0..1)"基线的直接实现；
• half-truncate RoPE（@YouJiacheng）：把head_dim // 4个有效频率向量复制到head_dim // 2长度后，再拼接等长的零向量，实现"只旋转一半维度"的半截断式 RoPE，配合基频调参使用；
• 使用nn.Buffer(..., persistent=False)声明非持久化缓冲区，cos/sin 表预计算好存为 buffer，训练中只做就地更新。

这一设计在当前仓库中已演进为 track_1_short/model/attention.py 中的Yarn类：它按 head 宽度（64/128）分别维护 cos/sin 因子表，支持"配对头"（paired heads）的2t/2t+1打包布局，并把"半截断"实现为ROTARY_DIM = 64常量（64 宽 head 全部维度旋转，128 宽 head 只旋转前半部分，其余保持固定）。

两套缩放：频率缩放（YaRN）与注意力缩放（attn_scale）

记录中最重要的调参结论是把两个常数做了经验修正：

1. YaRN 公式的 0.1 常数改为 0.2。原论文0.1 * log(curr/prev) + 1中的 0.1 在实证中不够用；
2. attn_scale 起点从 0.12 改为 0.1。原恒定缩放 0.12 改为从 0.1 起步、随窗口增长按对数累积上升，使训练全程的分布均值与原先相当，取值落在 0.1 到 0.14 之间。

# scale attention factor f in attn=softmax(f*qk) logarithmically with window size windows = list(dict.fromkeys(args.ws_schedule + [args.ws_validate])) scale_factors = [0.2 * math.log(curr / prev) + 1 for prev, curr in zip(windows[:-1], windows[1:])] # start with 0.1, inspired by 0.12 from @leloykun and learnable scalars used by @brendanh0gan attn_scales = list(accumulate([0.1] + scale_factors, lambda acc, factor: acc * factor)) self.attn_scales = dict(zip(windows, attn_scales))

注意scale_factors的构造方式：窗口序列去重（dict.fromkeys保留顺序）后，相邻两档窗口计算0.2 * log(curr/prev) + 1，再以 0.1 为起点做累乘得到每档窗口对应的注意力缩放系数。这样softmax(f * QK)中的f会随着窗口档位逐级放大，恰好抵消"key 数量增加导致注意力困惑度上升"的效应。0.1 的起点灵感来自 @leloykun 的 0.12 以及 @brendanh0gan 使用的可学习标量。

YaRN 频率更新的实现：alpha/beta 保持论文默认

YaRN 本身实现非常直接，该记录把alpha和beta保持在论文默认值 1 和 32（原论文针对 Llama 调参），并指出：YaRN 引起的频率更新在窗口从 3 切到 7、以及维度 5 到 10 之间最为显著。

def apply_yarn(self, old_window: int, new_window: int, alpha: int=1, beta: int=32): rotations = args.block_size * old_window * self.angular_freq / (2 * torch.pi) scaling_factor = old_window / new_window interpolation_weight = torch.clamp((rotations - alpha) / (beta - alpha), 0, 1) self.angular_freq *= scaling_factor + interpolation_weight * (1 - scaling_factor) t = torch.arange(self.max_seq_len, dtype=torch.float32, device=self.angular_freq.device) theta = torch.outer(t, self.angular_freq) self.rotary_cos.copy_(theta.cos()) self.rotary_sin.copy_(theta.sin())

逐行解读：

• rotations：把角频率折算成"在旧窗口长度上累计旋转了多少圈"，args.block_size * old_window是旧窗口对应的 token 数，除以2π得到圈数；
• scaling_factor：old_window / new_window，窗口放大时小于 1，表示需要把高频维度的旋转频率压缩以"延展"位置空间（NTK 式缩放）；
• interpolation_weight：按(rotations - alpha) / (beta - alpha)裁剪到 [0,1]，实现"维度旋转圈数在 alpha 以下完全按缩放因子压缩、超过 beta 完全不压缩、中间线性过渡"的 YaRN 插值；
• 最终angular_freq *= scaling_factor + interpolation_weight * (1 - scaling_factor)逐维度混频，并就地重建 cos/sin 表。

这一函数在当前仓库 track_1_short/model/attention.py 的Yarn.apply中继续存在，并且把注意力缩放更新也合并进同一次调用：self.attn_scale *= 0.2 * math.log(new_window / old_window) + 1——注意当前代码直接用math.log而不再依赖累积列表，逻辑等价于记录中的累乘公式。

用更长的窗口做最终验证：ws_validate

ws_validate参数允许在比训练更长的注意力窗口上做最终验证。该记录中训练最终窗口为 11，而验证窗口设为 13，对应的调度函数为：

def get_ws(step: int): if step == args.num_iterations: return args.ws_validate x = step / (1 + args.num_iterations) assert 0 <= x < 1 ws_idx = int(len(args.ws_schedule) * x) return args.ws_schedule[ws_idx]

要点是：训练步骤严格落在ws_schedule的离散档位上（按进度比例索引），唯独最后一步单独返回ws_validate，从而在不改变训练过程的前提下，让最终评估暴露模型在更长上下文下的真实能力。这一设计在当前仓库中演化为 track_1_short/schedule.py 的TrainingSchedule与 track_1_short/config.py 的WS_POST_YARN_EXT = 20（最终验证把长窗口再扩展到 20 个 block），并在 train_gpt.py 中通过limit_yarn_rebuild/complete_yarn_tables控制旋转表的重建时机：窗口变化只重建训练步实际读取的行，验证前再把整张表补齐，避免在计时区间内全量重建。

训练超参微调：cooldown_frac 与步数

该记录同时清理了训练收尾阶段的配置：

• cooldown_frac 从 0.45 提高到 0.5，以配合总步数从 1705 降到 1670，遵循"固定冷却步数"的经验法则（步数变少，冷却占比就得变大，才能保住同样数量的 LR 衰减步）；
• 降到 1695 步以下还有二级收益：少读一次（第 9 个）数据文件，大约节省 200ms 的墙钟时间——这类"步数-文件读取边界对齐"的微优化在百秒级竞速中直接换算成最终成绩。

窗口档位（long/short）从 (3/1) → (7/3) → (11/5) 逐级放大；当前仓库的 track_1_short/config.py 已把窗口粒度从 token 改为BLOCK_SIZE = 128token 的 block 单位（如window_sizes=(1,3)→(3,7)→(5,11)→(6,13)），并配套了 batch size 8→16→24→20→8 的调度，可见同一套"多阶段窗口放大 + YaRN 平滑"思想仍在延续。

实验结论：YaRN 消除了窗口切换的 loss 尖峰

记录用对照实验证实了两件事：

1. 没有 YaRN 时，注意力窗口从 3 突然放大到 7，验证 loss 会出现明显尖峰；引入 YaRN 频率缩放后尖峰消失，这是该 PR 能成立的关键证据；
2. 把最终验证窗口从 11 扩展到 13，loss 大约有0.0015 的提升。而且一个有意思的经验观察（@varunneal 提出）：奇数窗口表现显著更好。可能的解释包括：短窗ws_short = ws_long // 2的整除在奇偶长度下行为不同；短窗低于长窗的 50% 时表现更好；或模型学会了长/短窗口的固定比例、比例被大幅改变时表现变差——记录明确说明"可能还有完全不同的解释"，属于未定论经验。

消融实验（ablation）覆盖了四项：

• new_record：本 PR 完整配置；
• no_attn_scale：保持恒定 attn scale 0.12（关掉注意力缩放）；
• no_freq_scale：保持恒定旋转频率（基于 1024^(0..1)，关掉 YaRN 频率缩放）；
• prior_record：先前记录 + 更新后的步数 1705→1670 与 cooldown frac 0.5。

通过逐项开关，可以量化"频率缩放"与"注意力缩放"各自对最终 loss 的贡献，这是百秒级训练里做归因分析的标准打法。

统计验证：显著性检验与均值方差

该记录用 7 次重复运行做统计验证，证明 159.3s 的跑分成绩是可靠的：

import scipy.stats import torch accs = [3.2779, 3.2779, 3.2789, 3.2778, 3.2789, 3.2785, 3.2806] times = [159.447, 158.998, 159.467, 159.191, 159.503, 159.259, 159.468] print('p=%.4f' % scipy.stats.ttest_1samp(accs, 3.28, alternative='less').pvalue) # p=0.0053 print('acc:',torch.std_mean(torch.tensor(accs))) # acc: (tensor(0.0010), tensor(3.2786)) print('time:',torch.std_mean(torch.tensor(times))) # time: (tensor(0.1897), tensor(159.3333))

结论：验证 loss 均值 3.2786（标准差 0.0010），对 3.28 基准做单侧 t 检验p = 0.0053，统计显著；7 次运行耗时均值 159.33s（标准差 0.19s），方差极小，说明成绩可复现。

未来方向：记录中提出的四个开放问题

该记录末尾给出四点后续思考，至今仍有参考价值：

1. 优化器动量刹车：训练像"没有刹车的赛车"——窗口放大、模型"改变方向"时，若能有效压制优化器状态（momentum）中的动量项，或许能减少瞬时损失；作者在仅 Muon 参数上的初步实验结果是负面的；
2. 分布式 bos token 查找：8 个数据文件若每张 GPU 只负责 1 个文件而不是全部 8 个来定位 bos token 索引，在零开销假设下最多可省200ms * 7 ≈ 1.4s；
3. RoPE 起始角频率：把最大角频率定为每位置 1 弧度（57 度）看起来有些任意，但实验把该值升到 180 度并未带来性能提升；
4. 更细粒度的监控：每 125 步才画一次验证 loss 会掩盖窗口切换瞬间的 loss 尖峰，更细粒度的监控对这类竞速调优更有价值。

源码落点速查

这套"短窗起步 + 逐级放大 + YaRN 频率缩放 + attn_scale 对数累积 + 更长验证窗"的组合拳，是 modded-nanogpt 在约 159 秒内完成 124M 参数模型训练的关键一环；其"每个常数都有实证依据、每个改动都有消融对照"的做法，也值得任何追求极致训练效率的工程团队借鉴。

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