Andrej Karpathy 在 2024 年初用 llm.c 仓库训练 GPT-2（124M 参数）在 8×H100 GPU 上耗时约 45 分钟达到 3.28 验证损失。这一基准在短短三个月内被 Keller Jordan 及其合作者通过 modded-nanogpt 项目压缩至 3 分钟以内，此后记录继续下探至 2 分 20 秒，整体训练时间缩减超过 20%。这场被社区称为「NanoGPT 速度挑战」的竞赛，核心并非算力堆砌，而是对训练工程每一个环节的极致压榨。本文聚焦其中最具普适价值的突破点 ——Muon 优化器，从算法原理出发，剖析其正交化机制的实现细节，并给出可复现的参数配置与工程落地清单。
从 SGD 到 Muon：梯度正交化的工程动机
传统 SGD-Momentum 与 Adam 系列优化器在处理 transformer 语言模型的 2D 参数时，产生的更新矩阵往往呈现极高的条件数。这意味着更新方向高度集中在少数奇异值对应的奇异向量上，而大量「稀有方向」虽然对模型学习至关重要，却因更新幅度微弱而被抑制。Muon（MomentUm Orthogonalized by Newton-Schulz）的核心洞察在于：对 SGD-Momentum 产生的更新矩阵进行事后正交化处理，可有效放这些稀有方向的尺度，从而在保持单步计算开销可控的前提下提升样本效率。
正交化的数学目标是将更新矩阵 $G$ 替换为与其最近的正交或半正交矩阵 $UV^T$，其中 $USV^T$ 为 $G$ 的奇异值分解。最直接的方案是完整 SVD，但其在 GPU 上的计算开销过高，无法接受。Newton-Schulz 迭代提供了一种数值稳定且可完全在 bfloat16 精度下运行的替代方案，通过五步迭代即可将更新矩阵收敛至正交近似。值得注意的是，Muon 仅作用于 2D 参数（线性层权重、卷积核经展平后同样适用），标量与向量参数（如偏置、LayerNorm 权重）仍需配合 AdamW 使用。这一混合策略在实验中表现最佳，尤其在 embedding 层与输出分类头处，AdamW 的使用对最终性能有显著影响。
Newton-Schulz 迭代的系数调优与 FLOP 开销分析
Muon 论文给出的 Newton-Schulz 迭代系数为 $(a, b, c) = (3.4445, -4.7750, 2.0315)$，这一组数值经过约束优化得出：目标是在确保迭代收敛至 $[0.7, 1.3]$ 区间的同时，最大化 $a$ 以加速小奇异值的收敛。五步迭代对于当前主流模型规模已足够。更早的基线系数 $(2, -1.5, 0.5)$ 亦可工作，但收敛速度与最终精度略逊于调优后的版本。
针对工程落地，一个关键问题是 Muon 带来的额外计算开销究竟多大。设模型维度为 $m$，批次 token 数为 $B$，迭代步数为 $T$（通常取 5），则 FLOP 开销上界为 $Tm/B$。在 NanoGPT 速度挑战的典型配置下（$m=768$，$B=524288$），开销约为 $0.7%$；即使扩展至 Llama 405B 级别（$m=16384$，$B=16000000$），开销仍控制在 $0.5%$ 以内。这一 overhead 在端到端训练中几乎可忽略，前提是 Newton-Schulz 迭代本身已通过 Triton 等内核实现充分融合至计算图。
工程复现清单：从零构建 2 分钟训练流水线
基于 modded-nanogpt 仓库的公开记录与社区讨论，以下为可操作的复现参数与关键配置。首要步骤是环境准备：Python 3.10+，PyTorch 2.1+，Triton 3.0+，以及 CUDA 12.2 以上版本。依赖可通过 pip install -r requirements.txt 一键安装，其中核心依赖包括 triton、accelerate 与 muon（需从 GitHub KellerJordan/Muon 安装）。训练数据建议使用 FineWeb 预处理后的 tokenized 格式，单文件或分片存储不影响性能，关键在于预取与解码阶段的流水线并行度配置。
优化器配置采用混合策略：embedding 层与 lm_head 使用 AdamW（lr=6e-4，beta=(0.9, 0.95)，weight_decay=0.1），其余全部 2D 参数使用 Muon（lr=6e-4，momentum=0.95，nesterov=True，ns_steps=5）。学习率调度建议采用线性 warmup 20 步加余弦退火至 10% 峰值。梯度裁剪全局范数设为 1.0，混合精度训练启用 FP16 或 BF16，梯度累积步数根据显存调整以确保有效 batch_size 在 524288 token 左右。序列长度在复现初期可先使用 1024 以降低调试难度，后续逐步提升至 2048 甚至更长以逼近极限记录。
落地监控指标与常见瓶颈排查
训练过程中应重点监控三类指标：吞吐（tokens/sec）、损失曲线（train/val loss）以及各层梯度范数统计。Muon 的正交化效果可通过监测每步更新矩阵的 Frobenius 范数与奇异值分布间接验证 —— 理想情况下，正交化后的更新矩阵奇异值应集中在 1.0 附近，偏离度过大可能提示学习率设置不当或 momentum 配置有误。常见性能瓶颈包括：数据加载线程不足导致的 CPU 瓶颈（建议 prefetch_factor≥4，num_workers≥8）；Triton 内核未命中导致的注意力计算拖慢（可通过 PyTorch Profiler 确认）；以及梯度同步在多 GPU 场景下的通信重叠不足（检查梯度压缩与 NCCL 配置）。
若在单卡 4090 上复现，2 分钟目标需配合更激进的序列长度与 micro-batch 调度，社区经验表明 4090 复现 sub-5 分钟为较现实预期。若拥有 8×H100 或等效算力，则完整 2 分钟记录在正确配置下可达。过程中若出现 loss 不收敛或发散，首查 Muon 应用范围是否误将 embedding 层纳入（该层应使用 AdamW），次查学习率是否过高或 warmup 步数不足。
可迁移性讨论：Muon 在更大规模训练中的适用性
Muon 的设计初衷针对中小型 transformer（百 M 至数 B 参数），其在更大规模训练中的迁移性仍属开放问题。当前实验证据表明 Muon 在 1.5B 参数模型上仍保持样本效率优势，但 20B+ 级别尚无充分验证。核心挑战在于 Newton-Schulz 迭代的分布式实现 —— 当参数分片存储于多台机器时，正交化操作需跨节点通信完成，同步开销可能抵消其计算优势。社区正在探索分块正交化与流水线并行结合的方案，但尚未形成成熟最佳实践。对于生产级 LLM 训练，保守建议仍以 AdamW 为主，Muon 可作为实验性技术在小规模验证或特定子模块（如 FFN 层）中试点。
另一个值得关注的边界条件是微调与强化学习场景。Muon 的正交化机制对预训练阶段的高方差梯度更新有理论适配性，但下游任务的梯度分布可能不同，当前缺乏系统性研究。实践者若计划将 Muon 应用于微调，建议在验证集上对比 AdamW 基线，以实际指标而非理论推断为准。
结语
NanoGPT 速度挑战的意义远超刷新记录本身 —— 它建立了一套以「可复现基线 + 社区竞赛」为驱动的优化验证框架，使新技术的评估不再依赖研究者主观声称的基线调优水平。Muon 优化器是这一框架下最闪耀的成果之一，其 Newton-Schulz 正交化机制为二阶优化器的工程化提供了新思路。对于希望将相关技术落地的工程师，建议从单卡复现开始，逐步引入混合精度、数据流水线与多 GPU 扩展，在每个阶段明确监控吞吐与收敛指标的权衡，而非盲目追求极限数字。Muon 的工程价值不在于通杀所有训练场景，而在于它展示了一条「以算法洞察驱动工程极致化」的可行路径。
参考资料
- Keller Jordan. Muon: An optimizer for hidden layers in neural networks. https://kellerjordan.github.io/posts/muon/
- Keller Jordan. modded-nanogpt: NanoGPT (124M) in 2 minutes. https://github.com/KellerJordan/modded-nanogpt