用 Snoo（Sparse Nesterov Outer Optimizer）做外层优化器：modded-nanogpt GPT-2 Medium 赛道 23.28 分钟记录的技术拆解

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

本篇技术指南围绕 modded-nanogpt 仓库 track 2（GPT-2 Medium）赛道的一份官方记录展开：社区成员 @dominikkallusky 通过引入Snoo（Sparse Nesterov Outer Optimizer，稀疏 Nesterov 外层优化器），把赛道记录从 23.45 分钟推进到23.28 分钟，比基线少用约 60 个训练迭代（约 10 秒）。文章将完整讲解 Snoo 的设计动机、逐行源码实现、与 AdamW / Muon 的嵌套组合方式、超参数配置，以及官方提交中采用的统计学验证方法，读完后你可以把 Snoo 直接套用到自己的大模型训练管线中。

背景：track 2 赛道与 Snoo 的引入位置

modded-nanogpt 的主赛道（track 1）目标是"在 8 张 NVIDIA H100 上用最短时间把模型训到 FineWeb 验证集交叉熵 ≤ 3.28"，而track 2（GPT-2 Medium）把目标收紧到 ≤ 2.92，对应 Andrej Karpathy 的 350M 参数 llm.c 基线（见 README.md 中的 "Speedrun track 2" 章节）。track 2 的基线时间长达 5.8 小时，此后被社区逐步压缩到 23 分钟量级。

在 Snoo 记录之前，2025-09-11_SecondInputEmbed（PR#124）刚把记录推进到 24.07 分钟。Snoo 记录（PR#128）正是在此基础上叠加一个外层优化器，把迭代数从 5700 降到5640，并宣称相比前一记录改善约 60 步、约 10 秒的墙钟时间。其完整训练日志保存在 records/track_2_medium/2025-09-16_Snoo/000_01db7a67-f715-4114-a7b5-6bfe23bac1b1.txt，该日志内嵌了训练脚本全量源码，是本文所有源码引用的直接依据。

Snoo 是什么：一个 look-ahead（前瞻）动量包装器

根据 2025-09-16_Snoo 的 README，Snoo 被定义为：

Snoo is a look-ahead momentum-based wrapper that can improve the quality of large language models (LLM) and other models. Snoo implicitly smoothens the training trajectory and instills a bias towards flatter minima. Snoo is computationally efficient, incurring minimal overhead in compute and moderate memory usage.

即：

• Look-ahead 包装器：它不取代内层优化器（AdamW / Muon），而是"包在"内层优化器外面，周期性把内层优化器 K 步累计的参数位移当作梯度，交给一个外层 SGD 去"追赶"；
• 稀疏更新（Sparse）：不是每步都做外层操作，而是每隔 K 步才触发一次，因此计算开销极低；
• Nesterov 动量：外层 SGD 显式启用了 Nesterov 动量（nesterov=True），与 look-ahead 思想天然契合；
• 效果取向：文档声称它隐式平滑训练轨迹、把优化结果引向更平坦的极小值区域（这是记录文档给出的经验性表述，属于该 PR 的实验观察）。

从算法谱系上看，这与经典的 Lookahead 优化器（RAdam + Lookahead 式的"慢权重"追踪）一脉相承，但 Snoo 的写法更精简：只有一个"快/慢"参数缓冲区，配合外层 Nesterov SGD。

Snoo 的完整实现：源码逐行拆解

记录 README 只贴出了step()方法，而完整类定义（含__init__、state_dict、load_state_dict）在训练日志源码中（见 000_01db7a67-f715-4114-a7b5-6bfe23bac1b1.txt）。下面给出完整实现并逐段讲解。

构造函数与内部状态

class Snoo: @torch.no_grad() def __init__(self, model: nn.Module, lr: float, momentum: float, k: int) -> None: self.model = model self.lr = lr self.momentum = momentum self.k = k self.current_step = 0 self.outer_buf = [p.clone() for p in model.parameters()] # Cache the parameter list to avoid recreating it every step self.model_params = list(self.model.parameters()) self.optimizer = torch.optim.SGD( self.model.parameters(), lr=lr, momentum=momentum, nesterov=True, fused=True, )

要点：

• outer_buf是"慢权重"缓冲：初始化时对模型全部参数做一次clone()。由于该记录全程用 bfloat16 训练，这里的缓冲也保持了模型的原始 dtype；
• self.optimizer是一个fused Nesterov SGD，它并不消费真实梯度，而是消费"参数位移"；
• model_params被缓存为列表，避免每个 step 反复调用model.parameters()，属于面向性能的细节。

step()：稀疏的 look-ahead 更新

@torch.no_grad() def step(self) -> None: if self.current_step % self.k == 0: # Use cached parameter list and combine loops for better performance for p_new, p_old in zip(self.model_params, self.outer_buf): p_new.grad = p_old.data - p_new.data p_new.copy_(p_old, non_blocking=True) self.optimizer.step() # Update outer buffer in-place for p_new, p_old in zip(self.model_params, self.outer_buf): p_old.copy_(p_new, non_blocking=True) self.current_step += 1

