🤖 machine learning

DMuon: Efficient Distributed Muon Training with Near-Adam Overhead

The paper introduces DMuon, an open-source distributed implementation of the Muon optimizer that integrates seamlessly into existing training pipelines to drastically reduce optimizer-step latency and achieve near-AdamW efficiency, thereby enabling the scalable use of matrix-orthogonalization-based optimization in large foundation models.

Original authors: Vincent Chen, Starrick Liu, Regis Cheng, Dance Yang, Shalfun Li, Ryan Yu, Lucy Liang, Hang Su, Roy Gan, Hao Wang, Qian Wang

Published 2026-06-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Vincent Chen, Starrick Liu, Regis Cheng, Dance Yang, Shalfun Li, Ryan Yu, Lucy Liang, Hang Su, Roy Gan, Hao Wang, Qian Wang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are training a giant robot brain (a Large Language Model or an embodied AI) to learn. To teach it, you need an "optimizer"—a coach that looks at the robot's mistakes and adjusts its internal connections (weights) to make it smarter.

For years, the standard coach has been AdamW. It's like a team of thousands of tiny workers, each fixing one tiny screw on the robot at a time. It's fast and efficient, but it treats every screw independently.

Recently, a new, smarter coach called Muon was discovered. Instead of fixing screws one by one, Muon looks at entire panels of the robot (whole matrices of data) and adjusts them all together. This "group therapy" approach helps the robot learn faster and reach higher performance. However, there's a catch: Muon is incredibly slow to run on a distributed system (a cluster of many computers working together).

The Problem: The "All-Hands" Meeting Bottleneck
In a distributed training setup, the robot's brain is split across many computers (GPUs).

AdamW works like a remote team: Computer A fixes its screws, Computer B fixes theirs, and they don't need to talk to each other much.
Muon works like a committee: To fix a panel, every computer needs to see the entire panel first.

In a naive implementation of Muon, every single computer has to stop what it's doing, download the full data from everyone else, perform the complex math, and then send the results back. This is like holding a massive meeting where everyone reads the entire 1,000-page report before making a single decision. It takes so long that the "meeting" (the optimizer step) takes longer than the actual work (the learning steps). In fact, the paper notes this can cost 2x more time than the learning itself!

The Solution: DMuon (Distributed Muon)
The X Square Robot Team built DMuon, a new system that makes Muon run almost as fast as AdamW. They did this by changing how the team works, not by changing the math of the coach.

Here is how they solved the bottleneck using three main tricks:

1. The "Designated Expert" Strategy (Owner-Centric Execution)

Instead of everyone trying to fix every panel, DMuon assigns one specific computer to be the "Owner" for each panel.

Before: Everyone gathered the data, everyone did the math, everyone wasted time.
Now: Only the "Owner" computer gathers the data and does the heavy math. The other computers just wait or do other work.
Analogy: Imagine a construction site. Instead of every worker trying to build the whole house, one expert is assigned to the roof, another to the plumbing. The others don't try to build the roof; they just hand the expert the tools they need. This stops everyone from doing the same work twice.

2. The "Conveyor Belt" System (Communication Overlap)

While the "Owner" is doing the math, the other computers aren't just sitting idle. DMuon creates a pipeline:

Forward Pass: While the robot is "thinking" (processing data), the system quietly sends the next set of tools to the experts in the background.
Backward Pass: While the robot is "learning from mistakes," the system is already gathering the next batch of data for the experts.
Analogy: It's like a restaurant kitchen. The chef (the Owner) is chopping vegetables. Instead of waiting for the chef to finish before the waiter brings the next order, the waiter brings the next order while the chef is still chopping. The waiting time is hidden inside the chopping time.

3. The "Smart Assembly Line" (Batching and Tuning)

The math Muon does is complex. Sometimes the pieces are huge (like a giant wall), and sometimes they are tiny (like a small tile).

The Problem: If you try to process a tiny tile on a giant machine, the machine sits idle. If you process a giant wall, it takes forever.
The Fix: DMuon groups many tiny tiles together into a single "batch" so the machine stays busy. It also uses a "tuner" that automatically finds the fastest way to chop each specific shape of tile.
Analogy: Instead of sending a delivery truck for a single envelope, DMuon waits until it has a full truckload of envelopes and sends them all at once. It also changes the truck route depending on whether the destination is a city or a farm.

