MuLoCo: Communication-Efficient LLM Training

• MuLoCo is a distributed training algorithm that integrates the Muon optimizer, error-feedback, and 2-bit quantization to compress updates and achieve an 8× reduction in communication bandwidth.
• It employs an orthogonalization-based momentum method to ensure that aggressive compression recovers lost information, maintaining model quality.
• Empirical evaluations demonstrate that MuLoCo outperforms standard DiLoCo with AdamW by reducing network overhead while preserving convergence rates and accuracy in large-scale transformer models.

MuLoCo (short for "Muon inner optimizer DiLoCo") is an algorithmic enhancement to distributed low-communication training frameworks for LLMs. It combines the Muon optimizer, orthogonalization-based momentum, and error-feedback-augmented communication compression within the DiLoCo system. MuLoCo enables communication-efficient pre-training of transformer models by allowing parameter update deltas to be compressed down to 2 bits per element with minimal performance degradation, thus achieving an 8× reduction in communication bandwidth relative to standard DiLoCo at constant memory complexity (Thérien et al., 29 May 2025).

1. Background and Motivation

Distributed Low-Communication Training (DiLoCo) is a Local SGD variant designed for efficient data center-scale training of LLMs. DiLoCo operates by executing $H$ local optimization steps on each of $K$ parallel workers before performing an all-reduce operation to aggregate parameter updates. This paradigm reduces the communication frequency by a factor of $H$, but each communication step still transmits a complete precision copy of the model—incurring significant network overhead.

Prior approaches have attempted to further compress communicated updates using quantization or sparsification, but their effectiveness is limited by the properties of the local optimizer (commonly AdamW), which tend to produce parameter deltas that are not readily compressible. The MuLoCo approach substitutes the Muon optimizer for AdamW as DiLoCo’s inner optimizer. This substitution, especially when paired with standard compression schemes and error-feedback accumulators, achieves much higher compression ratios with negligible loss in model quality (Thérien et al., 29 May 2025).

2. Formal Definitions and Algorithmic Components

Let:

• $\theta \in \mathbb{R}^d$: model parameters,
• $\Delta\theta \in \mathbb{R}^d$: parameter delta after $H$ inner steps per worker,
• $e \in \mathbb{R}^d$: error-feedback accumulator,
• $C(\cdot)$: compression operator (e.g., Top-$k$ sparsification or quantization).

The Muon inner optimizer update for each coordinate block uses a single momentum accumulator $m$: $$m_t = \beta_1 m_{t-1} + (1-\beta_1)\nabla\ell(\theta_t)$$

$$\Delta\theta_t = -\eta_{\text{in}}\;\mathrm{orth}(m_t)$$

where $\mathrm{orth}(m_t)$ denotes an orthonormalized version of $m_t$ (via, e.g., quintic Newton–Schulz iteration), $\beta_1=0.9$, and $\eta_{\text{in}}$ is the local learning rate.

Error feedback is implemented as follows: after $H$ local steps, each worker computes its local delta $\Delta\theta_i = \theta_i^{(0)} - \theta_i^{(H)}$, updates the accumulator $e_i \leftarrow e_i + \Delta\theta_i$, compresses the sum $Q_i \leftarrow C(e_i)$, sends $Q_i$ instead of $\Delta\theta_i$, and then sets $e_i\leftarrow e_i-Q_i$. This mechanism ensures recovery of the information lost to aggressive compression in subsequent rounds.

3. MuLoCo Workflow

The MuLoCo workflow extends DiLoCo by introducing Muon-based inner optimization, error-feedback, and compression before the global synchronization step. The high-level steps per outer iteration are:

1. Each worker initializes $\theta_i$ from the global $\theta$.
2. Run $H$ inner Muon steps locally, updating $m$, orthogonalizing, and advancing $\theta_i$.
3. Compute local delta: $$\Delta\theta_i = \theta_{\text{before}} - \theta_{\text{after}}$$.
4. Error feedback: $e_i \leftarrow e_i + \Delta\theta_i$, $Q_i \leftarrow C(e_i)$, $e_i \leftarrow e_i - Q_i$.
5. All-reduce aggregated compressed deltas: $Q = (1/K) \sum_i Q_i$.
6. Outer update: SGD with Nesterov momentum using the aggregated $Q$.
7. Maintain or reset local states as per optimizer protocol.

