# 2 分钟训练 124M 参数 NanoGPT：Muon 优化器与工程复现实践 > 深入解析 Muon 优化器的 Newton-Schulz 正交化机制，及其在 NanoGPT 速度挑战中实现 1.35 倍加速的工程落地细节。 ## 元数据 - 路径: /posts/2026/01/24/2-minute-nanogpt-muon-optimizer/ - 发布时间: 2026-01-24T01:47:06+08:00 - 分类: [mlops](/categories/mlops/) - 站点: https://blog.hotdry.top ## 正文 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 ## 同分类近期文章 ### [MegaTrain全精度单GPU训练100B+参数LLM：梯度分片与optimizer状态重构技术路径](/posts/2026/04/09/megatrain-full-precision-single-gpu-training-100b-llm/) - 日期: 2026-04-09T01:01:41+08:00 - 分类: [mlops](/categories/mlops/) - 摘要: 深入解析MegaTrain如何通过主机内存存储、流水线双缓冲执行引擎与无状态层模板，实现单GPU全精度训练百亿参数大模型的核心技术细节与工程化参数。 ### [可验证的 RLHF 合成数据流水线与质量评估框架](/posts/2026/04/08/synthetic-data-rlhf-pipeline-verification-framework/) - 日期: 2026-04-08T23:27:39+08:00 - 分类: [mlops](/categories/mlops/) - 摘要: 基于 LLM 生成奖励模型训练数据，构建可验证的合成数据流水线与质量评估框架。 ### [单GPU全精度训练百亿参数LLM：显存优化与计算调度工程实践](/posts/2026/04/08/single-gpu-100b-llm-training-memory-optimization/) - 日期: 2026-04-08T20:49:46+08:00 - 分类: [mlops](/categories/mlops/) - 摘要: 深度解析MegaTrain如何通过CPU内存作为主存储、GPU作为瞬态计算引擎，实现单卡训练120B参数大模型的核心技术与工程细节。 ### [Gemma 4 多模态微调在 Apple Silicon 上的实践：MLX 框架适配与内存优化](/posts/2026/04/08/gemma-4-multimodal-fine-tuner-apple-silicon/) - 日期: 2026-04-08T12:26:59+08:00 - 分类: [mlops](/categories/mlops/) - 摘要: 在 Apple Silicon 本地运行 Gemma 4 多模态微调，聚焦 MLX 框架适配与内存优化工程参数，提供可落地的配置建议。 ### [极简自蒸馏SSD：代码生成中单次训练无过滤的工程实践](/posts/2026/04/05/embarrassingly-simple-self-distillation-code-generation/) - 日期: 2026-04-05T12:26:02+08:00 - 分类: [mlops](/categories/mlops/) - 摘要: 深入解析Simple Self-Distillation方法，探讨训练温度、截断策略与代码生成pass@1提升之间的参数映射关系。