The Results

The paper tested DMuon on real-world models, including a robot training model (Wall-OSS) and a language model (Qwen2.5).

Speed: DMuon made the "optimizer step" (the meeting) 6 to 163 times faster than the old, naive way.
Efficiency: The total time to train a step is now only 2% slower than the standard AdamW method.
Conclusion: DMuon allows teams to use the superior "group therapy" learning style of Muon without paying the massive time penalty. It turns a slow, clunky process into a smooth, efficient one, making it ready for production use in massive AI models.

In short, DMuon is the traffic control system that lets the smartest coach (Muon) run on a highway of computers without causing a traffic jam, finally making it practical to train the next generation of AI.

Technical Summary: DMuon – Efficient Distributed Muon Training with Near-Adam Overhead

1. Problem Statement

Matrix-aware optimizers, particularly Muon, have demonstrated superior convergence properties compared to element-wise optimizers like AdamW by applying Newton–Schulz (NS) iterations to the momentum-aggregated gradients of weight matrices. This approach drives the singular values of updates toward a uniform spectrum, offering significant compute efficiency gains in pretraining large models (e.g., LLMs and embodied foundation models).

However, deploying Muon in modern distributed training environments faces a critical systems bottleneck. Contemporary distributed training stacks (e.g., ZeRO, FSDP) operate under an implicit contract of element-wise optimization, where updates are computed independently on parameter shards. Muon violates this contract because its core NS iteration requires the full weight matrix to be present on a single device to compute the necessary matrix products (specifically $XX^\top$ ).

In a naive distributed implementation ("Gather-then-Compute"), every rank must:

All-gather the full reduced gradient matrix from all shards.
Perform the NS iteration locally.
Discard the results for shards it does not own.

This approach incurs two major costs:

Communication Overhead: Reconstructing full matrices at every step introduces collective communication that scales with model size and distributed width.
Redundant Computation: Every rank performs the same expensive orthogonalization, multiplying the computational cost by the number of ranks ( $D$ ).

In practice, these costs cause the optimizer step time in distributed Muon to rival or exceed the combined wall-clock time of the forward and backward passes, eroding the algorithmic advantages of Muon. This penalty is particularly severe for embodied model training (e.g., VLA), where forward-backward passes are computationally lighter relative to the optimizer step compared to standard LLM pretraining.

2. Methodology: DMuon System Design

DMuon is an open-source distributed implementation designed to close the performance gap between Muon and AdamW. It achieves this through a co-designed system architecture that integrates communication optimization, load balancing, and kernel-level efficiency, without modifying the underlying training framework (e.g., PyTorch/FSDP2).

2.1 Owner-Centric Execution Strategy

Instead of replicating computation across all ranks, DMuon assigns each matrix parameter to a single owner rank.

Owner Role: The owner maintains the authoritative optimizer state and computes the Muon update (Newton–Schulz iteration) exactly once per matrix.
Non-Owner Role: Other ranks materialize the parameter only when required for forward/backward computation and do not perform the orthogonalization.
Communication Flow:
- Forward: Owners publish updated parameters to non-owners asynchronously, pipelined with layer execution.
- Backward: Gradients are accumulated in packed buffers and reduced to the owner.
- Overlap: DMuon employs a two-stage hierarchy (intra-node and inter-node) with fine-grained weight layouts to overlap communication with computation, preventing collective contention.

2.2 Efficient Gram Newton–Schulz Iteration

To reduce the computational cost of the owner-side update, DMuon adopts the Gram Newton–Schulz formulation:

Gram Space Recurrence: Instead of operating in the original $m \times n$ space, the iteration is performed in the $m \times m$ Gram space ( $G = XX^\top$ ). This reduces complexity from $O(m^2n)$ to $O(m^3)$ when $m < n$ .
Symmetry Exploitation: Since $G$ is symmetric, DMuon uses a SYRK-style execution path to compute only the lower triangular portion, halving the arithmetic work of the dominant matrix multiplication.
Batching: To address GPU occupancy issues with small matrices, DMuon batches matrices of similar shapes and executes the Gram NS recurrence jointly, amortizing kernel dispatch overhead.

