1. 首頁 > 生活百科 > >

不加算力，只改一個算法：Muon在萬億MoE模型中最高2倍加速

人工智能ai

不加算力，只改一個算法：Muon在萬億MoE模型中最高2倍加速

文章圖片

不加算力，只改一個算法：Muon在萬億MoE模型中最高2倍加速

文章圖片

不加算力，只改一個算法：Muon在萬億MoE模型中最高2倍加速

文章圖片

不加算力，只改一個算法：Muon在萬億MoE模型中最高2倍加速

文章圖片

不加算力，只改一個算法：Muon在萬億MoE模型中最高2倍加速

文章圖片

不加算力，只改一個算法：Muon在萬億MoE模型中最高2倍加速

文章圖片

不加算力，只改一個算法：Muon在萬億MoE模型中最高2倍加速

文章圖片

不加算力，只改一個算法：Muon在萬億MoE模型中最高2倍加速

文章圖片

不加算力，只改一個算法：Muon在萬億MoE模型中最高2倍加速

文章圖片

不加算力，只改一個算法：Muon在萬億MoE模型中最高2倍加速

文章圖片

機器之心編輯部
在數值分析領域 ， Newton-Schulz 及其相關方法已被研究多年 ， 但大多數工作關注的是高精度計算、CPU 優化或方陣輸入 。
就在昨天 ， 普林斯頓大學、紐約大學的四位研究者提出 Gram Newton-Schulz ， 通過重構 Newton-Schulz ， 使其更適配 GPU 和大模型訓練場景 ， 在萬億參數 MoE 模型中可將優化器時間降低 40–50% 。

我們用一句話來總結 Gram Newton-Schulz 的核心思想：

不再直接在矩陣 X∈R^n×m 上迭代 ， 而是在更小的 Gram 矩陣 X X^?∈R^n×n 上迭代 ， 從而降低計算量并利用對稱矩陣計算優化 。
核心貢獻包括如下：
  • 將標準 Newton-Schulz 重寫為數學等價形式 ， 主要作用在 n×n 空間；
  • 提出 Naive Gram Newton-Schulz（基礎版本）；
  • 分析半精度下的不穩定性 ， 并提出重啟策略（Stabilized 版本）；
  • 實現對稱矩陣乘法 GPU kernel；
  • 構建 GramMuon ， 在實際訓練中實現 40–50% 加速且精度無損 。
下圖 1 為基于 NVIDIA B300 平臺 ， AdamW 與 Muon 在 LLaMA 不同模型規模下的優化器步驟耗時（實際墻鐘時間）對比 。

作者之一、普林斯頓大學本科生 Jack Zhang 表示 ， 「我們讓 Muon 在幾乎不增加額外成本的情況下 ， 運行速度最高提升 2 倍！Gram Newton-Schulz 可以作為 Muon 中 Newton-Schulz 的即插即用替代方案：我們觀察到驗證集困惑度幾乎沒有變化 ， 誤差在 0.01 以內 。 同時 ， 我們也分享了在這一算法穩定化過程中的大量探索 ， 確保在任何情況下都不犧牲訓練質量 。 」

作者之一、普林斯頓大學助理教授、Together AI 聯合創始人兼首席科學家 Tri Dao 表示 ， 「這是我最喜歡的一類工作 —— 線性代數洞察 + 高性能算子 。 我們花了幾個月時間 ， 深入研究 Muon 中間過程里這些矩陣的特征值和特征向量 ， 最終提出了一種簡單而優雅的算法 ， 讓這一思路真正落地 。 」

接下來我們來看 Gram Newton-Schulz 相較于標準 Newton-Schulz 方法的優勢 。 在此之前 ， 有必要先回顧一下 Muon 優化器 。
Muon 可以理解為在譜范數下的最速下降法：

其中：μ 表示動量系數 ， η 表示學習率 ， polar 表示極分解 。
極分解定義：若 X=UΣV^? ， 則：

由于精確計算昂貴 ， Muon 使用 Newton-Schulz 進行近似 。
標準 Newton-Schulz 的迭代形式如下：

本質上是對奇異值進行逐步歸一化 ， 使其趨向 1 。 不過 ， 標準 Newton-Schulz 存在計算瓶頸 。 每輪迭代包含三次矩陣乘法：
  • XX^?：2mn^2
  • A^2：2n^3
  • BX：2mn^2
總 FLOPs 如下：

這就導致了一些問題 ， 包括大量矩形矩陣乘法（GPU 效率低） 、未利用對稱性 ， 并成為優化器瓶頸 。
因此 ， 本文 Gram Newton-Schulz 核心思想是：將迭代從 X 轉移到 Gram 矩陣：

關鍵變換：任何奇數多項式

可寫為

這樣一來 ， 就將問題轉為了 1）對 R 進行矩陣多項式迭代、 本質是近似 Y^-1/2 。 優勢在于：維度從 n × m → n × n、可用對稱矩陣乘法 kernel 以及 減少矩形 GEMM 次數 。
不過 ， Naive Gram Newton-Schulz 在 float16 下不穩定 ， Gram 矩陣可能產生負特征值 ， 也會出現 loss spike 和 Inf 。

下圖為 Llama-430M 上的 Naive Gram Newton-Schulz：

因此有了穩定版的 Stabilized Gram Newton-Schulz：

核心改進在于使用 float16 加速：

并在第 2 步后重啟：

以及重新初始化 Gram 。
【不加算力，只改一個算法：Muon在萬億MoE模型中最高2倍加速】結果顯示 ， 在保持穩定性的同時 ， 仍然比原算法更快 。
標準 Newton-Schulz、Naive Gram 與 Stabilized Gram 的復雜度對比如下：

結果表明 ， 當 α>1（常見情況） ， Gram 方法明顯更便宜 ， FLOPs 降低最高約 42%–58% 。
最后來看下 Gram Newton-Schulz 在 Kimi K2 中的耗時：
Kimi K2 是一個萬億參數級的稀疏、細粒度 MoE 模型 ， 每一層包含 384 個專家 ， 隱藏層維度為 7168 ， 專家的中間層維度為 2048 。 由于模型正朝著更細粒度的 MoE 架構發展 ， 且 Kimi K2 使用 Muon 進行訓練 ， 因此這是一個評測 Gram Newton-Schulz 的理想場景 。
團隊將 Kimi K2 一次全局訓練步驟中暴露出來的 Newton-Schulz 的墻鐘時間 ， 近似為以下部分的總和：
  • 216 個專家的 up/gate/down 權重 ， 形狀為 2048 × 7168；
  • 1 個 dense 的 up/gate/down 權重 ， 形狀為 7168 × 18432 。
下圖為 NVIDIA H100 Hopper 硬件上 ， 在 Kimi K2 的流水線并行配置中 ， Gram Newton-Schulz 的速度是標準 Newton-Schulz 的 2 倍 。

下圖為 NVIDIA B300 Blackwell 硬件上 ， 在 Kimi K2 的流水線并行配置中 ， Gram Newton-Schulz 的速度同樣是標準 Newton-Schulz 的 2 倍 。

更多細節內容請參閱原博客 。
博客地址：https://dao-lab.ai/blog/2026/gram-newton-schulz/

    推薦閱讀