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

AIGC动态2小时前发布 almosthuman2014

53 0 0

文章摘要

【关键词】 GPU优化、Gram矩阵、大模型训练、Muon加速、算法稳定

在针对大模型训练中优化器效率瓶颈的探索中，普林斯顿大学与纽约大学研究团队提出 Gram Newton-Schulz 方法，旨在提升高维稀疏模型训练中关键矩阵运算的计算效率。其核心思想在于将传统迭代操作从原始矩阵 $X \in \mathbb{R}^{n \times m}$ 转移至更小维度的 Gram 矩阵 $XX^\top \in \mathbb{R}^{n \times n}$ 上执行，以降低整体计算量并更好地适配 GPU 对称矩阵算子的硬件特性。该方法不仅避免了标准 Newton-Schulz 中多次低效矩形矩阵乘法（GEMM），还通过重构迭代表达式将复杂度由 $O(mn^2 + n^3)$ 降为依赖对称结构的更高效率形式。

此重构显著减少了 FLOPs，理论最高可降低约 42%–58%，实测在 Kimi K2 这类万亿参数级 MoE 模型中实现 2 倍速度提升，优化器步骤时间减少 40–50%。团队进一步提出 Naive 版本并分析其在半精度下的不稳定性问题——包括 Gram 矩阵负特征值、loss spike 及 Inf 值现象。基于上述发现，作者设计 Stabilized Gram Newton-Schulz：引入 float16 加速计算机制，在第 2 步后加入重启策略，同时重置 Gram 矩阵初始状态以维持数值稳定性。实验验证该改进方案既保障了训练质量无损（验证集困惑度变化低于 0.01），又持续保持性能优势。

文章详细阐述了该方法在实际工程中的落地路径，包含高性能对称矩阵乘法的 CUDA kernel 实现，并构建名为 GramMuon 的模块。测试表明，在 NVIDIA B300 和 H100 平台下，Gram Newton-Schulz 在 LLaMA 不同规模模型及 Kimi K2 模型流水线并行配置中均表现优异；特别是在 Kimi K2 针对细粒度专家架构的训练场景中，因模型参数量巨大且涉及大量中间层张量变换，传统 Newton-Schulz 成为明显性能瓶颈，而 Gram 变体有效缓解了这一限制。

研究者特别强调，该工作融合线性代数深度理解与底层算子工程实践，体现了“理论—工程”协同创新思路：通过精准识别中间过程矩阵的特征值分布和奇异谱特性，进而设计出兼具简洁性与实用性的近似算法。团队公开分享相关经验与代码，支持其作为即插即用组件替换 Muon 内部现有 Newton-Schulz 的可行性，为未来超大规模模型训练提供切实可行的技术路线。