---
title: 'MuLoCo: Communication-Efficient LLM Training'
url: https://www.emergentmind.com/topics/muloco-algorithm
type: topic
---

# MuLoCo: Communication-Efficient LLM Training

MuLoCo (short for "Muon inner optimizer DiLoCo") is an algorithmic enhancement to distributed low-communication training frameworks for large language models. 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 [2505.23725].

## 1. Background and Motivation

Distributed Low-Communication Training (DiLoCo) is a Local SGD variant designed for efficient data center-scale training of large language models (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 [2505.23725].

## 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 [2505.23725].

## 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 [2505.23725]. 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 [2505.23725].

## 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 [2505.23725].

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

Source: https://www.emergentmind.com/topics/muloco-algorithm