万亿模型训练的免费午餐一个数学 trick 让 Muon 提速 50%。在万亿参数大模型的竞逐中训练效率的细微差距往往关乎巨大的算力成本。近期Kimi K2 与 GLM-5 等前沿语言模型开始广泛采用 Muon 优化器。对比 AdamWMuon 达到特定损失值所需的优化器步数更少但单步计算开销显著增加。这种开销主要来自 Newton-Schulz 正交化过程引入了早期优化器中不存在的三次方时间复杂度矩阵运算。〓 Muon 与 AdamW 单步实际运行时间的对比为突破该算力瓶颈普林斯顿大学 Tri Dao 团队Mamba 与 FlashAttention 核心作者联合纽约大学研究人员提出了 Gram Newton-Schulz 算法。在万亿参数 MoE 模型训练中该算法将正交化步骤的端到端耗时有效降低了 40% 至 50%。目前团队已将该算法开源作为标准 Muon 优化器的即插即用替换模块无需繁琐配置即可直接获得加速收益。文章标题GREPO: A Benchmark for Graph Neural Networks on Repository-Level Bug Localization文章链接https://dao-lab.ai/blog/2026/gram-newton-schulz/项目链接https://github.com/Dao-AILab/gram-newton-schulz核心算子链接https://github.com/Dao-AILab/quack/blob/main/quack/gemm_symmetric.py为什么标准算法这么慢传统优化器如AdamW执行逐元素操作时间复杂度为。Muon 等现代优化器需要进行正交化单步计算需要耗费的时间假设。Muon 的更新规则基于对动量矩阵的极分解由于极分解精确计算成本高昂Muon 采用 Newton-Schulz 多项式迭代进行近似现代 Transformer 架构特别是包含大量细粒度专家的 MoE的权重矩阵形状大多是不规则的矩形满足。标准 Newton-Schulz 需要在庞大的矩形矩阵上执行多次乘法矩形矩阵乘法完全主导了整体计算成本。〓 包含大量昂贵矩形矩阵乘法的标准 Newton-Schulz 伪代码更关键的是算法执行期间产生的诸多中间矩阵具备对称结构常规计算路线未能有效利用这一数学特性导致半数计算工作冗余。〓 标准 Newton-Schulz 在 Hopper 架构下的纯算子优化收益Gram矩阵的数学重构算法的核心在于转移迭代空间不再对庞大的矩形输入矩阵进行迭代而是转移至尺寸更小且对称的方形 Gram 矩阵。极分解可表示为。通过多项式的代数变换Newton-Schulz 迭代隐含了计算逆平方根的过程。核心迭代可转化为以下标量形式矩阵操作保留了奇异向量该标量迭代逻辑可直接推广至矩阵空间。空间转换后算法主体在的对称矩阵内运行极大削减了浮点运算次数。在典型的场景下相较于未优化的标准实现该方法理论上可节省 68% 的浮点运算次数。解决数值不稳定上述理论在精确算术下完全等价但在真实的半精度计算环境下会引发严重的数值不稳定。研究团队通过算法与硬件的协同设计化解了这一工程隐患。重启策略在 bfloat16 精度下计算 Gram 矩阵会产生由于浮点误差导致的伪负特征值。由于更新规则包含平方项初始微小的负特征值会随迭代呈指数级放大最终导致数值崩溃。研究团队引入了重启策略在算法执行至中途时利用当前的近似结果重新构造 Gram 矩阵消除累积的负特征值并重置状态。〓 引入重启策略后的特征值演变与稳定收敛代数重排在计算矩阵二次型时常规方法会显式加上单位阵。底层算子在执行该加法时会先在 float32 下计算并由于输出限制向 float16 截断导致后续乘法累积严重的精度损失。团队重排了代数逻辑将加法操作隐式融入后续计算即先算再在后续步骤中分配的运算全程在 float32 下保持高精度消除了隐藏的数值隐患。精度回退决策针对 bfloat16 的动态范围大但精度位数不足的问题算法在初始化阶段默认将输入张量转换为 float16。由于矩阵范数已被严格控制在 1 附近float16 在这一小区间内能提供更高的尾数精度进一步夯实了数值基础。〓 稳定的 Gram Newton-Schulz 算法伪代码权重拆分策略工程实现中团队特别强调了将 SwiGLU 架构中的和拆分后独立进行正交化。由于这两部分对激活函数的梯度贡献机制不同拆分处理不仅使 Llama-430M 的验证集困惑度优化了约 0.2更通过减半矩阵的小维度使得依赖复杂度的 Gram Newton-Schulz 获得了更显著的提速比例。〓 不同形状权重矩阵的单步耗时拆解极端矩形权重加速最显著定制三角形调度器在底层算子层面为最大化 Gram 矩阵对称性带来的计算红利团队基于 CuTeDSL 针对 Hopper 和 Blackwell 架构开发了定制算子。〓 对称矩阵乘法的三角形调度器示意其核心是一个三角形调度器仅将矩阵下三角区域的工作块分配给计算集群并在底层内存回写时转置复制至上三角位置确保负载均衡并消除冗余的内存访存。〓 定制对称算子与 cuBLAS 在不同架构下的 TFLOPS 吞吐量对比实验验证与工程收益在 Llama-430M、Qwen-600M、Gemma-1B 以及 10 亿参数规模的 MoE-1B 模型上进行的对比实验表明使用 Gram Newton-Schulz 与原版 Muon 的验证集困惑度差异严格控制在 0.01 以内。〓 不同模型上的验证集困惑度对比Hopper 架构在实际训练耗时方面新算法配合定制算子切实缩短了正交化的端到端耗时。在模拟 Kimi K2 模型特定流水线阶段的真实并行负载测试中Gram Newton-Schulz 实现了 2 倍的端到端正交化加速。〓 Kimi K2 流水线切片下的端到端耗时对比结语Gram Newton-Schulz 通过底层的数学逻辑重构与针对性的数值稳定性修复以及 GPU 架构级别的定制算子优化打通了现代优化器在大规模并行训练中的效率瓶颈。这为极度消耗算力的矩阵正交化问题提供了一条可行路径也再次印证了算法与硬件协同设计的实用价值。目前研究团队已将 Gram Newton-Schulz 完整开源。在实际工程应用中唯一需要微调的超参数仅为重启迭代的节点。为此开源库中提供了一个自动化调参脚本只需输入一组多项式系数即可自动分析并建议最优的重启节点。这套兼具理论深度与工程可用性的工具为受限于算力瓶颈的大模型训练提供了一份切实可用的优化方案。更多阅读#投 稿 通 道#让你的文字被更多人看到如何才能让更多的优质内容以更短路径到达读者群体缩短读者寻找优质内容的成本呢答案就是你不认识的人。总有一些你不认识的人知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁促使不同背景、不同方向的学者和学术灵感相互碰撞迸发出更多的可能性。PaperWeekly 鼓励高校实验室或个人在我们的平台上分享各类优质内容可以是最新论文解读也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个让知识真正流动起来。稿件基本要求• 文章确系个人原创作品未曾在公开渠道发表如为其他平台已发表或待发表的文章请明确标注• 稿件建议以markdown格式撰写文中配图以附件形式发送要求图片清晰无版权问题• PaperWeekly 尊重原作者署名权并将为每篇被采纳的原创首发稿件提供业内具有竞争力稿酬具体依据文章阅读量和文章质量阶梯制结算投稿通道• 投稿邮箱hrpaperweekly.site• 来稿请备注即时联系方式微信以便我们在稿件选用的第一时间联系作者• 您也可以直接添加小编微信pwbot02快速投稿备注姓名-投稿△长按添加PaperWeekly小编现在在「知乎」也能找到我们了进入知乎首页搜索「PaperWeekly」点击「关注」订阅我们的专栏吧·