Research map / Large language models / Reasoning and chain of thought

MuLoCo: Muon is a Practical Inner Optimizer for DiLoCo

ICML 2026 · Large language models › Reasoning and chain of thought · approach: Efficient LLM serving · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes MuLoCo, which replaces DiLoCo's AdamW inner optimizer with Muon, showing better worker scaling, compression robustness, and larger critical batch sizes for LLM pretraining up to 15B parameters.

Problem

DiLoCo's performance degrades as the number of workers grows and its critical batch size is limited; the choice of inner optimizer may be responsible.

Approach

Replace DiLoCo's AdamW inner optimizer with Muon (orthonormalized steps) so that worker pseudogradients are more directionally aligned and have less interference during averaging, giving better worker scaling, larger critical batch sizes, and robustness to compression.

Key results

• At 3.1B scale, MuLoCo's critical batch size is 2-4x larger than DiLoCo's at any worker count K.
• At 15B scale, K=1 MuLoCo matches the best baselines while using a 16M-token batch (4-16x larger), and K=16 MuLoCo matches final performance of DP AdamW on language model evaluation benchmarks.
• At iso-loss, K=1 MuLoCo reaches the loss of DP AdamW in roughly 4.5x fewer sequential FLOPs at observed scales.

Surprising findings

• K=1 DiLoCo improves more over its DP AdamW baseline than K=1 MuLoCo does over DP Muon, attributed to Muon's already well-aligned steps.
• Scaling exponents and the monotonic trend of improvement relative to data-parallel with model size were not consistently observed, contrary to prior DiLoCo scaling findings.

Limitations

• Theoretical analysis is acknowledged as speculative and not directly linked to performance.
• Scaling-law extrapolation to 15B is uncertain; 15B was not swept to optimality.
• Top-k sparsification overhead from index communication makes high sparsity less practical than 2-bit quantization.

Experimental setup

Datasets: Nemotron-CC

Benchmarks: held-out Nemotron-CC validation loss

Metrics: final evaluation loss percent increase over data-parallel baseline critical batch size cosine similarity of pseudogradients wall-clock training time scaling-law residuals

Baselines: DiLoCo DP AdamW DP Muon Streaming DiLoCo GPA

Builds on: DiLoCo Muon Streaming DiLoCo SNOO Gemma 3

Closest papers by meaning

• Towards Understanding the Power and Limits of the Muon Optimizer: A River-Valley Perspective (NeurIPS 2026 · similarity 0.89)
• Communication-Efficient Language Model Training Scales Reliably and Robustly: Scaling Laws for DiLoCo (NeurIPS 2025 · similarity 0.88)
• Sharp Capacity Scaling of Spectral Optimizers in Learning Associative Memory (NeurIPS 2026 · similarity 0.87)
• Powering Up Zeroth-Order Training via Subspace Gradient Orthogonalization (NeurIPS 2026 · similarity 0.87)
• From SGD to Muon: Adaptive Optimization via Schatten-p Norms (NeurIPS 2026 · similarity 0.87)
• AMUSE: Anytime Muon with Stable Gradient Evaluation (NeurIPS 2026 · similarity 0.86)