IT资源栈FlashAttention 联合作者 Tri Dao 发布了一种名为 Gram Newton-Schulz 的全新算法，旨在显著降低大语言模型训练的计算成本。该研究针对 Muon 优化器在 Kimi、GLM 等万亿参数模型训练中遇到的计算瓶颈进行了深度优化。Muon 虽通过正交化提升了优化质量，但其核心的 Newton-Schulz 迭代过程涉及大量昂贵的矩阵乘法。新方法通过迭代处理小型的对称 Gram 矩阵，大幅减少了浮点运算量，并针对 Hopper 和 Blackwell GPU 架构开发了专门的对称矩阵乘法内核。同时，团队引入“重启”策略解决了半精度运算中的数值不稳定性问题。实验结果显示，该技术在万亿参数 MoE 模型上能将优化器步长缩短 40% 至 50%，且不影响模型验证困惑度。
事件分析
随着大模型参数规模向万亿级迈进，优化器的计算开销已不容忽视，Muon 凭借卓越的收敛性逐渐成为 Kimi 等前沿模型的选择，但其高昂的计算成本限制了推广。Tri Dao 的这项工作通过数学变换利用矩阵对称性，并定制底层 GPU 内核，在不牺牲精度的前提下实现了显著的“免费午餐”式性能提升。这标志着大模型训练的效率优化正从通用的算力堆叠转向算法与硬件架构的深度协同设计，特别是针对 MoE 架构的高 rectangular 矩阵特性优化，有望成为未来降低大模型训练成本的关键技术路径。
💡 核心观点：算法与底层硬件的深度协同优化正成为大模型降本增效的突破口，Muon 优化器的性能瓶颈已被打破。
原文链接：Hacker News
评论前必须登录！
立即登录 注册