不加算力，只改一个算法：Muon 在万亿 MoE 模型中最高 2 倍加速

📌 一句话摘要

本文介绍了 Gram Newton-Schulz 算法，通过将 Muon 优化器的迭代过程转移到更小的 Gram 矩阵上，在万亿级 MoE 模型训练中实现了最高 2 倍的优化器加速。

📝 详细摘要

普林斯顿大学和纽约大学的研究者（包括 Tri Dao）提出了 Gram Newton-Schulz (GNS) 算法，旨在解决 Muon 优化器在 GPU 上进行极分解近似时的计算瓶颈。该方法的核心是将原本在权重矩阵上的迭代重构为在更小的 Gram 矩阵上进行，从而大幅减少 FLOPs 并能利用高效的对称矩阵乘法算子。针对 FP16 训练中的数值不稳定问题，研究团队提出了稳定版方案及重启策略。实验表明，在 Kimi K2 等万亿参数 MoE 模型中，该算法可将优化器耗时降低 40-50%，在 NVIDIA H100 和 B300 平台上实现显著的墙钟时间加速，且验证集精度无损。

💡 主要观点

- 核心思想是将迭代从权重矩阵转移到 Gram 矩阵。 通过数学等价变换，将原本在 n×m 空间上的迭代转移到 n×n 空间，利用维度缩减显著降低计算量。

针对 GPU 硬件特性进行了深度优化。 将低效的矩形矩阵乘法（GEMM）替换为高效的对称矩阵乘法 kernel，充分释放 GPU 算力。

解决了半精度训练下的数值稳定性难题。 引入了稳定版算法及重启策略，有效处理了 FP16 格式下 Gram 矩阵可能出现的负特征值和溢出问题。

在万亿级 MoE 模型中实现了即插即用的无损加速。 在 Kimi K2 等真实大规模训练场景中，优化器耗时缩短一半，且验证集困惑度误差保持在 0.01 以内。

💬 文章金句

- 不再直接在矩阵 X 上迭代，而是在更小的 Gram 矩阵 XXᵀ 上迭代，从而降低计算量并利用对称矩阵计算优化。

Gram Newton-Schulz 可以作为 Muon 中 Newton-Schulz 的即插即用替代方案：我们观察到验证集困惑度几乎没有变化，误差在 0.01 以内。
这是我最喜欢的一类工作 —— 线性代数洞察 + 高性能算子。

📊 文章信息

AI 评分：92

来源：机器之心

作者：机器之心

分类：人工智能

语言：中文

阅读时间：7 分钟

字数：1583

标签： Muon 优化器, Newton-Schulz, MoE 模型, 模型训练优化, 线性代数

阅读完整文章

不加算力，只改一个算法：Muon 在万亿 MoE 模型中最高 2 倍加速

🤖 問 AI