4. Communication and Memory Analysis

MuLoCo’s principal advantage is the aggressive reduction of communication volume:

Variant	Bits Sent per Communication	Memory Complexity
DiLoCo (AdamW, no compression)	$32 \times |\theta|$	$2M$ (AdamW) $+M$ outer $=3M$
DiLoCo+EF (adds error accumulator)	$32 \times |\theta|$	$4M$
MuLoCo (Muon, no error feedback)	$32 \times |\theta|$	$M$ (Muon) $+M$ outer $=2M$
MuLoCo+EF (2-bit quant. + error feed)	$2 \times |\theta|$	$3M$

Here $M = |\theta|$. With 2-bit quantization, MuLoCo's communication cost is $N\times2|\theta|$ bits ($N$ = number of communications), representing an 8× reduction relative to 16-bit AdamW-DiLoCo for the same $H$. The memory cost is also lower than AdamW-based DiLoCo due to Muon's single momentum buffer (Thérien et al., 29 May 2025).

5. Convergence and Theoretical Guarantees

MuLoCo inherits the convergence guarantees of Local SGD with error-feedback. Under standard smoothness and bounded variance assumptions, and with unbiased compression $C(\cdot)$, the convergence rate is $O(1/\sqrt{KT})$ where $K$ is the number of workers and $T$ is the number of total steps. This rate includes additional constants accounting for compression error but is fundamentally unchanged in qualitative terms by the introduction of Muon or aggressive compression (Thérien et al., 29 May 2025). No new theorems are introduced for MuLoCo, but the theoretical basis references established results (e.g., Karimireddy et al., 2019).

6. Empirical Results

MuLoCo’s empirical evaluation is conducted using a 12-layer decoder-only transformer (post-layernorm), 220M parameters, on FineWeb-EDU using $K=8$ workers, $H=30$ local steps, and a total batch of $2^{18}$ tokens. Baselines include data-parallel AdamW/Muon and DiLoCo (AdamW), with or without compression/error-feedback.

Key findings:

• Without compression, MuLoCo slightly outperforms DiLoCo and matches Muon in the data-parallel regime.
• With error feedback, both Top-$k$ sparsification ($k$ in $[1\%, 10\%]$) and 2-bit quantization further stabilize training.
• MuLoCo combined with error feedback and 2-bit quantization achieves a lower test loss than standard AdamW-DiLoCo (16-bit), while reducing total bits sent by a factor of 8.
• All results are reported at wall-clock-matched workloads, confirming the practical benefits of aggressive compression enabled by Muon (Thérien et al., 29 May 2025).

7. Implementation Recommendations

The MuLoCo approach is compatible with modern accelerator-based pipelines (JAX/Optax, PyTorch). Recommended practices include:

• Applying Muon as the inner optimizer to hidden layers, with optional retention of AdamW for embeddings/output layers.
• Enabling error feedback with $\beta\approx0.9$ for MuLoCo ($\beta\approx0.7$ for AdamW-DiLoCo).
• For maximal bandwidth reduction, employ 2-bit quantization. Alternatively, Top-$k$ sparsification with $k$ in $[1\%, 10\%]$ is effective.
• Hyperparameters: for MuLoCo, $\eta_{\text{in}}\approx2.68\times10^{-3}$, min/max LR ratio $=0.1$; AdamW-DiLoCo $\eta_{\text{in}}\approx3\times10^{-4}$. Outer SGD uses $\eta_{\text{out}}=0.8$, Nesterov $\mu=0.9$.
• Pipeline: wrap DiLoCo’s inner loop with Muon updates, apply compression and error-feedback prior to all-reduce, then carry out the outer SGD step (Thérien et al., 29 May 2025).

This structure supports practitioners seeking communication-efficient, large-scale LLM training with improved compressibility and memory efficiency relative to prior distributed frameworks.