Searcharxiv⌕ Search

arXiv · 2611.07084

DMuon-2: Scaling Matrix Optimizers in Data Parallelism with Topology-Aware Matrix Ownership and Communication

Abstract

Matrix optimizers such as Muon update whole weight matrices, which conflicts with the element- or row-wise sharding of data-parallel runtimes. Once parameters are sharded as well, matrix optimizers incur extra communication or repeated computation. We present DMuon, a matrix-granular optimizer runtime for PyTorch FSDP (ZeRO-2/3) and HSDP. DMuon assigns each matrix a unique owner, which performs the Newton--Schulz (NS) iteration and optimizer update on the full matrix, and balances the load across owners by computational cost. On HSDP, ownership is global across replicas, so the cluster keeps a single copy of gradients and optimizer state, and updated parameters are published across nodes as soon as they are ready, overlapping subsequent computation. This layout makes every collective sparse and variable in length, so DMuon schedules these collectives as their data becomes ready and routes them along the cluster's two-level network. On up to 256 H100 GPUs, across three dense LLMs and two VLA models, DMuon matches or exceeds the training throughput of PyTorch FSDP2 with AdamW, reaching 1.05$\times$ that of FSDP2 with AdamW at the median on 256 GPUs. The optimized parameter gathering and gradient reduction absorb the extra communication and computation a matrix optimizer requires. Against a gather-then-replicate Muon baseline, DMuon's throughput is 2.1 times as high at the median.

Explore related subjects

Keep this discovery

Explore connections, maps & timelines

BibTeXRIS

Starrick Liu, Ruizhi Chengze, Lufang Chen, Michael Zhao, Ran Yu, Xiaofan Li, Dance Kuzi, Yue Zhang, Zijuan Lin, SongAo Han, YongHua Chen, Ruyi Gan, Hao Wang. 2026-11-06. DMuon-2: Scaling Matrix Optimizers in Data Parallelism with Topology-Aware Matrix Ownership and Communication. https://arxiv.org/abs/2611.07084

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

Online Graph Exploration under T-Interval Connectivity

We study deterministic exploration by a single agent in $T$-interval-connected graphs, a standard model of dynamic networks in which, for every time window of length $T$, the intersection of the graphs within the window is connected. The agent does not know the window size $T$, nor the number of nodes $n$ or edges $m$, and must visit all nodes of the graph. We consider two visibility models: in $KT_0$, the agent observes only its current node's identifier and port information; in $KT_1$, it additionally observes the identifiers of its currently available neighbors. We investigate two fundamental questions: the minimum window size that guarantees exploration, and the optimal exploration time under sufficiently large window size. For both models, we show that a window size $T = Ω(m)$ is necessary, and we present deterministic algorithms whose required window size is $O(ξ(n,m)\cdot m + n \log^2 n)$, where $ξ(n,m) = \frac{\ln n}{1 + \ln m - \ln n}$. This window-size bound is tight for a wide range of $m$, in particular when $m = n^{1+Θ(1)}$. Separately, for exploration time, we prove lower bounds of $Ω((m - n + 1)n)$ in the $KT_0$ model and $Ω(m)$ in the $KT_1$ model, even when the dynamic graph is $\infty$-interval-connected (i.e., even when all snapshots share a common spanning tree), and our algorithms match these time bounds up to a polylogarithmic factor. These bounds are also tight when parameterized solely by $n$: $Θ(n^3)$ for $KT_0$ and $Θ(n^2)$ for $KT_1$.

cs.DC↗

Rendezvous under Variable Disorientation:The Algorithmic Power of Fixed Unit Distance

We study deterministic Rendezvous of two robots with lights under weak geometric and synchronization assumptions. Our focus is on Variable Disorientation (VD), where each robot may arbitrarily rotate, reflect, and rescale its local coordinate system at every \Look, and on VD+Fixed-Unit-Distance (FUD), where each robot's local unit distance remains fixed. We consider the one-sided visibility models FSTA and FCOM, and schedulers ranging from the energy-restricted RSYNCH and R-RSYNCH to fully asynchronous executions. Under VD, we prove color-independent impossibility results for FSTA under SSYNCH and for FCOM under ASYNCH, and establish sharp two-color boundaries under RSYNCH and R-RSYNCH. Under VD+FUD, we introduce periodic distance classes, which turn multipl icative changes of the physical distance into predictable cyclic phase shifts. This technique yields a two-color self-stabilizing non-$L$-Rendezvous algorithm for both FSTA and FCOM under RSYNCH and R-RSYNCH even with Non-Rigid movement. Under Rigid movement, periodic distance classes further yield an optimal three-color non-quasi-self-stabilizing algorithm for FSTA under SSYNCH, a four-color FSTA algorithm under $M$-atomic ASYNCH, and a four-color FCOM algorithm under CM-atomic ASYNCH. By guarding the latter phase mechanism with two additional colors and replacing its initial swap by passive initialization, we obtain a six-color FCOM algorithm under full ASYNCH, improving the previous twelve-color upper bound by a factor of two. These results show that a fixed local metric can serve as persistent cyclic memory, and clarify how metric stability, one-sided visibility, and scheduler atomicity jointly determine Rendezvous solvability and color complexity.

cs.DC↗

Decentralized Dynamic Stable Matching

We study decentralized computing of the player-optimal stable matching under unannounced player arrivals and departures. Without a common clock, arrival/departure events information, or direct communication, players can only make individual proposals to arms and observe individual accept/reject outcomes. We measure mismatch cost --- the number of rounds in which the realized matching differs from that round's POSM, as the time horizon $T$ grows. For known preferences, we develop a deterministic polynomial-time self-calibrating protocol with mismatch cost $O(t_*+K^2\log T)$, greatly improving the existing upper-bound $O(T^{1-1/K}(\log T)^{2/K})$. $K$ is the number of arms and $t_*$ the last event time in the problem instance. Each player runs the protocol using their own local time and proposal-outcome history, without knowing~$t_*$ or $T$. More generally, configuring our protocol with a non-decreasing polynomial-time-computable function $g(t)=Ω(t)$ yields mismatch cost $O(g(t_*) + K^2 N_g(T))$, where $N_g(T)$ counts iterations of $t\mapsto t+g(t)$ below $T$. $g(t)$ reflects the protocol's convergence speed. The relationship between $g(t)$ and~$N_g(T)$ reflects the tradeoff between convergence speed and asymptotic mismatch cost. We show that any randomized protocol with convergence speed $g(t)\geq 1$ has worst-case expected cost $Ω(t_* + N_g(T))$. Thus, up to a $K^2$ factor, our protocol with $g(t)=t$ is optimal among protocols with convergence speed $O(t_*)$. For unknown preferences on both sides, each player and arm observe their own sampled utilities after being matched. We develop a decentralized learning protocol with expected mismatch cost $O(t_*+K^4Δ_{\mathrm{min}}^{-2}\log T)$, where $Δ_{\min}$ measures the preference gaps. A complementary static-setting lower-bound adapted from the literature shows that dependencies on $T$ and $Δ_{\min}$ are tight.

cs.DC↗