arXiv:2510.16981cs.LGmath.OC2025-10被引 33

提出块周期正交化优化器,提速同时保持训练稳定性。

MuonBP: Faster Muon via Block-Periodic Orthogonalization

  • 分块独立正交化+周期性全局正交,降低通信开销。
  • 80亿参数模型下比原版Muon提升8%吞吐量,性能无损。
  • 适合大规模分布式语言模型训练,兼容ZeRO等并行策略。

梯度正交化能显著加速梯度下降。Muon优化器(Jordan, Jin等,2024)结合一阶动量与梯度正交化,在语言模型训练中相比Adam/AdamW大幅提升数据效率。然而在模型并行时,梯度正交化需额外的梯度分片收集与分发操作,相较坐标式优化器(如AdamW)引入5%-10%吞吐损失。为此,我们提出带有块周期正交化的MuonBP(MuonBP),在各设备上独立对矩阵分片执行正交化,并周期性进行全量正交以保障大规模训练稳定性。我们推导了从基线到MuonBP的学习率调整方法,并给出该算法的收敛性保证。关键的是,理论指出需使用两个步长:一个用于分块正交化步骤,一个用于全量正交化步骤。该方法简单,仅需极少超参数调整,迭代复杂度接近基线Muon,且每轮吞吐量媲美坐标式方法(如AdamW)。在八路张量并行和ZeRO优化器状态分片下训练80亿参数模型时,MuonBP相比原版Muon实现8%吞吐提升,性能无下降。

原文摘要 · Abstract (English)

Gradient orthogonalization is a simple strategy that shows great utility in speeding up gradient descent. The Muon optimizer (Jordan, Jin, et al., 2024) combines gradient orthogonalization with first-order momentum and achieves significant improvement in data efficiency over Adam/AdamW (Loshchilov and Hutter, 2019) for language model training. However, when using model parallelism, gradient orthogonalization introduces additional overhead compared to coordinate-wise optimizers (such as AdamW) due to additional gather and scatter operations on gradient matrix shards from different devices. This additional communication can amount to a throughput hit of 5%-10% compared to Adam/AdamW. To remedy this, we propose Muon with Block-Periodic Orthogonalization (MuonBP), which applies orthogonalization independently to matrix shards on each device and periodically performs full orthogonalization to maintain training stability at scale. We show how to adjust the learning rate from the baseline to MuonBP and give convergence guarantees for this algorithm. Crucially, our theory dictates that we use two stepsizes: one for the blockwise orthogonalization steps, and one for the full orthogonalization steps. Our method is simple, requires minimal hyperparameter adjustments, and achieves competitive iteration complexity compared with baseline Muon while providing per-iteration throughput comparable to coordinate-wise methods such as AdamW. When training an 8B model with eight-way tensor parallelism and ZeRO optimizer state sharding, MuonBP achieves 8% throughput increase compared to Muon with no degradation in performance.

优化器模型并行训练加速

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。