MuonBP: Faster Muon via Block-Periodic Orthogonalization

TL;DR

MuonBP addresses the throughput gap of Muon under model parallelism by introducing block-periodic orthogonalization, where local blockwise updates run on each device and periodic full orthogonalization maintains stability. It formalizes two learning rates, $\eta_{full}$ and $\eta_{block}$, and a period $P$ to interpolate between BlockMuon and Muon, supported by convergence guarantees involving a harmonic-mean smoothness $\bar{L}_{BP}$. Empirically, MuonBP recovers Muon’s data-efficiency while achieving up to an 8% increase in throughput on 8B-scale models and reducing optimizer-communication by a factor of 5 in periodic steps, with BlockMuon alone proving unstable at scale. Overall, MuonBP offers a simple, practical path to faster large-scale pretraining without sacrificing convergence or accuracy.

Abstract

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.

Figures (12)

• Figure 1: Validation loss as a function of orthogonalization period for different TP degrees (280M model).
• Figure 2: Parameter norm vs iteration of competing methods.
• Figure 3: 8B model validation perplexities. Comparison of Muon, BlockMuon, and MuonBP across wall-clock time. For a target validation perplexity our method is $\sim10-13\%$ faster in terms of the wall-clock time to reach it, and for a given time point before the learning rate decay our method results in $\sim5-7\%$ lower perplexity compared to the baseline.
• Figure 4: 960M model. Comparison of baseline, block, and periodic orthogonal block methods across training steps and wall-clock time.
• Figure 5: 1.2B model. Comparison of baseline, block, and periodic orthogonal block methods across training steps and wall-clock time.

Theorems & Definitions (17)

• Theorem 1
• Lemma 1: Dual of the Block-Spectral Norm
• Theorem 2: Convergence of MuonBP
• Lemma 2: Dual of the block-spectral norm
• proof
• Lemma 3
• proof
• Lemma 4
• proof
• Lemma 5: Descent Lemma