Surprising Discovery: A Chinese embodied‑AI company solves the distributed Muon bottleneck

Hello, I am PaperAgent, not an Agent! This year Chinese models DeepSeek‑V4 and Kimi‑K2 have shown strong performance, and both papers reveal that they both adopt the Muon optimizer instead of AdamW.

Muon, however, introduces a fatal problem: distributed training adds extra computation and communication, making the end‑to‑end step time 2.2 times slower than AdamW.

The recent DMuon paper demonstrates a solution that compresses Muon's distributed overhead to almost the same level as AdamW, accelerating embodied‑model training by 30%.

Surprisingly, the DMuon work comes from a domestic embodied‑model company, Zibian Robot.

Title: DMuon: Efficient Distributed Muon Training with Near‑Adam Overhead
Team: Zibian Robot, X‑Square‑Robot
Link: https://arxiv.org/pdf/2606.27153
Code: https://github.com/X-Square-Robot/dmuon

01 A headache‑inducing problem

Muon is a matrix‑aware optimizer that treats each weight matrix as a whole and uses Newton‑Schulz iteration to equalize singular‑value distribution, unlike element‑wise AdamW.

Both Kimi‑K2 and DeepSeek‑V4 use Muon, and Moonlight reports that Muon can be roughly twice as compute‑efficient as AdamW.

Nevertheless, Muon cannot scale in distributed environments because its Newton‑Schulz step requires the full matrix on each GPU, forcing every GPU to gather, compute, and then scatter the matrix, which dramatically inflates both communication and computation.

The paper cites a shocking figure: without optimization, Muon's overhead is twice the cost of forward + backward passes.

02 DMuon’s solution

Zibian Robot’s DMuon system keeps Muon's mathematical semantics while reducing distributed overhead to near‑AdamW levels.

It achieves this through three core designs:

1. Owner‑style execution – each matrix is assigned a single “owner” GPU that performs the Newton‑Schulz iteration; other GPUs simply synchronize the result, eliminating redundant computation.

2. Computation‑aware load balancing – DMuon first benchmarks every matrix shape and candidate batch size, then feeds the measured runtimes to a MILP solver that minimizes the slowest rank’s completion time; if the problem is too large, it falls back to a greedy search.

3. Overlapped communication and computation – DMuon pipelines parameter broadcast and gradient reduction so that forward and backward stages overlap with network traffic, preventing idle network time.

Forward‑side overlap: inter‑parameter broadcast and intra‑parameter broadcast are staged to overlap with main compute.

Backward‑side overlap: gradient reduction interleaves with next‑layer parameter broadcast, avoiding congestion.

4. Space‑efficient Newton‑Schulz – DMuon uses a Gram‑based variant that reduces the dominant matrix‑multiply from O(m²n) to O(m³) when m < n.

Additional optimizations include:

Specialized kernel for symmetric matrices that computes only the lower triangle.

Batching of multiple small matrices to improve GPU utilization.

Autotuning of tile shapes and pipeline depth per matrix shape, with results cached.

03 A solid performance record

Experiments on four models – WALL‑OSS‑0.5 (Zibian Robot’s base model), π0 (Physical Intelligence VLA), WALL‑WM (world model), and Qwen2.5 (text‑only LLM) – show that DMuon’s end‑to‑end training speed is only 2% slower than AdamW.

On 128 A800 GPUs, the optimizer step for WALL‑OSS‑0.5 takes 19 ms with DMuon, versus 1851 ms for naïve distributed Muon, a 97.4× reduction.

Scaling from 1 to 256 GPUs, DMuon’s throughput follows the AdamW curve, while unoptimized Muon falls far behind at 256 GPUs.

Conclusion

The key takeaway is that low‑level infrastructure such as optimizers has traditionally been a “self‑served” area for large‑model teams, but this time a robotics company, Zibian Robot, solved the distributed Muon bottleneck and is feeding those engineering advances back into the broader AI infrastructure stack.

This signals that embodied‑AI firms are evolving from model users to technology providers, adding a new pole to the AI technology landscape.