算法语义可以拆成四步理解：

1. 触发条件（稀疏性）：只有current_step % k == 0时才执行外层逻辑，其余迭代完全零开销，这正是"稀疏"二字的来源；
2. 构造伪梯度：对每个参数执行p_new.grad = p_old.data - p_new.data。这里p_old是慢缓冲（look-ahead 追踪的快照），p_new是内层优化器跑了 K 步后的当前参数，两者的差就是这 K 步的净位移——把它当作"梯度"注入当前参数；
3. 回退到慢权重：p_new.copy_(p_old)把当前参数重置回慢缓冲的值，然后调用外层 Nesterov SGD 的optimizer.step()，由 SGD 的动量机制决定本次最终落点（Nesterov 会在位移方向上前瞻一步，从而"看向"下一步的下降方向）；
4. 刷新慢缓冲：第二次循环把更新后的p_new写回p_old，完成一个 look-ahead 周期。

整个step()被@torch.no_grad()包裹，不参与 autograd，与内层优化器解耦。也可以从另一个视角看待：内层优化器（AdamW/Muon）负责"快速探索"，外层 Snoo 负责"周期性纠偏回锚点"，两者共同平滑训练轨迹。

断点续训支持：state_dict 与 load_state_dict

def state_dict(self): state_dict = { "current_step": self.current_step, "lr": self.lr, "momentum": self.momentum, "k": self.k, "outer_buf": [p.clone() for p in self.outer_buf], "optimizer_state_dict": self.optimizer.state_dict(), } return state_dict def load_state_dict(self, state_dict): self.current_step = state_dict["current_step"] self.lr = state_dict["lr"] self.momentum = state_dict["momentum"] self.k = state_dict["k"] for p_src, p_dst in zip(state_dict["outer_buf"], self.outer_buf): p_dst.copy_(p_src) self.optimizer.load_state_dict(state_dict["optimizer_state_dict"])

• 状态完整覆盖了步数、超参数、慢缓冲快照和外层 SGD 的动量状态，保证k周期的相位在断点续训后不丢失；
• 训练日志里也确实在 warmup 后用它做状态复位：先跑 10 步随机数据 warmup 内核，再initial_state保存、恢复（见 000_01db7a67...txt），确保 warmup 不污染正式计时。

与 AdamW + Muon 的嵌套组合：完整配置实例

Snoo 的价值在于"包住"多个内层优化器。训练日志中的装配方式如下（000_01db7a67...txt）：

hidden_matrix_params = sorted( (p for p in model.blocks.parameters() if p.ndim >= 2), key=lambda x: x.size(), reverse=True, ) embed_params = [*model.embed1.parameters(), *model.embed2.parameters(), *model.value_embeds.parameters()] scalar_params = [model.scalars] head_params = [model.lm_head_w] adam_param_groups = [ dict(params=head_params, lr=1 / 320), dict(params=embed_params, lr=0.3), dict(params=scalar_params, lr=0.015), ] inner_optimizers = [ torch.optim.AdamW(adam_param_groups, betas=(0.8, 0.95), eps=1e-10, weight_decay=0.0, fused=True) ] inner_hidden_optim = Muon(hidden_matrix_params, lr=0.027, momentum=0.95, rank=rank, world_size=world_size) inner_optimizers += [inner_hidden_optim] outer_optim = Snoo(model, lr=0.68, momentum=0.37, k=28) all_optimizers = [outer_optim] + inner_optimizers

关键参数速查表：

从结构可以推断，Snoo 的outer_buf覆盖所有模型参数，因此它对 AdamW 管理的参数和 Muon 管理的参数一视同仁地做周期回锚——这与后续 2025-09-30_SmoothedSnooMedium 的说明一致：Snoo 作用于每一个参数，而非仅 Muon 参数。

在训练循环中的调用顺序

训练循环里，内层优化器和外层优化器是顺序调用的（000_01db7a67...txt）：

for opt in inner_optimizers: torch.futures.collect_all(opt2futures[opt]).wait() opt.step() outer_optim.step() model.zero_grad(set_to_none=True)

• 先等内层优化器对应的梯度 all-reduce 完成，再执行inner_optimizers的step()；
• 然后调用outer_optim.step()——注意它不会真正消费梯度，只在步数满足% k == 0时改写参数；
• 最后zero_grad(set_to_none=True)，因此外层伪造的p_new.grad会在下一次反向传播时被覆盖，不会污染真实梯度。

另外训练循环还配合了两类调度：学习率采用"先稳定、后 cooldown"线性衰减（cooldown_frac=0.7），Muon 动量在 300 步内从 0.85 线性 warmup 到 0.95；滑动窗口大小按三次曲线从 128 块增长到 1792 块（见 000_01db7a67...txt）。Snoo 的超参数（lr=0.68、momentum=0.37、k=28）在日志中是固定值，未见随 step 变化的调度。

实验结果与统计学验证

官方 30 次运行的统计量

