Gram Newton-Schulz：针对 Muon 优化器的快速、硬件感知型算法解读

背景

随着 Kimi K2 Thinking1 和 GLM-5.2 等最先进语言模型的训练，Muon 优化器正逐渐成为行业首选。与广泛使用的 AdamW 相比，Muon 在达到相同损失值时所需的优化步数更少，这意味着它在优化质量上具有显著优势。然而，这种优势是以更高的单步计算成本为代价的。

Muon 的单步计算开销主要源于其核心的 Newton-Schulz 正交化过程。这是一个立方时间复杂度的矩阵运算，这是传统优化器（如 SGD 或 AdamW）所不具备的。传统优化器通常执行元素级操作（如更新动量或根据二阶矩重新缩放），对于大小为 $n \times m$ 的权重矩阵，其时间复杂度为 $O(mn)$。相比之下，Muon 以及 Scion、Dion、SOAP、Shampoo、SPlus 等现代优化器，利用正交化或更高阶的预条件技术来计算权重更新，这涉及矩阵乘法，时间复杂度为 $O(mn^2)$（假设 $n \le m$）。

尽管 $O(mn^2)$ 的运行时是此类算法不可避免的成本，但在 FLOPs（浮点运算次数）和实际墙钟时间（wall clock time）上仍有巨大的优化空间。根据训练设置（全局批次大小、集群规模和并行设置），Newton-Schulz 过程占据了端到端训练总耗时的 2% 到 17%。随着模型规模的扩大，特别是对于拥有大量细粒度专家（fine-grained experts）的混合专家（MoE）架构，这种开销增长迅速。

核心内容

为了解决上述瓶颈，研究团队提出了 Gram Newton-Schulz，这是对标准 Newton-Schulz 例程的重构，旨在将万亿参数 MoE 模型（如 Kimi K2）中的优化器时间减少高达 50%。

1. 标准 Newton-Schulz 的局限性

通常实现的 Newton-Schulz 例程存在以下主要缺陷：

• 矩形矩阵乘法主导成本：它使用多达十次 $n \times m$ 矩阵的乘法，每次消耗 $2mn^2$ 个 FLOPs。大多数流行架构的权重矩阵是矩形的（$m \gg n$），而 MoE 架构的矩阵甚至更加细长。因此，矩形矩阵乘法成为了主要的成本来源。
• 未利用对称结构：计算过程中产生的许多中间矩阵是对称的，但标准实现并未利用这一结构优势，导致一半的计算工作是冗余的。
• 硬件适配不足：它使用 cuBLAS 进行批量矩阵乘法/加法，而 cuBLAS 并未针对最新的 Hopper GPU 架构进行完全优化。

此前的改进工作多集中于优化多项式系数或归一化步骤以减少迭代次数，或针对对称矩阵乘法使用专用例程，但未能全面解决上述问题，尤其是针对高度矩形矩阵和 GPU 硬件特性的优化。

2. Gram Newton-Schulz 的核心创新

Gram Newton-Schulz 的核心思想是改变迭代对象：不再直接在矩形输入矩阵 $\mathbf{X} \in \mathbb{R}^{n \times m}$ 上进行迭代，而是迭代小的正方形对称 Gram 矩阵 $\mathbf{XX^\top} \in \mathbb{R}^{n \times n}$。这一转变带来了以下好处：

• 降低 FLOP 成本：由于 $n \times n$ 矩阵远小于 $n \times m$ 矩阵，计算量显著下降。
• 启用专用内核：允许更多地使用针对对称矩阵优化的 GEMM（通用矩阵乘法）内核。

3. 算法演进与稳定性

研究团队提出了两个阶段的算法演进：

• Naive Gram Newton-Schulz（朴素 Gram Newton-Schulz）： 将标准 Newton-Schulz 重写为数学上等价的形式，但主要在 $n \times n$ 矩阵空间内操作。除了预处理步骤（形成 $\mathbf X \mathbf X^\top$）和后处理步骤（乘以 $\mathbf X$）需要矩形矩阵乘法外，其余迭代步骤均作用于较小的对称矩阵。

• Stabilized Gram Newton-Schulz（稳定化 Gram Newton-Schulz）： 在对 Naive 版本进行数值性质研究时，发现使用半精度浮点运算时存在数值不稳定性，特别是 Gram 矩阵中可能出现虚假的负特征值。为此，团队引入了一种“重启”（restarting）策略，即在算法中途重建 Gram 矩阵，从而解决了不稳定性问题。

4. 硬件加速实现

为了充分利用最新一代 GPU（Hopper 和 Blackwell 架构），团队使用 CuTeDSL 实现了自定义的对称矩阵乘法内核。这些内核针对特定硬件架构进行了深度优化，达到了最先进的性能水平。

5. 最终成果：GramMuon

通过将 Muon 的 Newton-Schulz 例程替换为 Gram Newton-Schulz，团队开发了名为 GramMuon 的优化器。实验表明，GramMuon 将正交化步骤的运行时间减少了 40-50%。更重要的是，训练语言模型时保持了稳定性，且优化质量与标准版本相比，验证集困惑度（perplexity）差异在 0.01 以内，实现了近乎“免费午餐”的性能提升。

关键要点

• 性能提升显著：GramNewton-Schulz 将 Muon 优化器的正交化步骤运行时间减少了 40-50%，在万亿参数 MoE 模型中效果尤为明显。
• 算法重构逻辑：通过将迭代对象从矩形矩阵 $\mathbf{X}$ 转换为较小的对称 Gram 矩阵 $\mathbf{XX^\top}$，大幅降低了计算复杂度并利用了硬件对对称矩阵乘法的优化能力。
• 数值稳定性保障：针对半精度计算中的数值不稳定问题（如虚假负特征值），引入了“重启”策略，确保了算法在实际训练中的鲁棒性。
• 硬件感知优化：使用 CuTeDSL 为 Hopper 和 Blackwell GPU 架构编写了专用的对称矩阵乘法内核，突破了通用库（如 cuBLAS）的性能瓶颈。
• 无损优化质量：GramMuon 在保持与标准 Muon 几乎相同的优化质量（验证困惑度差异 < 0.01）的前提下，实现了训练速度的大幅提升。
• 开源贡献：团队发布了数学等价、数值稳定且速度更快的 Muon 替换实现，以及独立的快速 GPU 对称矩阵乘法内核。

意义与影响

Gram Newton-Schulz 的提出标志着大模型训练优化器工程化的一个重要进步。

首先，它解决了 Muon 等高级优化器在大规模训练中的可扩展性瓶颈。随着模型向万亿参数和 MoE 架构发展，计算效率至关重要。GramMuon 通过算法层面的重构和硬件层面的深度定制，证明了在保持优化质量的同时，可以显著降低训练成本。

其次，该工作展示了“算法-硬件协同设计”的价值。单纯优化算法系数或依赖通用数学库已触及天花板，而针对特定矩阵结构（如对称性）和特定 GPU 架构（如 Hopper/Blackwell）的定制内核，能够释放巨大的性能潜力。

最后，这一成果为后续研究提供了新的范式。它表明，对于涉及高度矩形矩阵和高阶预条件的优化算法，通过变换问题维度（从 $n \times m$ 到 $n \times n$）并利用数值线性代数中的结构特性，可以实现显著的加速。这对于 Kimi、GLM 等采用 Muon 优化器的前沿模型训练，以及未来更复杂的大模型架构，都具有重要的参考价值和实际应用意义。