2.3 Computation-Aware Load Balancing

Because Muon update costs vary significantly with matrix shape, a naive assignment (e.g., round-robin) leads to straggler ranks. DMuon formulates owner assignment as a Mixed-Integer Linear Program (MILP):

Profiling: During initialization, the system benchmarks the execution time of various batch sizes and kernel schedules for each distinct matrix shape.
Optimization: The MILP solver minimizes the predicted makespan (maximum load) across owner ranks, ensuring the heaviest matrices are distributed to balance the workload.
Fallback: For extremely large search spaces, a greedy heuristic is used to bound initialization time.

2.4 Shape-Adaptive Kernel Autotuning

DMuon utilizes a DSL-driven autotuning workflow (using TileLang and CUTE DSL) to generate high-performance kernels.

It searches over tile shapes, block sizes, and software pipeline configurations specific to the target hardware and matrix shapes.
The optimal configuration is cached and reused for subsequent steps, ensuring that the kernel schedule is always near-optimal for the specific problem dimensions encountered in transformer training.

3. Key Contributions

Fine-Grained Communication Optimization: An owner-centric runtime that minimizes communication overhead via hierarchical scheduling and fine-grained weight layouts, preserving exact optimizer semantics.
Shape-Adaptive Execution Stack: A high-performance engine combining batched Gram Newton–Schulz, symmetry-aware kernels, and autotuning to handle heterogeneous matrix workloads efficiently.
Computation-Aware Load Balancing: A measured-cost optimization approach (MILP) that assigns matrix owners based on actual runtime profiling rather than theoretical FLOP counts.
Drop-in Integration: A module that integrates into existing training pipelines (e.g., FSDP2) with minimal effort, requiring no framework-level modifications.
Production-Scale Validation: Successful deployment on both robotics (embodied) and language model training workloads.

4. Experimental Results

The authors evaluated DMuon on four production workloads: Wall-OSS-0.5, Pi0, Wall-WM, and Qwen2.5-7B, scaling from 8 to 256 A800 GPUs.

Optimizer Step Time: DMuon achieves a 6.85× to 163.00× speedup in optimizer-step time compared to a naive "Gather-then-Compute" distributed Muon implementation.
End-to-End Step Time: DMuon achieves a 1.48× to 3.01× speedup in total step time.
Overhead vs. AdamW: The per-step latency of DMuon is brought to near-AdamW levels, with an average overhead of only +2% compared to AdamW across the evaluated workloads.
Scaling: On Wall-OSS, DMuon maintains throughput close to AdamW as the number of GPUs increases, whereas naive Muon would suffer from severe communication bottlenecks.

Speedup Breakdown:

Symmetric Gram Kernel: 48% of the speedup (halving arithmetic in dominant products).
Owner Scheduling & Load Balancing: 32% (removing $D$ -fold redundant computation).
Auto-tuning & Batching: 16% (reducing kernel launch overhead for small matrices).

5. Significance and Claims

The paper claims that DMuon successfully resolves the systems mismatch between matrix-aware optimizers and distributed training infrastructure. By eliminating redundant computation and optimizing communication, DMuon allows the convergence benefits of Muon (such as higher compute efficiency and better optimization geometry) to translate into practical wall-clock training speedups.

The authors position DMuon not as a new algorithm, but as a necessary systems enabler that makes Muon a viable, drop-in replacement for AdamW in large-scale production training. The work demonstrates that the "penalty" of matrix-aware optimization can be reduced to negligible levels, enabling efficient scaling for both Large Language Models and Embodied Foundation Models (VLA).

The paper explicitly notes that DMuon does not change the update rule or convergence properties of Muon; it strictly addresses the systems overhead. Consequently, the benefits are most pronounced in distributed settings, while single-GPU training sees a more modest improvement (approx. 2× optimizer speedup) due to the lack of distributed redundancy to eliminate.