记录 README 附上了 30 次独立运行的完整数据表，并在表后给出了权威统计输出（注意 README 表格表头"Train time | Val Loss"与数值列实际有错位，应以代码打印的统计为准）：

df_nanogpt_med_5640['train_time'].mean() = 1404029.9333333333 # 约 23.40 分钟 df_nanogpt_med_5640['train_time'].std() = 1523.797296985019 # 约 1.52 秒 df_nanogpt_med_5640['train_time'].count()= 30 df_nanogpt_med_5640['train_time'].min() = 1402608.0 df_nanogpt_med_5640['train_time'].max() = 1410229.0 df_nanogpt_med_5640['loss'].mean() = 2.919616033333334 df_nanogpt_med_5640['loss'].min() = 2.917969 df_nanogpt_med_5640['loss'].max() = 2.921223 df_nanogpt_med_5640['loss'].std() = 0.0007156313095892154 scipy.stats.ttest_1samp(..., 2.92, alternative='less').pvalue = 0.0032015979421488247

要点解读：

• 30 次运行平均墙钟时间约1404 秒（≈ 23.4 分钟），标准差仅约 1.5 秒，重复性极好；
• 平均验证损失2.919616，低于赛道目标 2.92，且 30 次运行全部落在 2.918～2.921 区间（极差约 0.003），波动很小；
• 官方用scipy.stats.ttest_1samp做单侧检验（alternative='less'，原假设均值为 2.92），得到p = 0.0032 < 0.01，满足赛道"统计显著性 p < 0.01"的提交要求。

单次运行的训练曲线

训练日志 000_01db7a67...txt 记录了 5640 步全程的验证损失曲线，关键里程碑如下：

该次运行最终val_loss = 2.919060，峰值显存56505 MiB（约 55.2 GiB），与统计表的 2.9190～2.9196 区间吻合。日志环境为 Python 3.10.18、PyTorch 2.9.0.dev20250823+cu126、8 张 NVIDIA H100 80GB。

后续演化：Snoo + Muon 更新平滑（记录 #14）

Snoo 被接受后很快成为 track 2 的标准组件。紧接着的记录 2025-09-30_SmoothedSnooMedium 把它与 PR#129 的"Muon 更新 EMA 平滑"叠加（即把平滑后的 Muon 更新作为 Snoo 的内层优化器），迭代数进一步从 5640 降到5590，记录时间推进到23.08 分钟。该记录的方法论笔记还包含两个对 Snoo 的补充理解：

• 若k=1且外层 SGD 不用 Nesterov 动量，Snoo 与单纯的"更新平滑"在数学上等价——区别仅在于 Snoo 作用于所有参数而 PR#129 只作用于 Muon 参数；
• 该记录用 1000 次 bootstrap 抽样（每组 40 个样本）评估 p 值落在 0.01 以下的概率，报告 Snoo 单列（5640 步配置）的该比例为 32.10%，叠加平滑后提升到 85.40%——这是对"为什么 Snoo 本身只减少 60 步就足够稳健"的统计侧面印证。

这也说明 Snoo 在仓库后续版本中持续被复用：在 2025-10-04_GPT2MediumLayerReuse 和 2025-11-02-Smear-MTP 的训练脚本中都以相同的Snoo(model, lr=0.68, momentum=0.37, k=28)配置出现。

在自己训练管线中使用 Snoo 的实践建议

1. 即插即用：Snoo 类的依赖只有 PyTorch 标准库（torch.optim.SGD+ 参数克隆），可以原样复制到任何 PyTorch 训练脚本，不需要额外依赖；
2. 参数选择：仓库验证过的配置是lr=0.68, momentum=0.37, k=28；其中k直接决定稀疏度（28 意味着外层开销只占 1/28 的步数），lr应显著大于内层学习率；
3. 调用顺序：务必在内层优化器step()之后调用outer_optim.step()，并在zero_grad之前，以免伪造的grad影响后续反向传播；
4. 兼容性：fused=True的 SGD 要求 CUDA 支持，若在 CPU 或旧 GPU 上运行可去掉fused标志；state_dict/load_state_dict已完整实现，可安全用于 checkpoint；
5. 与 EMA 类方法的关系：若你已经在用参数 EMA 或更新平滑，Snoo 的收益可能与它们部分重叠（后续记录正是在二者组合中继续取得收益），建议用 ablation 判断是否叠加。

小结

Snoo 用不到 40 行代码实现了一个"稀疏、前瞻、动量包装"的外层优化器：内层 AdamW/Muon 负责快速优化，外层每隔 k 步用参数位移做一次 Nesterov 纠偏，从而以极低的计算与中等内存开销平滑训练轨迹。在 modded-nanogpt 的 track 2 赛道上，它把迭代数从 5700 压缩到 5640，记录推进约 60 步（约 10 秒），并以 30 次运行、p=0.0032 的显著性通过 2.92 验证损失门槛。其完整实现与运行日志均可在 records/track_2_medium/2025-09-16_Snoo 目录下查阅，作为直接可复现的参考实现。

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