Conceptio › Archive › arXiv CS

arXiv CSopen access

HyperParallel-FSDP: Topology-Aware Fully Sharded Training with Layout-Driven Muon on Ascend SuperPods

· arxiv_cs

arXiv CS · Papers · License: Open Access

Open Source ↗Direct PDF ↓

cloud distributed-computing parallel-computing

distributed computing, parallel computing, cloud

arXiv:2609.21594v1 [cs.DC] 18 Sep 2026

HyperParallel-FSDP: Topology-Aware Fully Sharded Training with Layout-Driven Muon on Ascend SuperPods Mo Sun∗

Yifan Yao∗

Yanwei Liu

Zhejiang University

Huawei Technologies Co., Ltd

Luobin Liu

Zhenzhang Yang

Kaisheng Wang

Huawei Technologies Co., Ltd

Xiangyu Meng

Chen Li

Xizheng Pang

Huawei Technologies Co., Ltd

Huilan Li

Xinglei Xu

Yushi Cui

Huawei Technologies Co., Ltd

Xinyao Lin

Kaiqi Chen

Jie Zhang

Zhejiang University

Zeke Wang

Teng Su

Zhejiang University

Huawei Technologies Co., Ltd

Abstract Declarative SPMD programming, in which a sharding description attached to tensors drives all distributed execution, decouples parallelization from model code. The evaluated eager stack built on PyTorch’s native distributed tensor, however, dispatches every operator below the autograd engine, paying per-operator dispatch and metadata-handling cost on each training step, and offers no low-cost end-to-end validation path for the same plan used in production. FSDP communication and distributed Muon expose two further mismatches on two-tier supernode networks, whose intra- and inter-supernode interconnects differ by an order of magnitude: the evaluated fully-sharded configuration materializes parameters through explicit packing and unpacking passes, and the increasingly adopted Muon optimizer orthogonalizes whole matrices, which conflicts with parameter sharding. Our key insight is that distributed tensors only need to describe sharding semantics at the tensor API boundary above autograd (i.e., the function-level interception point before the autograd engine records operations), so that differentiation and kernel execution always observe plain tensors. We present HyperParallel-FSDP, a training system built on this insight with three designs: (1) Dual-Mode Distributed-Tensor Execution, where we propose one declarative sharding plan to drive both a production mode that resolves layouts once at plan-application time with zero steady-state per-operator DTensor dispatch overhead, and a validation mode that propagates layout metadata end-to-end ∗ Contributes equally.

with fail-fast checking, closed by a gradient-equivalence harness. (2) Topology-Aware Fully-Sharded Data Parallelism, where we propose zero-copy per-parameter collectives for the dominant leading-dimension-sharded layouts within a supernode, fused reduction across supernodes, and a fourstep cross-layer backward pipeline that keeps the backward path free of layer-level waits on slow links. (3) Layout-Driven Distributed Muon, including communication groups derived from sharding semantics, deduplication of orthogonalization across shard and replica domains, and shape-fused batched Newton–Schulz iterations. We evaluate HyperParallel-FSDP on Atlas 900 A3 SuperPoD from 16 dies to 384 physical cards (768 ranks). The topology-aware fully-sharded layers sustain a 505Bparameter MoE at 421k tokens/s at 768 ranks with FSDP communication occupying 2.9% of step time, which reduces mean step time by 29.7% against the evaluated PyTorch FSDP2 configuration at 16 dies and by 25.5% against Megatron DDP with distributed optimizer at 128 ranks, while per-step losses track the baseline at Pearson 𝑟 > 0.999997 over 1,000 steps. The layout-driven distributed Muon reduces profiler step time by 5.4~16.0% against DMuon, MatrixFSDP, and TorchTitan FlexShard/DistMuon integrations at 16 NPUs. HyperParallel-FSDP is open source at https://atomgit.com/mindspore/hyper-parallel.

Introduction

Large language model training has moved to clusters of 103 ~104 accelerator dies organized as supernodes, and at this scale training efficiency is decided by how communication is organized, not by raw compute [8, 22]. The hardware premise has changed with it: modern NPU supernodes

Generative AI tools were used to assist with manuscript drafting, organization, and language refinement. All AI-assisted content was reviewed and verified by the authors, who take full responsibility for the final manuscript.

connect hundreds of dies through a high-bandwidth fabric, while inter-supernode traffic traverses a conventional datacenter network with roughly an order of magnitude lower bandwidth, producing a pronounced two-tier network hierarchy [28]. The programming models in use today sit at two poles, neither satisfactory on this hardware. Megatron-class frameworks deliver high performance, but their parallelization is deeply coupled with model code. For instance, tensor, pipeline, and context parallelism are hand-wired into the model definition, together with custom checkpointing, which is unfriendly to LLM algorithm developers [22]. Declarative SPMD programming, exemplified by TorchTitan and NeMoAutoModel, offers a compelling advantage: the model is decoupled from the parallel strategy, which is expressed as data over a device mesh [14, 16]. Yet adoption remains limited in the production stacks we study, and the existing realizations explain why: built on PyTorch’s native distributed tensor, they perform operator-level dispatch below the autograd engine and pay per-operator dispatch and metadata-handling overhead for the whole training run and expose no separate validation mode that checks the same static plan used by production, leaving manual debugging as the fallback [1, 19]. We argue that performance-versus-usability trade-off is removable. FSDP communication and distributed Muon remain mismatched to supernodes. First, the evaluated PyTorch 2.9 FSDP2 configuration communicates in fused units surrounded by explicit packing and unpacking passes [20, 27]. This fusion amortizes message-launch overhead on generic Ethernet; in communication-rich supernode configurations, host-side fusion alone may provide limited end-to-end benefit. Second, the increasingly widespread Muon optimizer that is validated at frontier scale by Kimi K2’s MuonClip, Kimi K3, DeepSeek-V4, and the GLM series [3, 5, 10–12, 15] orthogonalizes whole matrices, which conflicts with sharded parameter layouts. To bridge this gap, existing distributed realizations restore whole-matrix optimizer inputs through various mechanisms, including bucketed ZeRO-1-style assignment, explicit per-parameter owner routing, ownershaped ZeRO-3 placement, strategy-specific static partitioning, and storage-to-compute resharding [2, 4, 15, 17, 24]. In contrast, HyperParallel-FSDP’s realization differs in how its communication domains are derived and bounded (§8): communication domains and deduplication follow the distribution semantics of general multi-dimensional meshes rather than a fixed per-strategy assignment; parameters whose matrix plane is unsharded require no Muon-specific optimizer communication over the expert-sharding axis; and replica-domain deduplication is confined to topology-aligned subgroups (aligned with machine boundaries under the configured rank mapping) with deterministic recomputation across them, so that under the canonical mesh-to-topology

mapping (Figure 1) the evaluated configuration sends no optimizer traffic over the slow interconnect tier. In this paper, we present HyperParallel-FSDP, a training system that resolves the challenges with one decision and two mechanisms built on it. (1) Dual-Mode DistributedTensor Execution (§4): the enabling decision is to intercept at the tensor API layer above autograd rather than at the operator dispatch layer below it, so a single declarative sharding plan drives both a production mode free of steadystate per-operator DTensor dispatch and a validation mode with fail-fast layout-contract checking, closed by a gradientequivalence harness. (2) Topology-Aware Fully-Sharded Data Parallelism (§5): per-parameter collectives that are copy-free for the dominant layouts within a supernode, fused reduction across supernodes, and a four-step backward pipeline whose layer-level path contains no wait on slow communication, where residual work is settled once at the end of the backward pass. (3) Layout-Driven Distributed Muon (§6): communication domains and two-level deduplication derived automatically from distribution semantics, so redundant orthogonalization is eliminated within each shard group and each topology-aligned replica subgroup on general multidimensional process meshes. We evaluate HyperParallel-FSDP on Atlas 900 A3 SuperPoD from 16 dies to 768 ranks, through dual-mode equivalence and host-overhead measurements (§7.2), two whole-system fully-sharded comparisons. Compared to PyTorch FSDP2 at 16 dies, HyperParallel-FSDP reduces mean step time by 29.7%; compared to Megatron DDP with distributed optimizer at 128 ranks, HyperParallel-FSDP reduces by 25.5% while per-step losses track the baseline at Pearson 𝑟 > 0.999997 over 1,000 steps. Also, a 505B-parameter MoE scale demonstration at 384 physical cards (§7.3), and distributed Muon comparisons against DMuon, MatrixFSDP, and TorchTitan FlexShard/DistMuon integrations at 16 NPUs, where profiler step time drops by 5.4~16.0% and the owner-resident baselines surface a memory imbalance and a recompute-exposed gradient-lifetime bug that the shard-native design avoids (§7.4). In summary, this paper makes the following contributions: • A dual-mode distributed-tensor abstraction built on pre-autograd, API-level interception, including zero steady-state per-operator DTensor dispatch in production mode and fail-fast verification in validation mode, together with a declarative sharding planner that derives strategies automatically from model structure (§4). • A topology-aware fully-sharded layer combining zerocopy per-parameter collectives for eligible leadingdimension layouts, fused cross-supernode reduction over per-step buffer views, and a four-step backward pipeline with single-point settlement (§5).

• A distributed Muon whose communication domains and two-level deduplication are derived automatically from distribution semantics, exposed to the algorithm developer as a local-tensor interface backed by an internally planned global execution, with shape-fused batched orthogonalization and pipelining (§6). • A multi-scale evaluation on Atlas 900 A3 SuperPoD (16 dies to 384 physical cards) covering dual-mode execution equivalence, two whole-system fully-sharded comparisons with 1,000-step loss alignment, a 505Bparameter MoE scale demonstration, and distributed Muon comparisons against three baseline integrations with expert coverage and baseline-issue disclosures (§7).

Background and Motivation

2.1

Supernode Architecture and Two-Tier Networks

Supernode 0 r0

which approximates msign(𝐺) = 𝑈𝑉 ⊤√︁from the SVD, and applying an RMS-matched scaling 0.2 max(𝑚, 𝑛) so that the AdamW learning-rate scale transfers directly [10, 15]. The per-step coefficients (𝑎𝑘 , 𝑏𝑘 , 𝑐𝑘 ) define the variant: HyperParallel-FSDP defaults to an asymmetric five-step variant and retains the legacy constant-coefficient quintic as a compatibility option (§6.6). Two properties are decisive for distribution. First, Muon applies only to matrices: twodimensional parameters directly, higher-order parameters sliced along their matrix planes, and one-dimensional parameters not at all (they are delegated to AdamW). Second, orthogonalization is global: the Gram factor couples one matrix dimension in its entirety, so a shard produced by any row- or column-wise cut cannot in general be orthogonalized in isolation. Muon has been validated at frontier scale, including Kimi K2’s MuonClip variant [11], Kimi K3 [12], DeepSeek-V4 up to the 1.6T V4-Pro [3], and the GLM-4.5/5 series [5], and is now supported by DeepSpeed and Megatron-LM [22]; it has become an increasingly adopted alternative to AdamW for frontier pretraining. 2.4

Distributed Tensors and Autodiff

PyTorch’s native distributed tensor performs operatorlevel distributed dispatch below autograd [1, 19]: every operator invocation passes through dispatch and distributedmetadata handling, while sharding propagation may hit cached Python/C++ fast paths and redistribution is inserted only when the propagated layout requires it. This generality comes at two costs: a per-operator dispatch and metadata overhead paid on every training step even on the warm path, and a single execution mode that provides no validate/production dual-mode contract sharing one static plan. Declarative SPMD frameworks such as TorchTitan and NeMo-AutoModel build on this primitive [14, 16]: they

The Muon Optimizer and Newton–Schulz Iteration

Muon updates a weight matrix 𝑊 ∈ R𝑚×𝑛 by accumulating momentum 𝐺, orthogonalizing the update direction with a Newton–Schulz (NS) iteration 𝑋 0 = 𝐺/(∥𝐺 ∥ 𝐹 + 𝜖),

Figure 1. Two-tier supernode network and the hybridsharding mapping. The shard dimension lives inside a supernode’s high-bandwidth fabric; the replica dimension crosses the slower datacenter tier. Every communication policy in this paper is derived from this asymmetry. (AG/RS: allgather/reduce-scatter.)

Fully Sharded Data Parallel and HSDP

𝑋𝑘+1 = 𝑎𝑘 𝑋𝑘 + 𝑏𝑘 (𝑋𝑘 𝑋𝑘⊤ )𝑋𝑘 + 𝑐𝑘 (𝑋𝑘 𝑋𝑘⊤ ) 2𝑋𝑘 ,

Supernode 1

Fully sharded data parallelism (FSDP) shards parameters, gradients, and optimizer states across the data-parallel group, allgathers each parameter on demand, releases it after use, and reduce-scatters gradients in the backward pass [21, 23, 27]. Hybrid sharded data parallelism (HSDP) adds a replica dimension: parameters are sharded within each group and replicated across groups, with a cross-replica gradient all-reduce following the intra-group reduce-scatter [18]. The canonical mapping onto supernode hardware places the sharding dimension inside a supernode and the replica dimension across supernodes, matching communication cost to link speed [8] (Figure 1). 2.3

Shard dim (intra-supernode): AG / RS Replica dim (intersupernode, ∼ 10× slower): fused gradient all-reduce

Modern NPU supernodes connect hundreds of accelerator dies through a high-bandwidth intra-supernode fabric and an inter-supernode datacenter network, which features an order of magnitude higher latency and lower bandwidth [28]. This two-tier hierarchy is the hardware premise of every design decision in this paper: within a supernode, singleparameter messages are already efficient and latency is low; across supernodes, message count and fusion dominate. In this paper, we refer to a machine (equivalently, a node) as a physical host, and multiple machines group as a supernode. We refer to high-bandwidth domain (fast tier) as all intramachine links, and everything leaving the supernode as the slow tier. 2.2

(1) 3

Model Structure + Device Mesh + Declarative User Overrides

achieve the decoupling of parallelization from model code that Megatron-style frameworks lack [22], but inherit the dispatch overhead everywhere.

One Static Sharding Plan (§4.1)

2.5 Motivation: Design Pain Points This paper is motivated by the three main problems of current PyTorch: First, the packing and unpacking passes around the fused collectives of the evaluated PyTorch 2.9 FSDP2 configuration burn memory bandwidth linearly in the number of layers; in communication-dominated configurations, hostside fusion alone may provide limited end-to-end benefit (§5.1). Second, naive distributed Muon repeats the orthogonalization once per shard rank, and again per replica under hybrid sharding; moreover, small matrices executed one by one are host-bound in dispatch rather than arithmetic (§6.1). Third, waiting on inter-supernode reductions inside the backward pass forces faster intra-supernode collectives to queue behind the slowest traffic (§5.4).

[PyTorch]

auto-derived, override-combined, statically linted: perparameter placements, boundary transitions, layout contracts; one plan drives both production and validation modes topology layouts contracts

Dual-Mode DistributedTensor Execution (§4) boundary contracts; production / validation modes [PyTorch]

Topology-Aware FSDP/HSDP (§5) two-tier collectives from the mesh-totopology mapping; backward pipeline [PyTorch / MindSpore]

Layout-Driven Distributed Muon (§6) grouping, owners, and dedup domains from layouts; Muon + AdamW [PyTorch]

Shared Platform Abstraction + Collective Runtime: PyTorch / MindSpore backends; vendor collective-communication library

System Overview Figure 2. Architecture of HyperParallel-FSDP: First is one static sharding plan derived from the model structure, the device mesh, and declarative user overrides. The plan materializes the placements and contracts consumed, directly or indirectly, by three subsystems: (1) Dual-mode distributedtensor execution (boundary contracts; production and validation modes share the same plan), (2) topology-aware FSDP/HSDP (the two-tier mapping shapes its collective domains), and (3) layout-driven distributed Muon (layouts shape its grouping and dedup domains). Backend badges mark where each component runs: the core distributedtensor and fully-sharded layers run on both backends, while the planner, the above-autograd interception, and the distributed optimizer are implemented on the PyTorch backend and evaluated in this paper. The platform abstraction and collective runtime at the bottom are shared infrastructure, not a pipeline stage; symmetric memory (not shown) serves only specific fused TP/EP kernels. (TP/EP: tensor/expert parallelism.)

HyperParallel-FSDP is organized around a single static artifact (Figure 2): a declarative sharding plan derived from the model structure (§4.1) and combined with declarative user overrides. The plan materializes the placements and boundary contracts consumed by three subsystems either directly or indirectly. The dual-mode distributed-tensor runtime (§4.3) executes its boundary transitions: in production mode parameters are permanently materialized as plain local tensors and only precompiled boundary collectives execute; in validation mode layout metadata propagates end-to-end with fail-fast contract checking. The topology-aware fullysharded layer (§5) reads the plan’s distribution semantics to run per-parameter zero-copy collectives within a supernode and fused reduction across supernodes. The distributed optimizer layer (§6) derives grouping, compute owners, and dedup domains from the plan’s placements, running layoutdriven Muon for matrix parameters chained with AdamW for the rest. The three subsystems share a platform abstraction and a collective runtime. The platform abstraction spans the PyTorch and MindSpore backends. The core distributedtensor and fully-sharded layers run on both, while the dualmode planner, the above-autograd interception (§4.4), and distributed Muon are implemented on the PyTorch backend; this paper evaluates the PyTorch 2.9/Ascend path, and the MindSpore backend is an engineering capability rather than an equally evaluated target of the three contributions. The collective runtime is the vendor collective-communication library. Separately, a symmetric-memory facility serves only specific fused tensor- and expert-parallel kernels; it is not a shared dependency of the three designs.

Three design principles recur throughout. Model/system decoupling: the strategy is declarative data, injected without touching model code. Explicit, verifiable semantics: every distribution fact lives in the plan and is checkable in validation mode. Topology awareness everywhere: communication policy, dedup domains, and pipelining are all derived from the two-tier network structure rather than tuned per model. Tables 1 and 2 position HyperParallel-FSDP against the PyTorch-native stack and prior distributed Muon along these principles.

Table 1. Distribution semantics and fully-sharded communication: HyperParallel-FSDP versus the evaluated PyTorch DTensor + FSDP2 configuration (version to be pinned in §7.1). “–” marks a feature not provided by that baseline configuration.

Distribution semantics Declarative sharding plan Zero steady-state per-operator DTensor dispatch Fail-fast validation mode Fully-sharded communication Zero-copy per-parameter collectives (leading-dim layouts) Two-tier (supernode-aware) policy No layer-local wait on slow tier

Evaluated PyTorch DTensor + FSDP2 config.

HyperParallelFSDP

manual –

auto-derived ✓

–

✓

–

✓

manual mesh

✓

–

✓

shard-only

✓

–

✓

Semantic roles by name rules

Parameters grouped by owning module

Structural types from member roles

6. Static Linting

5. Override Merging Declarative file: auto / none / when

4. Template Instantiation Placements + I/O contracts per boundary type

Figure 3. The six logical passes of the sharding planner. The plan is a static, data-level artifact: machine-checkable before a single collective is launched. which device-mesh dimension, and the annotations are intertwined with model code. We observe, however, that the sharding of a supported model family is determined almost entirely by the semantic role each parameter plays (e.g., columnparallel projection, normalization scale, and expert weight), rather than by architecture-specific detail. HyperParallelFSDP therefore derives the sharding plan statically from the module graph. A sharding plan assigns to every parameter a placement over the device mesh and to every module boundary a pair of input/output layout contracts. The planner produces this plan in six logical passes (Figure 3). (1) Role Classification: each parameter is classified into a fixed semantic-role taxonomy (e.g., column-wise or row-wise projection, normalization, embedding, routed or shared expert, fused QKV) by segment-aware matching of its qualified name against a small rule set, optionally refined by per-architecture overrides. (2) Boundary Grouping: parameters are grouped by their owning module, and each group forms a candidate communication boundary. (3) Boundary-Type Inference: each boundary is classified into a structural type (attention, MLP, normalization, embedding, output head, MoE routing, MoE expert) from the roles of its members. (4) Template Instantiation: boundary templates, one per boundary type, map roles to concrete placements and generate the input/output layout contracts for every active mesh dimension. (5) Override Merging: a declarative override file, with the sentinels auto (derive automatically), none (leave local), and when (conditional clause), is merged over the derived plan, so users state only what deviates from the default. (6) Static Linting: the merged plan is checked for divisibility of sharded dimensions, assignment of every parameter to exactly one boundary specification, and complete coverage of trainable parameters; any violation aborts before a single collective is launched. The role taxonomy and boundary templates cover the supported Transformer families (dense attention/MLP,

4 Dual-Mode Distributed-Tensor Execution and Autodiff This section presents how HyperParallel-FSDP executes and validates a sharding plan. We first describe how the plan is derived automatically from the model structure (§4.1) and how boundary communication is precompiled when the plan is applied (§4.2). We then present the production and validation execution modes (§4.3), explain the architectural decision that makes them possible by intercepting above autograd rather than below it (§4.4), and close with the numerical harness that turns the two modes into a verification loop (§4.5). 4.1

3. BoundaryType Inference

Sharding Plan: Per-parameter placements + boundary layout contracts

Moonlight distr. HyperParallelMuon FSDP –

2. Boundary Grouping

Divisibility, single-boundary assignment, coverage

Table 2. Distributed Muon: HyperParallel-FSDP versus Moonlight’s distributed Muon [15] (bucket-based, ZeRO1-style). “–” marks a feature not provided by the baseline.

Comm. domains from layout semantics Shard + replica two-level dedup Shape-fused batched orthogonalization

1. Role Classification

Declarative Sharding Planning

Manually authored, model-specific parallelization plans scale poorly with model diversity: each new architecture requires its own annotation of which parameter is sharded along

One Sharding Plan: Placements + Logical Boundary Transitions + Contracts

MoE, and MLA-style factorized projections); parameters that match no rule fall back to the override file, so coverage degrades gracefully to manual specification rather than failing silently. The effect is that the repeated per-model parallelization rules of manually authored plans are centralized into a small set of naming rules. Because the plan is a static, data-level artifact, it becomes machine-checkable before execution. 4.2

The only behavioral fork: one mode bit

Precompiled Boundary Communication

The layout transformations between adjacent modules are the only communication the sharding plan’s forward contracts induce in production-mode execution. HyperParallelFSDP materializes every boundary’s transformation into a precompiled transition at plan-application time: Transitions in the supported set are lowered to concrete collective sequences (tensor-parallel collectives further to differentiable local-tensor collective primitives), which are cached and replayed unchanged at every step; all other transitions retain a generic distributed-tensor redistribution fallback, so expressiveness is never gated by the lowerer. This design occupies a deliberate middle point between two extremes. Eager per-operator dispatch, as in PyTorch’s native distributed tensor, retains per-operator distributed dispatch and metadata handling at every operator invocation; whole-graph compilation, as in XLA-style auto-parallel systems, removes dispatch cost but forfeits eager debuggability and model flexibility. HyperParallel-FSDP keeps eager execution within modules. In validation mode, an operator still passes through interception, argument preprocessing, cache-keyed layout inference, local execution, and output checking, with the cache eliminating only repeated layout inference, while production mode runs none of this. Meanwhile, HyperParallel-FSDP precompiles communication across module boundaries, which is where all production traffic induced by the plan’s forward layout contracts occurs. Two scoping remarks complete the picture. First, a fallback transition in production mode executes as generic redistribution directly on plain local tensors through the same collective library, so it never re-enters the wrapper machinery; only validation mode routes it through the generic distributedtensor path so that layouts remain tracked. Second, communication inside user-declared opaque regions, such as expert all-to-all routing, context-parallel attention, and custom kernels, belongs to those regions’ own declared contracts, not to the planner’s forward-contract traffic accounted for above. 4.3

Production Mode

Validation Mode

Parameters permanently unwrapped to plain local tensors at plan-application time

Layout metadata propagates end-to-end through three-phase dispatch

Zero Per-Operator DTensor Dispatch

Fail-Fast, Precise Error Locations

forward = local compute + precompiled boundary collectives

Boundary input/output contracts checked

Gradient-Equivalence Harness (§4.5) One training step in both modes; replica-axis reduction simulated; gradients must agree

Figure 4. Dual-mode execution. Both modes share the same plan, the same logical boundary transitions, and the same contracts (concrete primitives may differ at fallback transitions), so a validate-mode pass is evidence about what production executes; the gradient-equivalence harness closes the loop at the value level. differ where a transition takes the generic redistribution fallback (production can use the local lowerer where validate redistributes through the generic distributed-tensor path). In production mode, parameters are permanently materialized as plain local tensors when the plan is applied. The forward pass performs purely local computation stitched together by the precompiled boundary collectives; no layout inference, metadata propagation, or wrapper dispatch executes at all, so production mode pays no steady-state per-operator DTensor dispatch. In validation mode, tensors carry their layout metadata end-to-end. Every operator invocation propagates layouts through the distributed-tensor semantics, and every module boundary checks its input and output contracts; a mismatch raises an immediate, precisely located error rather than a downstream numerical corruption. Because both modes execute the same plan with the same logical boundary transitions and contracts, a validate-mode pass over a model is strong evidence that production mode executes the intended sharding: verification is a property of the plan, not of a separate code path. User-injected regions, including expert all-to-all routing, context-parallel attention, and custom kernels, must declare themselves under a region axiom: Either transparent, in which case their outputs are checked against the downstream contract; or opaque, in which case they are treated as entrusted black boxes whose layout obligations are explicit in the plan.

Production and Validation Modes

A single sharding plan, with its logical boundary transitions and contracts, drives two execution modes (Figure 4): a production mode and a validation mode (code mode name validate); the only behavioral fork is one mode bit in the execution plan. The two modes share the plan, the logical transitions, and the contracts; the concrete primitives may 6

4.4

API-Level Sharding Semantics: Above Autograd, Not Below

that the resulting gradients agree. This closes the loop: Declared (the plan) = verified (validation mode) = what production runs. The harness is a design and engineering verification mechanism over the supported model architectures, not a complete numerical-equivalence proof; §7.2 presents the end-to-end evidence: over the evaluated 60-step horizon the two modes execute identical device work with identical loss and gradient-norm trajectories, and §7.3 adds a 1,000-step whole-system run whose per-step losses track the baseline at Pearson 𝑟 > 0.999997.

The dual-mode design is enabled by one architectural decision: HyperParallel-FSDP intercepts at the tensor API layer above autograd by a function-level callback that fires before the autograd engine records the operation (PyTorch’s __torch_function__), whereas PyTorch’s native distributed tensor dispatches at the operator layer below autograd, inside the recorded graph (the exact Python/C++ dispatch path is version-dependent [19]). Four consequences follow. First, the distributed tensor describes only forward sharding semantics. It is an auxiliary, API-level description of how data is laid out, rather than the model’s core data structure. Second, graph construction, differentiation, and kernel execution always observe plain local tensors. The wrapper never enters the operator layer, so the autograd engine and the device kernels require no awareness of distribution whatsoever. Third, switching between production and validation modes is a metadata-level operation: because no operatorlevel machinery depends on the wrapper, production mode can discard it entirely at plan-application time. Fourth, distribution semantics and differentiation semantics are orthogonal by construction. Forward layout inference handles distribution by a three-phase dispatch that preprocesses arguments, infers the output layout through a cached rule table, and expands to local execution; ordinary local autograd handles differentiation. No distributed autograd engine is needed. Gradients materialize on local shards, where the data-parallel hooks reduce-scatter them (§5); residual gradients along tensor-parallel replica axes are reduced exactly once at the end of the backward pass, preserving the mathematical semantics of a replicated computation. The placement algebra itself is correspondingly richer than the shard/replicate/partial trio of the native design: it supports strided and ragged sharding, uneven (ceil-chunk) sharding for dimensions not divisible by the mesh size, and shard-aware random number generation with per-shard offset tracking so that stochastic layers remain bit-reproducible under resharding. 4.5

Topology-Aware Fully-Sharded Parallelism

This section describes how HyperParallel-FSDP materializes sharded parameters and reduces gradients over a two-tier HSDP mesh. The design premise is the two-tier network of supernode hardware (§2.1), a two-tier communication hierarchy in which, within a supernode, links are fast enough that single-parameter messages are already efficient; across supernodes, bandwidth is an order of magnitude scarcer and message count dominates. Every choice below follows from treating the two tiers differently: per-parameter zero-copy collectives within a supernode (§5.1), fused reduction across supernodes (§5.2), a backward schedule whose layer-level path never waits on the slow tier (§5.3, §5.4), and the engineering capabilities that make the scheme practical (§5.5). 5.1

Per-Parameter Zero-Copy Communication

The evaluated PyTorch 2.9 FSDP2 baseline configuration (§7.1) [20, 27] implements fused-unit communication with explicit packing. For all-gather, a copy-in pass packs every shard of a unit into one contiguous buffer before the collective, and a copy-out pass unpacks the result back into per-parameter storage afterwards; for reduce-scatter, gradient shards are likewise packed before the collective, and the fused output is then copied back (or exposed as views) per parameter. The copy-out pass sits on the critical path because downstream computation depends on it. We state this data path as configured in our evaluation; without trace evidence we do not claim that every copy is exposed on the critical path. Fusion amortizes message launch overhead, which especially matters on general-purpose Ethernet, for which the design was made. Within a supernode, this trade inverts. Links are highbandwidth and low-latency, so a single parameter’s shard is already a large, efficient message: all-gather is prefetched, while gradient reduction is issued as an asynchronous reducescatter, so communication can overlap with computation. Fusion’s launch savings therefore provide little benefit in this regime while its copy cost is real. HyperParallel-FSDP therefore issues collectives per parameter and eliminates the copies for the leading-dimension sharded layouts that

Gradient Equivalence Validation

Layout checking alone does not establish that production mode computes the same values as the semantics describe. HyperParallel-FSDP therefore ships a numerical harness that lifts validation mode from layout checking to value checking: it runs one training step in both modes, simulates within validation mode the replica-axis gradient reduction that production performs inside the data-parallel hooks, and asserts

(a) Fused Path (evaluated PyTorch 2.9 FSDP2 configuration): packing and view reconstruction around an all-gather 𝑝 1 shard 𝑝 2 shard 𝑝 3 shard

Copyin (pack)

Fused input buffer

Fused output buffer

Copyout (unpack)

𝑝 1 ..𝑝 4 fullparameter buffers

Copy kernels occupy the compute stream

𝑝 4 shard

(b) Per-Parameter Zero-Copy Path (HyperParallel-FSDP): no data movement around the collective 𝑝𝑖 shard (the parameter’s own storage)

Persistent perparameter buffer; full 𝑝𝑖 is a view of it

Input is the shard itself: zero copyin; output lands in the parameter’s own buffer: zero copyout; collective payload volume identical to (a)

Figure 5. All-gather data paths within a supernode, for the evaluated PyTorch 2.9 FSDP2 configuration and HyperParallel-FSDP. (a) The fused scheme of the evaluated baseline surrounds the all-gather with packing and view-reconstruction passes over device memory. (b) HyperParallel-FSDP issues one collective per parameter and removes both copies for leading-dimensionsharded layouts; the full parameter is a view of its persistent buffer. (a) Split policy: many small zero-copy collectives inside, one fused collective across

(b) Fusion without copies: reduce-scatter writes buffer views

Inter: one fused all-reduce (async) Supernode 0 r0

Supernode 1 r3

Intra: per-param AG/RS, zero-copy

RS 𝑝 1

RS 𝑝 2

RS 𝑝 3

𝑝 1 view

𝑝 2 view

𝑝 3 view

Intra: per-param AG/RS, zero-copy

pad

one all_reduce over the whole buffer (SUM), asynchronous Buffer base 512B-aligned; total length padded, pre-zeroed: padding is neutral under SUM; gradients read back from views. Zero copyin, zero copyout

Figure 6. Two-tier hybrid sharding. (a) Per-parameter zero-copy collectives run within each supernode; a single fused allreduce crosses supernodes per unit. (b) The fusion itself is copy-free: per-parameter reduce-scatter outputs land directly in their corresponding views of one contiguous buffer with a 512-byte-aligned base address and a padded total length, and one collective reduces the whole buffer. (AG/RS: all-gather/reduce-scatter.) dominate in practice: the all-gather input is each parameter’s own shard, and the output is written directly into a persistent per-parameter buffer of which the materialized full parameter is merely a view; reduce-scatter is symmetric. (Parameters sharded on other dimensions take a chunkand-concatenate copy-out fallback.) The collective payload volume is unchanged; the path removes the surrounding device-memory packing and unpacking traffic in the eligible layouts (Figure 5).

and reduce-scatter within each supernode (zero-copy), one fused all-reduce across supernodes per unit (Figure 6). Crucially, the fusion itself is zero-copy. A contiguous gradient buffer is allocated per unit in each step; the allocation’s base address is 512-byte aligned and the total buffer length is padded to a 512-byte multiple. Padding is zeroinitialized and is therefore neutral under SUM. Each reducescatter writes its output directly into the corresponding view of the fused buffer. One fused all-reduce then covers the whole buffer, and each parameter reads its final gradient back from its own view, thus no parameter ever moves. Intra-supernode traffic trades zero copies for bandwidth, inter-supernode traffic trades fusion for message count, and the shared buffer views make the two policies compose. On the lower-bandwidth inter-supernode tier, aggregation reduces launch and message-count overhead in the evaluated configuration.

5.2 Two-Tier HSDP: Fused Reduction without Copies Across supernodes the trade inverts again: at the intersupernode bandwidths and message sizes of the configurations evaluated here, fewer and larger messages are generally preferable. HyperParallel-FSDP maps hybrid sharded data parallelism onto the two tiers with a split policy: HyperParallel-FSDP performs per-parameter all-gather 8

(a) Evaluated PyTorch FSDP2 configuration: fused units + same-layer RS/AR chain Bwd 𝐿𝑁 + grad copyin

Bwd 𝐿𝑁 −1 + grad copyin

Bwd 𝐿𝑁 −2 + grad copyin

···

AG 𝐿𝑁 + copyout

AG 𝐿𝑁 −1 + copyout

AG 𝐿𝑁 −2 + copyout

···

RS 𝐿𝑁 (fused)

RS 𝐿𝑁 −1 (fused)

RS 𝐿𝑁 −2 (fused)

···

AR 𝐿𝑁

AR 𝐿𝑁 −1

AR 𝐿𝑁 −2

Compute

Copy kernels ride the compute stream; AR chained to samelayer RS; per-layer AR events re-checked at many points: slow tier leaks into backward

···

(b) HyperParallel-FSDP: four-step post-backward pipeline; the slow tier is issue-only and settles once at the end of backward Compute

Bwd 𝐿𝑁

PB𝑁

Bwd 𝐿𝑁 −1

PB𝑁 −1

PB𝑁 −2

Bwd 𝐿𝑁 −2

···

AG 𝐿𝑁 per-param

AG 𝐿𝑁 −1 per-param

AG 𝐿𝑁 −2 per-param

···

RS 𝐿𝑁 zero-copy

RS 𝐿𝑁 −1 zero-copy

RS 𝐿𝑁 −2 zero-copy

···

Fused AR 𝐿𝑁

Fused AR 𝐿𝑁 −1

···

End-of-backward settlement: wait-all ARs, one write-back pass

RS of 𝐿ℓ overlaps bwd of 𝐿ℓ −1 ; AR issued one layer late, never waited on mid-backward

Time (backward of 𝐿𝑁 → 𝐿𝑁 −1 → · · · → end of backward)

Figure 7. Backward communication timelines. (a) The evaluated PyTorch 2.9 FSDP2 configuration chains each layer’s allreduce to its own reduce-scatter and surrounds collectives with copy kernels, so slow inter-supernode latency enters the backward path. (b) HyperParallel-FSDP issues per-parameter reduce-scatter at each layer’s post-backward step and defers the fused all-reduce asynchronously by one layer; nothing waits on the slow tier until a single end-of-backward settlement. (PB: post-backward step; AR: all-reduce.) 5.3

Algorithm 1 Post-backward step of layer ℓ (four-step pipeline; layer indices in forward order, backward runs 𝑁 → 1)

Four-Step Backward Communication Pipeline

The backward pass must overlap these collectives without layer-local waits on the slow tier. HyperParallel-FSDP schedules it with a deadline-driven argument: a layer’s reducescatter has a deadline one layer of backward computation away (the optimizer will need it soon), while the fused allreduce’s deadline is the end of the whole backward pass. Fast communication therefore goes first; slow communication consumes the slack. Concretely, each layer’s post-backward step performs exactly four actions (Algorithm 1, Figure 7): (1) wait for the previously completed layer’s reduce-scatter, which has already overlapped with this layer’s backward computation and is approximately free; (2) settle gradients of parameters that need no cross-replica reduction; (3) issue this layer’s reduce-scatter with its output written directly into the corresponding fused-buffer view; and (4) asynchronously issue the previously completed layer’s fused all-reduce, whose completion handle is parked in a global pending queue and never waited on here. At the end of the backward pass, a single settlement performs the wait-all and writes all final gradients back in one pass. The backward path therefore contains no layer-local wait on the slow inter-supernode tier; the end-of-backward settlement enters the step’s critical

1: wait RSℓ+1 ⊲ overlapped during layer ℓ’s backward 2: settle gradients of no-reduction parameters of layer ℓ 3: issue RSℓ into fused-buffer views ⊲ zero-copy 4: issue-async ARℓ+1 ; park handle in pending queue

⊲ at end of backward: wait-all pending ARs; single write-back pass path only if the slowest all-reduce outlasts the entire backward hiding window. By contrast, the evaluated PyTorch 2.9 FSDP2 configuration of §7.1 issues the all-reduce immediately after the same layer’s reduce-scatter and maintains per-layer completion events that are re-checked at many points, letting slow-traffic latency leak into the backward path. 5.4

Interaction with TP/CP/EP Communication Domains

In hybrid parallelism, the backward path mixes communication domains of very different speeds: tensor-parallel allgather and reduce-scatter and expert/context-parallel allto-all run within the supernode, while the fused all-reduce 9

crosses supernodes and is the slowest communication in the system. Any backward-path synchronization on the slow reduction forces the fast collectives issued synchronously on the compute stream to queue behind it, so slow traffic leaks into the critical path and is amplified layer by layer. The schedule of §5.3 removes every layer-local wait on the slow tier from the backward path by construction (issueonly, with a single end-of-backward settlement), so the fast domains observe it only through that final settlement.

a general conclusion; we therefore choose a gather-based organization in this deployment regime. Existing distributed realizations restore whole-matrix optimizer inputs through materially different mechanisms: Moonlight’s bucketed ZeRO-1-style assignment, DMuon’s explicit per-parameter owner routing, MatrixFSDP’s ownershaped ZeRO-3 placement, Canzona’s strategy-specific static partitioning, and FlexShard’s storage-to-compute resharding. They differ in when communication occurs, where optimizer state resides, and how computation is scheduled. Against these realizations (Table 2, Table 3; versions as inspected, runtime revisions pinned in §7.4), HyperParallel-FSDP combines three properties in one design: (i) communication domains and redundancy derived automatically from distribution-layout semantics on general multi-dimensional meshes, including a case that requires no Muon-specific optimizer communication over the expert-sharding axis when no sharding axis intersects the matrix plane; (ii) replica-domain dedup within topology-aligned subgroups, whose size is determined from the local device count and is machine-aligned under the standard contiguous rank assignment (§6.5), with deterministic recomputation across subgroups, so that under the canonical mesh-to-topology mapping (Figure 1) the evaluated configuration sends no optimizer traffic over the slow interconnect tier; and (iii) shape-fused batched orthogonalization inside a local-tensor programming model backed by one static plan. DMuon already provides batched orthogonalization and placement-aware tensor-parallel handling, and computes each matrix exactly once globally; its owner domains are assigned over a fixed data-parallel mesh rather than derived from layout semantics, and its routing spans both interconnect tiers. MatrixFSDP places each whole matrix on one data-parallel rank (its parameter owner, holding the matrix and its optimizer state) under ZeRO-3 sharding, so its optimizer step issues no matrix collective; the owner placement is specific to the data-parallel plane and excludes tensor-parallel-fragmented matrices. Canzona statically assigns whole parameters to data-parallel ranks for communication-free updates and batches tensor-parallel fragments into asynchronous intra-node micro-groups; its domains are fixed per parallel strategy rather than derived from general layout semantics. TorchTitan’s FlexShard, a general storage-to-compute redistribution substrate, is contrasted in §6.2.

5.5 Engineering Capabilities The mechanism above composes with the practical requirements of production training: uneven (ceil-chunk) sharding on the leading dimension; strided composition with preexisting tensor-parallel layouts; mixed precision with fp32 master gradients; bucket-granularity gradient accumulation; CPU offloading; recompute-aware prefetch suppression and explicit prefetch control; and a dual PyTorch/MindSpore backend behind the platform abstraction of §3.

Layout-Driven Distributed Muon

Muon replaces the element-wise second-moment scaling of AdamW with an orthogonalization of the momentum matrix by NS iteration, and has been adopted by a growing list of frontier models (§2.3). This section describes how HyperParallel-FSDP scales Muon to multi-dimensional sharded layouts. We first state why the interaction of matrix semantics with sharding is hard (§6.1) and the design stance that shapes our answer (§6.2): §6.3~§6.7 then present its components: Grouping parameters from layout semantics (§6.3), deduplicating orthogonalization across shard domains (§6.4) and replica domains (§6.5), batching the NS iteration by core shape (§6.6), and pipelining the whole optimizer step (§6.7). 6.1

Challenges: Matrix Semantics vs. Sharded Layouts

Orthogonalization is intrinsically whole-matrix: the Gram factor 𝐺 ⊤𝐺 couples the columns of the momentum matrix 𝐺 in their entirety (and 𝐺𝐺 ⊤ the rows), so no shard of 𝐺 can be orthogonalized in isolation. Sharding, however, cuts the matrix apart, and the naive gather-then-compute on every rank remedy repeats the NS computation 𝑆-fold across an 𝑆-way shard group, multiplied a further 𝑅-fold across 𝑅 replicas under hybrid sharding. Worse, the gather, the NS iteration, and the result return are all new critical-path work inside the optimizer step, and per-matrix execution is hostbound for the many small matrices a transformer contains. An all-to-all-based realization is theoretically redundancyfree, but in our target deployment it creates an unfavorable connection and resource footprint on the network interfaces. This deployment constraint motivates our choice rather than

6.2

Design Stance: a Local Programming Model over an Internal Global Plan

The organizing principle of HyperParallel-FSDP’s Muon is that distribution is an execution concern, not an algorithmic one. The algorithm is written against local tensors: each rank sees only its gradients, its momentum, the matrices it must orthogonalize, the Newton-Schulz iteration, and the parameter update. The distribution layer answers two questions only: How each parameter is currently sharded, and 10

Table 3. Design contrast of distributed Muon realizations (versions as inspected, runtime revisions pinned in §7.4): the six dimensions most relevant to HyperParallel-FSDP’s design stance, against DMuon [2], MatrixFSDP [4], Canzona [24], and TorchTitan’s FlexShard/DistMuon [14, 17]. The full eleven-dimensional comparison matrix is in Appendix A. HyperParallel-FSDP Muon

DMuon

MatrixFSDP

Canzona

FlexShard DistMuon

Domain deriva- From layout semantics on tion general meshes (matrixplane predicate)

Owner assignment over the DP mesh; nested TP ownership from placements

Owner-shaped ZeRO-3 shards: one whole-matrix owner per 2-D weight

Logical assignment decoupled from physical layout; static per-strategy partition

User-declared compute layouts on named DeviceMesh axes

State dedup scope

Shard + replica two-level; replica state reduced to 1/𝑅𝑠 in topology-aligned subgroups

One authoritative owner per matrix (parameter + momentum)

Whole matrix and state on the owner; empty shards elsewhere

Whole-parameter state at the Owner-style state on the statically assigned rank compute layout (Owned, BlockShard)

Cross-topology strategy

Replica subgroups from Two-stage intra-/inter-node the local device count hierarchy with XOR owner (contiguous rank blocks, slots; spans both tiers machine-aligned under the standard rank mapping; a configuration premise, §6.5); deterministic recomputation; no slow-tier optimizer traffic under the canonical mapping (Fig. 1)

Optimizer step local; the backward reduction spans the DP mesh

DP step zerocommunication; TP reconstruction kept intra-node

Optimizer-step collectives

Fused all-gather + relay broadcast; none over the expert-sharding axis when the matrix plane is unsharded

None. The backward reduce- DP: none; TP: fused all-to-all Packed all-to-all redistribuscatter lands the input on the micro-groups tion owner

Owner-to-all broadcast / all-to-owner reduce

Topology from the userchosen mesh and layouts; no automatic machineboundary policy described

Batching / Load Shape-fused batched NS; balance greedy owner balance; ≤1 batch in flight

Gram-NS batching with Global owner planner balSYRK kernels and autotunancing resident bytes and ing; measured (MILP) owner optimizer work assignment

𝛼 -balanced LPT (DP); micro- BucketConfig grouping group scheduling with and ordering for packedrollback (TP) redistribution overlap

Mesh coverage

FSDP2/HSDP data-parallel mesh + nested TP

ZeRO-1 DP + TP; optimizeragnostic (Muon, Shampoo, SOAP)

General multi-dimensional meshes; expert-axis case needs no Muon-specific optimizer communication

ZeRO-3 data parallelism; TP-fragmented matrices excluded

which process groups that sharding spans. It is never asked to re-express the optimizer’s internal decisions (which rank computes a given matrix, and in what order) as a second, optimizer-facing layout system. The global organization exists, but inside the execution engine rather than in the algorithm’s vocabulary. The ranks joining each collective, the compute owner of every matrix, per-rank send and receive volumes, collective and bucket ordering, and buffer sizing and reuse are all fixed by a static plan when the optimizer is built, and replayed unchanged at every step. One step traverses: local gradient → replica-owner deduplication → momentum update → fused all-gather → load-balanced compute-owner assignment → batched NS → packed relay broadcast → slice extraction and local update → replica broadcast. An internal global plan; an external local programming model. Two consequences follow. First, communication primitives follow the hardware rather than the abstraction. Fused all-gather plus owner broadcast (i.e., hierarchical relay across mesh axes) matches the fast paths of collectivecommunication libraries, fixes collective order statically, makes buffer sizes and lifetimes pre-plannable, and avoids the rank-pair queue-pair pressure of all-to-all collectives on the slow tier (§6.1); the policy can be retargeted per

Layouts on named mesh axes; DistMuon flat matrixbatch: BlockShard on one non-unit axis only

platform without touching the algorithm interface. Second, the distributed skeleton is decoupled from the algorithm it carries: the momentum rule, the orthogonalization itself (step count, coefficients, epsilon), input reshaping and matrix extraction, RMS-matched scaling, and post-update processing are all pluggable or configurable stages, so legacy Muon, the asymmetric five-step variant, per-parameter coefficient sets, alternative momentum rules, and QK-clip-style post-processing all reuse the same owners, buckets, buffers, and communication plan. Table 3 contrasts this stance with the realizations of §6.1 along the dimensions most relevant to it (Appendix A gives the full matrix). The sharpest contrast is with FlexShard, TorchTitan’s general substrate for distributed Muon [14, 17]. FlexShard gives the optimizer a second layout system: a perparameter compute layout (block-sharded or owner-assigned) is declared next to the storage layout, and a general reshard planner lowers the storage-to-compute transition onto packed all-to-all collectives. The abstraction is complete and serves many optimizers; its costs are a second layout vocabulary for algorithm developers and an all-to-all communication profile. DMuon assigns each matrix a single authoritative owner over the data-parallel mesh and reaches 11

Parameter + Placements (either backend)

near-AdamW overhead through hierarchical communication, Gram-space batched Newton-Schulz, and measured load balancing; its domains are assigned over a fixed mesh hierarchy rather than derived from layout semantics, and its publication and reduction traffic spans both interconnect tiers. MatrixFSDP and Canzona likewise keep a conventional optimizer interface, but their domains follow fixed ownershaped ZeRO-3 placement, respectively a static DP/TP assignment, rather than the layout semantics of a general multidimensional mesh. HyperParallel-FSDP makes a different trade: specializing to Muon’s fixed structure enables statically pre-plannable communication and hardware-friendly collectives, at the cost of a general optimizer layout abstraction. The distributed tensor stays at its most valuable boundary: Describing how parameters are distributed, not everything the optimizer then does with them. 6.3

Extract layout (shape-free): which mesh axes shard which tensor dims, along which tensor dim; which axes replicate Does any sharding axis intersect the matrix plane (last two dims)?

Yes e.g., FSDP-sharded Linear [𝑜𝑢𝑡, 𝑖𝑛] on dim 0: each rank holds rows only

→ shard-domain dedup (§6.4): fused gather, owner NS, relay, slice

No e.g., MoE expert [𝐸, 𝑖𝑛, 𝑜𝑢𝑡 ] sharded on 𝐸 : each rank holds whole expert matrices

→ local batched NS (§6.6): no Muon-specific optimizer communication over the expertsharding axis, no redundancy

1-D / scalar parameters → AdamW (chained optimizer; RMS matching transfers the learning-rate scale; decoupled weight decay is applied separately)

Layout-Semantics-Driven Parameter Grouping

HyperParallel-FSDP requires no user-supplied process groups. For every parameter, the optimizer extracts from the distribution metadata which mesh axes shard it (along which tensor dimension) and which axes replicate it; one code path serves both HyperParallel-FSDP’s own distributed tensors and PyTorch-native ones, because both expose the shard/replicate placement metadata this grouping path requires. The grouping rests on one predicate (Figure 8): does any sharding axis intersect the matrix plane, i.e., the last two tensor dimensions on which NS operates? If not (the canonical case is the expert axis of a 3-D expert weight, which slices between matrices rather than through them), every locally held matrix slice is already complete, and orthogonalization is purely local, i.e., no Muon-specific optimizer communication over the expert-sharding axis, and is redundancy-free. The governing principle: communication cost is determined by which matrices are cut and by the sizes of the cut matrix planes, not by raw parameter count alone. One-dimensional parameters leave Muon entirely and are delegated to AdamW by a chained optimizer. The expert case deserves one precision. A fused expert parameter of shape [𝐸, 𝑀, 𝑁 ] denotes 𝐸 independent [𝑀, 𝑁 ] matrices: Muon must run the NS iteration per expert matrix, and the three-dimensional parameter must never be flattened into a single [𝐸, 𝑀𝑁 ] matrix for orthogonalization. 6.4

yes

Figure 8. Layout-semantics-driven grouping. Communication cost is determined by which matrices are cut and by the sizes of the cut matrix planes, not by raw parameter count alone: only matrix-plane-sharded parameters gather; expert-axis-sharded weights, which are the most numerous under expert parallelism, require no Muon-specific optimizer communication over the expert-sharding axis. each matrix is orthogonalized exactly once per shard group while all matrices of a batch still progress in parallel across the group. (3) Relay broadcast: owners pack their updates and broadcast them asynchronously, dimension by dimension, along the sharding axes. (4) Slice and apply: each rank extracts its own slice and applies the update with weight decay through fused elementwise kernels. Momentum stays sharded throughout, so optimizer-state memory never reflows. 6.5

Replica-Domain Deduplication

Replica axes contribute the second, multiplicative redundancy: 𝑅 identical copies of every replicated matrix would otherwise be orthogonalized 𝑅 times. HyperParallel-FSDP deduplicates within each topology-aligned replica subgroup (Figure 9(b)): momentum buffers exist only on one replicaowner rank per subgroup, only replica owners orthogonalize, and the updated parameters are broadcast to peers hierarchically and asynchronously, one replica dimension at a time. Writing 𝑅𝑠 for the subgroup size, the Muon momentum state of eligible matrix parameters is reduced to 1/𝑅𝑠 of the onecopy-per-replica baseline, and each matrix is orthogonalized once per subgroup, which is 𝑅/𝑅𝑠 times in total across the original replica group, degenerating to a single global computation when the group is not split (𝑅𝑠 = 𝑅).

Shard-Domain Deduplication

For parameters whose matrix plane is sharded, HyperParallelFSDP eliminates the 𝑆-fold redundancy in four steps (Figure 9(a)). (1) Fused gather: one collective per sharding axis packs a whole batch of momentum shards into an aligned staging buffer (bf16, uneven splits supported). (2) Owner assignment: each matrix is assigned to exactly one rank of the shard group by greedy size balancing (its compute owner); only the owner materializes the full matrix and runs NS, so 12

(a) Shard Domain: Each matrix orthogonalized exactly once per shard group 2. Owner Assignment

1. Fused Gather

Naive

One collective per sharding axis packs a whole batch (bf16, aligned, uneven splits)

Gather, then every rank runs NS over all 𝑆 matrices: 𝑆 × redundant compute

Greedy size balance; only the owner materializes the full matrix and runs NS

3. Relay Broadcast Owners pack updates; broadcast per sharding axis, asynchronous

4. Slice-Apply Each rank extracts its own slice; momentum stays sharded throughout

(b) Replica Domain: Within each subgroup, one of 𝑅𝑠 identical copies computes; subgroups sized from the local device count Updated-parameter broadcast

Replica 0

Replica 1

No momentum buffer, no NS

Replica 2 = owner Momentum + NS + apply: computes once

Replica 3 No momentum buffer, no NS

Updated parameters broadcast hierarchically, one replica axis at a time; Muon state of eligible matrix parameters reduced to 1/𝑅𝑠 (𝑅𝑠 : subgroup size); every matrix computed once per topology-aligned replica subgroup

Figure 9. Two-level deduplication of orthogonalization. (a) Shard domain: a fused gather feeds greedy owner assignment, so each matrix is orthogonalized exactly once per shard group instead of once per rank. (b) Replica domain: only owner ranks hold momentum and compute; the updated parameters are broadcast hierarchically, and replica groups are split into contiguous rank blocks sized from the local device count (aligned with machine boundaries under the standard rank assignment), so replica-domain traffic stays on the fast tier; shard-domain collectives stay intra-supernode under the canonical mapping (Figure 1). The dedup domain’s topology alignment is a configuration premise, not physical-topology probing: the realization derives the subgroup size from the per-machine device count by selecting the largest divisor of the replica-group size not exceeding it, with no split at all when the group already fits on one machine, and partitions each replica group’s sorted ranks into contiguous blocks of that size. This coincides with machine boundaries under the standard contiguous rank-to-node assignment of the evaluated launchers. Note that a machine (node) is defined as one physical host inside a supernode, so the boundary lies strictly inside the supernode’s high-bandwidth domain. A rank permutation that interleaved machines, or a replica group smaller than one machine yet spanning two, would violate the premise undetected; making the partition read the physical rank-tonode mapping (and fail fast when a subgroup would cross a slow boundary) is implementation future work, and the evaluated configuration’s alignment is checked with the link-tier traffic measurements of our experiment plan. Within each subgroup, deduplication and broadcast run locally, while across subgroups HyperParallel-FSDP uses deterministic re-computation. Since gradients have already been reduced identically (§5.2) and the NS iteration is deterministic, every subgroup derives numerically equivalent results with no communication (identical inputs and a deterministic iteration; bitwise identity additionally requires a fixed kernel set and software stack, which our deployments pin). Under this premise, replica-domain traffic never crosses the slow tier. Shard-domain collectives inherit the topology of their sharding axes: under the canonical mesh-to-topology

mapping of Figure 1 where sharding axes inside the supernode, replica axes across supernodes, the collectives remain intra-supernode, so the evaluated configuration sends no Muon-specific optimizer traffic over the slow tier. A mapping that placed a sharding axis across machines would need an analogous bound on the shard domain, which the current realization does not enforce automatically. 6.6

Shape-Fused Batched Orthogonalization

Per-matrix NS launches starve the device for small matrices. HyperParallel-FSDP groups parameters by their core matrix shape (Figure 10(a)): tensors of shape [𝐴, 𝐵], [𝐴, 1, 𝐵], and [𝑁 , 𝐴, 𝐵] share the (𝐴, 𝐵) group, are concatenated along a batch axis, orthogonalized by one batched NS iteration (a single sequence of batched matmuls), and split back. Parameter storage layout and Muon’s mathematical matrices are deliberately decoupled. One physical parameter can hold several logical matrices, e.g., the heads of a fused QKV projection, the gate and up matrices packed into an MoE FC1 weight, the two query factors of an MLA projection, or the key and value factors of its KV-up projection. Rather than pushing these model semantics into the distribution layout, HyperParallel-FSDP defines the matrix as Muon sees it in the model adapter: a zero-copy view reshape, an arbitrary reversible split/concatenate/interleave transform applied to the NS input, and a restore mapping that writes the logical-matrix updates back into the physical parameter. Model-structure changes thereby stay in the model adapter and never propagate into placement metadata, the sharding planner, or the communication runtime. 13

(a) Shape Fusion: One batched kernel per core-shape group

𝑝 1 [𝐴, 𝐵 ] 𝑝 2 [𝐴, 𝐵 ] 𝑝 3 [𝑁 , 𝐴, 𝐵 ] 𝑝 4 [𝐴, 1, 𝐵 ]

Concat on batch axis Same core shape (𝐴, 𝐵)

One Batched NS: 5 steps, batched matmuls

Merged [𝐾, 𝐴, 𝐵 ]

Split back; apply per-parameter updates

bf16, workspace allocated once per batched invocation; transpose-first when rows>cols; reshape hook for fused QKV; cap 229 elements scaled by shard-group size

(b) Batch-Level Pipelining: Broadcasts sink beneath compute Final drain: wait all

Compute Batch 1: NS (largest)

Batch 2: NS

Batch 3: NS

Comm.

Broadcast batch 1 (async)

Broadcast batch 2

···

Batches packed in descending size; at most one batch in flight; admitting a new batch settles the oldest; peraxis relay removes host-side bubbles

Time

Figure 10. Batched orthogonalization and pipelining. (a) Parameters sharing a core matrix shape ([𝐴, 𝐵], [𝐴, 1, 𝐵], [𝑁 , 𝐴, 𝐵]) are concatenated and orthogonalized by one batched Newton-Schulz kernel. (b) Each batch’s broadcast is issued asynchronously as soon as the batch finishes, so the timeline degenerates to mostly continuous compute; the bounded in-flight queue may settle an earlier batch before admitting the next one, and the final drain waits for all remaining work. Algorithm 2 One distributed Muon step

The kernels are engineered for the regime: bf16 arithmetic with workspace allocated once per batched invocation rather than per matrix; a transpose-first rule when rows exceed columns so the Gram product falls on the small dimension; two selectable coefficient sets (the legacy quintic and an asymmetric five-step variant); and memory-safe batching capped at 229 elements, scaled down by the shard-group size. For the many small matrices a production MoE contains, batching replaces dozens of host-bound per-matrix launches with a single batched kernel. 6.7

1: partition parameters by the matrix-plane predicate (§6.3) 2: complete path: momentum update → batched NS →

apply ⊲ no Muon-specific optimizer communication over the expert-sharding axis 3: for all batches 𝐵 of matrix-plane-sharded params., descending size do 4: fused-gather momentum shards of 𝐵 per sharding axis 5: assign matrices of 𝐵 to owners (greedy size balance) 6: owners: batched NS over core-shape groups 7: issue-async relay broadcast of 𝐵 along shard axes 8: subgroup shard ranks: slice from relay; apply with decay 9: replica owners: issue-async per-axis broadcast ⊲ ≤ 1 in flight; new batch settles oldest 10: end for 11: final drain: wait all outstanding broadcasts

Batch-Level Pipelining

Finally, the optimizer step itself is pipelined (Figure 10(b)). Parameters are packed into batches in descending size order (i.e., large matrices first) so their result broadcasts enjoy the longest hiding window. Each batch’s replica broadcast is issued asynchronously as soon as the batch finishes, with at most one batch in flight. Admitting a new batch settles the oldest, so an intermediate batch may wait briefly at that point; dimension-by-dimension relay removes host-side bubbles. The step timeline degenerates to a continuous stream of momentum and orthogonalization compute with communication sunk beneath it; a final drain waits for all remaining work (Algorithm 2).

study runs the production and validation modes of one sharding plan at two micro-batch sizes and compares device work, memory, numerics, and step time. The fully-sharded study reports two whole-system comparisons against PyTorch FSDP2 at 16 dies and against Megatron DDP with distributed optimizer at 128 ranks, plus one baseline-free scale demonstration on a 505B-parameter MoE. The Muon study compares three baseline integrations at 16 NPUs and discloses the implementation issues the comparisons surfaced. We do not perform per-component ablations and do not attribute whole-system differences to individual mechanisms; where

Evaluation

We evaluate HyperParallel-FSDP on Ascend hardware at three scales: a 16-device system, a 128-rank cluster, and a 384card (768-rank) cluster. The evaluation is organized by the three design threads: dual-mode distributed-tensor execution (§7.2), topology-aware fully-sharded parallelism (§7.3), and layout-driven distributed Muon (§7.4). The tensor-runtime 14

Table 4. Evaluated configurations. A: dual-mode tensor-runtime study (production vs. validation mode at MBS 1 and 2 under one sharding plan; §7.2). B: 16-die pure-FSDP comparison against PyTorch FSDP2 (§7.3). C: 128-rank dense-model comparison against Megatron DDP with distributed optimizer (§7.3); the baseline column of B/C is described in the text. D: 505B-parameter MoE scale demonstration without a baseline (§7.3). Parallel meshes describe the HyperParallel-FSDP side. Field

A: dual-mode

B: 16-die FSDP

C: 128-rank dense

D: 505B MoE

Model

Qwen3-30B-A3B (48 layers, 128 experts, top-8)

Qwen3-30B-A3B (same as A)

7BV2 dense (39 layers, GQA, SWA, MTP-4)

505B MoE (51 layers, MLA, 384 experts, top-8)

Hardware

16 Ascend NPU devices (16 ranks)

Ascend A3 system, 16 NPU dies

8 nodes × 16 dies = 128 ranks, Ascend 910C

48 nodes, 384 Ascend 910C cards × 2 dies = 768 ranks

Parallelism TP4 · EP4 · CP1; DP-shard 4; EDP-shard pure FSDP, shard 16 (HyperParallel- 4 (TP=EP=CP=PP=1); reshard after FSDP) fwd/bwd; prefetch depth 1

TP4 · DP32; HSDP 16-shard × 2replica; reshard-after-forward; prefetch depth 2

TP4 · EP16 · DP192 · PP1; non-expert shard 32 × 6 replicas; expert shard 6 × 2 replicas; prefetch depth 2

Workload

seq 4096; GBS 32; MBS 1 and 2; 131,072 tok/step

seq 4096; GBS 32; MBS 2; 131,072 tok/step

seq 4096; GBS 384; MBS 4; 3 gradaccum; 1,572,864 tok/step

seq 4096; GBS 768; MBS 4; 3,145,728 tok/step

Precision

bf16 model dtype; mixed-precision policy off

bf16 parameters; HCCL bf16

bf16 compute; fp32 softmax/logits/gradient reduction

bf16 parameters; fp32 master gradients and reduction

Recompute

off

full, layers 0~26

off (both systems)

full/block, 47 of 51 layers

Optimizer

Muon + AdamW groups

Muon (both systems)

Muon (all-gather strategy) + AdamW groups

Software

torch_npu 2.9.0; CANN 8.5.1 torch 2.9.0+cpu; torch_npu 2.9.0; CANN 8.5.1; Transformers 5.13.0; HyperParallel-FSDP at bada2df4 with the archived working-tree patch

Measurement two mode pairs, each 60 steps; window steps 2~60 excluding profilerexport step 13 (𝑛 = 58 per run)

PyTorch 2.6.0; torch_npu 2.6.0.post5; PyTorch 2.6.0; torch_npu 2.6.0.post5; CANN 8.5.0 CANN 8.5.2

60 steps; window steps 8~60 (𝑛 = 53); 1,000 steps; window steps 100~999 single-step rank-0 profiler (𝑛 = 900); one run per system

the compared configurations differ, the delta is disclosed next to the result. 7.1

uses the same global batch size of 32 and completes 60 steps from the same code and data. We compare production mode (validate_placement=false), which materializes plain local tensors when applying the plan, against validation mode (validate_placement=true), which retains DTensors for layout propagation and boundary checking. We report the mean over steps 2–60 after excluding profiler-export step 13 (𝑛 = 58 per run); the fixed workload contains 131,072 tokens per optimizer step, so throughput is derived from step time.

Experimental Setup

Table 4 fixes the four evaluated configurations; the subsections below add per-experiment detail and disclose baselineside configuration deltas next to each result. All experiments run on Atlas 900 A3 SuperPoD with HCCL collectives. MFU Accounting. All reported MFU figures use one shared accounting rule with coefficients 1 for forward computation and 2 for backward computation: causal attention of sequence length 𝐿 is counted over the 𝐿(𝐿+1)/2 effective query–key positions of the lower triangle, not 𝐿 2 ; MoE layers are counted over the FLOPs of the activated (routed) experts rather than all experts. MFU is reported for experiments C and D only: in experiment C both systems derive MFU from the same per-chip peak denominator (≈353 TFLOP/s, recovered from the logs’ TFLOP/s-to-MFU ratio), so the two values are directly comparable; experiment D quotes its training log’s own MFU accounting. Experiments A and B report step time and shape-derived nominal throughput instead. 7.2

window steps 30~50 (𝑛 = 21); no baseline

Production vs. Validation. Production is faster in both micro-batch regimes (Table 5). At MBS 2, validation increases mean step time from 6.070 to 6.244 s (+2.873%) and reduces throughput by 2.792%. At MBS 1, validation increases step time from 9.823 to 10.506 s (+6.960%) and reduces throughput by 6.508%. Thus, when the smaller micro-batch makes host dispatch and operation submission more prominent, the absolute production–validation gap grows from 0.174 to 0.684 s (3.92×), and the relative step-time gap grows from 2.873% to 6.960%. The profiler provides direct supporting evidence: relative to production, validation adds 23,713 CPU operator events at MBS 2 and 47,431 at MBS 1 while adding no device kernels or collectives. This is the expected signature of validation’s per-operator DTensor and metadata path; production removes that steady-state dispatch path after applying the plan.

Dual-Mode Distributed-Tensor Accuracy and Performance

Protocol. Experiment A (Table 4) trains Qwen3-30B-A3B on the same 16 devices in two matched mode pairs, one with micro-batch size (MBS) 2 and one with MBS 1. Each run 15

Table 5. Production and validation performance at two micro-batch sizes. Parentheses show the validation-mode change relative to production at the same MBS. Throughput is nominal, computed from 131,072 tokens per step. The primary window is steps 2~60 excluding profiler-export step 13 (𝑛 = 58 per run). MBS

Mode

Mean step (s)

Median step (s)

Tokens/s

Peak alloc. (GiB)

Peak HBM (MiB)

2 2

Production Validation

6.070 6.244 (+2.873%)

6.021 6.178

21,593 20,990 (−2.792%)

36.649 36.649

43,683 43,683

1 1

Production Validation

9.823 10.506 (+6.960%)

9.800 10.562

13,344 12,475 (−6.508%)

24.152 24.152

29,911 30,000

7.3

Cross-MBS Comparison. Holding the global batch and parallel mesh fixed, MBS 1 executes twice as many microbatch forward/backward passes per optimizer step as MBS 2. Relative to MBS 2, its mean step time is 61.8% higher in production (9.823 vs. 6.070 s), but 68.3% higher in validation (10.506 vs. 6.244 s). The corresponding throughput losses are 38.2% and 40.6%, respectively. Device activity scales in the same direction: the profiled non-communication kernel count rises from 57,560 to 105,288 and the collective count from 9,554 to 18,854. The smaller MBS reduces framework peak allocation from 36.649 to 24.152 GiB, but exposes more host-side dispatch/submission overhead; validation pays that overhead on every operator, so its performance separates further from production. The clean post-profile window (steps 14~60) reproduces the trend: validation throughput is 2.597% lower at MBS 2 and 6.638% lower at MBS 1.

Topology-Aware FSDP/HSDP Comparison

This subsection reports two whole-system comparisons (experiments B and C in Table 4; Table 6) and one baselinefree scale demonstration (experiment D). We report wholesystem differences only; profiler observations are presented as mechanism evidence consistent with the design rationale of §5.5, not as per-feature attribution. Experiment B: 16-Die Comparison against PyTorch FSDP2. Both sides train Qwen3-30B-A3B from the same checkpoint under a pure-FSDP mesh (shard 16; TP, EP, CP, and PP all equal to 1) with identical model, data, and workload. Over the 53-step steady window, mean step time drops from 3.720 s to 2.614 s (−29.7%), a nominal-throughput gain of 42.3% at the fixed step shape; peak allocated memory falls by 1.03 GB and peak reserved by 2.94 GB; both runs reach loss 1.888 at step 60. A single-step rank-0 profiler trace corroborates the log result (stage time 3,715.1 ms vs. 2,586.7 ms, −30.4%) at identical logical payload (7.670 GB all-gather, 3.817 GB reduce-scatter): exposed communication falls from 916.7 ms to 255.4 ms (−72.1%) and the communication-overlap rate rises from 70.81% to 85.32%. The trace is consistent with the design rationale of §5.5: PyTorch FSDP2’s fused path surrounds its collectives with 451.1 ms of packing/unpacking device work (copy_in/copy_out), and 97 of its 125 all-gathers plus all 49 of its reduce-scatters have buffer lengths not aligned to 512 bytes. Within PyTorch FSDP2’s own all-gathers, the non-aligned group achieves 51.13 GB/s weighted HCCS bandwidth versus 101.96 GB/s for the aligned group (−49.9%), while HyperParallel-FSDP’s non-aligned bytes are negligible (0.0001% of its HCCS traffic). HyperParallel-FSDP instead issues 1,087 all-gathers and 531 reduce-scatters per step, pays ≈81.7 ms of extra per-parameter post-reduction division and an order of magnitude more host-side launch/event time (153.2 ms vs. 9.5 ms of collective-API host self time), and shows higher steady jitter (CV 4.14% vs. 2.39%). Two configuration deltas are disclosed: prefetch depth 1 (HyperParallel-FSDP) vs. 0 (PyTorch), and expandable allocator segments enabled on the PyTorch side only. The comparison is therefore between the two evaluated configurations as whole systems: at this workload, the device-side critical-path gains outweigh the host-side cost of per-parameter dispatch.

Numerical and Device-Work Equivalence. Within each MBS pair, all 60 logged foundation-loss, total-loss, gradientnorm, and learning-rate values match exactly at logging precision. The rank-0 profiles also match within each pair in non-communication kernels, collective types and counts, and per-link transfer volume. Framework peak allocated and reserved memory are identical between modes; sampled HBM is identical at MBS 2 and differs by only 89 MiB at MBS 1. The larger validation gap is therefore consistent with its extra host execution path, rather than a change in accelerator work, framework memory footprint, or training numerics. This 60-step result complements the one-step gradient-equivalence harness (§4.5); it is not a long-horizon convergence study, which §7.3 provides separately. Measurement Scope. Each MBS uses one productionthen-validation pair and includes a single profiled step. Individual traces are sensitive to communication overlap and idle-time variation, but the aggregate and clean postprofile windows agree in both configurations. Therefore, rather than treating the exact percentages as a hardwareindependent causal constant, We use the results to show the measured scaling relationship. A larger mode gap when dispatch/submission work is more prominent. A controlled comparison against PyTorch’s native DTensor runtime remains outside the current evidence set. 16

Table 6. Whole-system fully-sharded results. Experiment B: Qwen3-30B-A3B, pure FSDP on 16 dies, torch_npu 2.9.0, steady window steps 8~60 (𝑛 = 53); baseline is PyTorch FSDP2. Experiment C: 7B-class dense model, 128 ranks, TP4/DP32, PyTorch 2.6.0, steady window steps 100~999 (𝑛 = 900); baseline is Megatron DDP with distributed optimizer, and its memory rows are maxima over the 32 monitored ranks (of 128). Experiment B throughput is nominal, derived from the fixed 131,072-token step shape (the raw token counters read zero in both logs). Profiler rows are single-step rank-0 measurements of experiment B. Disclosed configuration deltas are listed in the text. B: 16-die pure FSDP Metric Mean step time Speedup Throughput MFU Step-time CV Peak allocated HBM Peak reserved HBM

C: 128-rank dense

PyTorch FSDP2 HyperParallel-FSDP Megatron DDP

HyperParallel-FSDP

3.720 s 2.614 s ( − 29.7%) 6,185.2 ± 33.0 ms 4,606.7 ± 21.8 ms ( − 25.5%) 1.00 × 1.42 × 1.00 × 1.343 × 35.23 ktok/s 50.14 ktok/s (+42.3%) 254.30 ktok/s 341.43 ktok/s (+34.3%) not reported 35.65% 48.02% (+12.37 pp) 2.39% 4.14% 0.53% 0.47% 51.035 GB 50.005 GB ( − 1.03 GB) 48.516 GB 41.569 GB ( − 14.3%) 58.145 GB 55.201 GB ( − 2.94 GB) 51.793 GB 49.150 GB ( − 5.1%)

Experiment B, single-step rank-0 profiler: Stage time 3,715.1 ms Exposed communication 916.7 ms Communication overlap rate 70.81% HCCS weighted bandwidth 56.78 GB/s

2,586.7 ms ( − 30.4%) 255.4 ms ( − 72.1%) 85.32% 125.22 GB/s (2.21 × )

Experiment C: 128-Rank Comparison against Megatron DDP. The second comparison uses a different baseline family at larger scale: Megatron DDP with distributed optimizer (42 overlapped gradient buckets over the full 32-way DP group), 128 ranks, TP4/DP32, a 7B-class dense model, 1,000 steps. HyperParallel-FSDP runs as 16-way shard × 2-replica HSDP with reshard-after-forward and prefetch depth 2. Mean step time drops from 6,185.2 ms to 4,606.7 ms (−25.5%, 1.343×), cluster throughput rises from 254.30 to 341.43 ktok/s (+34.3%), and MFU from 35.65% to 48.02% (+12.37 pp, same per-chip peak denominator). Worst-observed peak allocated memory falls from 48.516 GB to 41.569 GB (−14.3%). Disclosed deltas: the HyperParallelFSDP run pads the vocabulary to 153,600 vs. 151,552 (+0.19% parameters, on the HyperParallel-FSDP side), gradientaccumulation fusion is enabled on the Megatron side only, memory is monitored on 32 of 128 ranks, and each system was run once, so cross-run cluster noise is not characterized.

Experiment D: 505B-Parameter Scale Demonstration, No Baseline. A 505B-parameter MoE (51 layers, MLA, 384 experts, top-8) trains on 384 physical Ascend 910C cards (768 ranks) with TP4/EP16/DP192, non-expert state sharded 32-way with 6 replicas and expert state sharded 6-way with 2 replicas. The steady window (steps 30~50, 𝑛 = 21) runs at 7.465 s/step with a 0.75% coefficient of variation, 421.4k tokens/s (36.41B tokens/day), 67.43 TFLOP/s per logical die, and 19.1% MFU under the run’s logging accounting. The breakdown recorded for this configuration attributes 65.1% of step time to compute, 31.4% to exposed communication, and 3.5% to idle; FSDP communication accounts for only 2.9% of step time, while expert-parallel token dispatch (19.5%) and optimizer communication (7.8%) dominate the communication critical path. At this scale the fully-sharded layer’s own communication is a minor cost and the optimization pressure moves to expert-parallel and optimizer traffic. No matched baseline exists for this configuration, so we report trainability, stability, and the achieved breakdown only, with no relative speedup claim.

Long-Horizon Numerical Alignment. Experiment C doubles as the paper’s long-horizon numerical evidence. Over all 1,000 aligned steps, the four MTP-head training losses of the two systems correlate at Pearson 𝑟 > 0.999997 with stepwise mean absolute differences of 0.00203–0.00313; neither run records a NaN, a skipped step, or a loss-scale change; and final-step losses differ by at most 0.0015 (relative difference ≤ 0.034%). Gradient norms correlate at 𝑟 = 0.9478 with alternating-sign differences. The two implementations are not bit-identical because different sharding changes the FP32 reduction order; however, no convergence-speed or stability degradation is observed. These losses cover the four MTP heads of this model only; no downstream evaluation is included.

7.4

Layout-Driven Distributed Muon Comparison

We compare HyperParallel-FSDP’s distributed Muon against the three baseline realizations we integrated and ran: DMuon, MatrixFSDP, and TorchTitan’s FlexShard/DistMuon. All three comparisons share one platform (16 Ascend A3 NPUs with HCCL, PyTorch 2.6.0, CANN 8.5.2, 16-way fullysharded data parallelism with TP=CP=EP=PP=1, sequence length 4,096, global batch 32, micro batch 2, 60 training steps) and align the Muon/AdamW parameter split, the NS iteration count (5), the Muon learning rate (3×10−5 , momentum 0.95, Nesterov), and the AdamW settings across the two sides. Performance is measured from single-step profiler traces (rank 0, step 11, one profile per pair) with 17

Table 7. Distributed-Muon comparisons at 16 NPUs (16-way fully-sharded data parallelism, profiler step 11, rank 0; one profile per pair). Stage is the profiler’s full-step time; equivalent-throughput gains are +5.74%, +19.04%, and +6.60% respectively. The optimizer range is Computing + exposed communication + free time inside the optimizer step; the MatrixFSDP range is an approximate manual timeline selection. Per-comparison configuration deltas are disclosed in the text. Comparison

Model

vs. DMuon

30-layer dense Qwen3-14B- full, 30/30 Base layers 24-layer dense Qwen3-14B- full, 24/24 Base layers 48-layer MoE Qwen3-30B- full, 27/48 A3B layers

vs. MatrixFSDP vs. FlexShard/ DistMuon

HyperParallel- Δ stage FSDP stage

Optimizer range (baseline → HyperParallel-FSDP)

4,974.978 ms

4,705.004 ms

−5.43%

317.6 → 605.2 ms

4,548.031 ms

3,820.726 ms

−15.99%

≈394 → ≈488 ms

2,816.295 ms

2,641.921 ms

−6.19%

834.2 → 527.6 ms

Recompute Baseline stage

the same computing/exposed-communication accounting as §7.3, not from log step time. Table 7 summarizes the three comparisons; the distributed Muon of §6 additionally trains inside experiments C and D above, both of which optimize with Muon.

NS coefficients while HyperParallel-FSDP runs shape-fused batched NS with the asym5 coefficients; both use equivalent RMS-matching update scaling, five BF16 NS iterations, and FP32 master weights. FlexShard/DistMuon. The DistMuon comparison controls the fully-sharded layer itself: both sides run on HyperParallel-FSDP’s FSDP with an identical, forced parameter classification (336 Muon and 195 AdamW parameters of the MoE model), isolating the optimizer as the variable. DistMuon decouples storage from compute layout through a per-step pack → all-to-all → NS → all-to-all → unpack redistribution with a double-buffered transfer stream. Its pipeline is effective because communication overlap is higher than HyperParallel-FSDP’s (90.64% vs. 87.22%) and exposed communication lower, but the redistribution adds more work than the overlap recovers: the optimizer range grows to 834.2 ms against 527.6 ms (−36.75% for HyperParallel-FSDP), and the net stage difference is −174.374 ms (−6.19%). Two disclosures: in this comparison HyperParallel-FSDP’s owner assignment was the simpler per-sub-batch greedy variant, so the baseline fielded the more sophisticated (cross-bucket, NScost-based) planner because the model’s regular, repeated matrix shapes limit the tail-balancing benefit such a planner can buy; and the MoE run confirms that the FlexShard adapter consumed HyperParallel-FSDP’s distributed tensors on the native code path, not a dense-tensor fallback.

DMuon. DMuon assigns each Muon matrix a global owner that holds the full parameter, receives the FP32 gradient by reduction, and runs Newton–Schulz locally; parameter publication and gradient reduction are embedded as owner broadcast/reduce hooks in the forward, recompute, and backward paths. Its optimizer range is therefore shorter than HyperParallel-FSDP’s (317.6 vs. 605.2 ms), but the non-optimizer path is 557.6 ms longer (4,657.4 vs. 4,099.8 ms): under full-layer recompute its 212 dedicated parameters trigger 634 BF16 broadcasts (about three materializations per parameter per step), and its total communication, although lower (2,201.3 vs. 2,307.7 ms), overlaps less (78.55% vs. 85.16%), leaving 472.1 ms exposed against HyperParallel-FSDP’s 342.5 ms. The net is −269.974 ms of stage time (−5.43%). A disclosed secondary delta: DMuon normalizes the NS input in FP32 and iterates in FP16, while HyperParallel-FSDP runs the whole NS path in BF16; both sides keep FP32 optimizer state, and this dtype-path difference does not account for the stage gap. MatrixFSDP. MatrixFSDP encodes whole-matrix ownership into the ZeRO-3 layout itself: one rank holds each Muon matrix, the others hold empty shards, and the backward reduction delivers the full gradient to the owner. Its optimizer range is correspondingly communication-free and ≈94 ms shorter than HyperParallel-FSDP’s, but the owner broadcast/reduce traffic moves into the model path (93.47% of its communication time), where only 43.81% of communication overlaps with compute, which is significantly lower than HyperParallel-FSDP’s 82.25%, leaving 782.6 ms exposed against 268.4 ms. The net is −727.305 ms (−15.99%). This comparison is end-to-end over the two full implementations: the baseline runs per-matrix torch.optim.Muon with fixed

Issues Surfaced in the Baseline Integrations. We report three problems our experiments exposed in the baselines, as observed in the pinned integration trees. (i) DMuon, gradient lifetime under recompute: its custom asynchronous reduce path cleared param.grad before the reduce stream had finished reading the BF16 source gradient; under full-layer recompute the caching allocator can reuse that storage before the enqueued cast and reduction consume it, silently polluting the reduced gradient (observable as grad_norm or loss NaN). We fixed the integration with grad.record_stream(reduce_stream); the reported 18

DMuon numbers include the fix, and the episode is concrete evidence that embedding optimizer communication in model hooks makes the optimizer responsible for cross-stream tensor lifetimes that a unified shard lifecycle manages centrally. (ii) DMuon, memory imbalance: owner-resident full matrices plus FP32 master and momentum state concentrate memory on owners (owner-resident capacities span 603M~1,386M elements, a 2.30× max-to-min ratio); at a 40-layer configuration DMuon fails with an HCCL memory-allocation OOM while HyperParallel-FSDP trains the same configuration. (iii) MatrixFSDP, coverage and allocator footprint: the current integration accepts only strictly two-dimensional parameters (torch.optim.Muon enforces the check at construction and at each step) and its demo gates out non-dense models and any TP/CP/EP/PP > 1, so it cannot run the fused three-dimensional expert parameters of the MoE workload at all; and at an equal live-tensor peak (≈24.9 GB allocated on both sides) its allocator reserved high-water is 55.281 GB against HyperParallel-FSDP’s 31.850 GB (−42.4%). Structurally, both owner-resident baselines couple the optimizer to the FSDP parameter lifecycle, i.e., broadcast/reduce hooks in forward, recompute, and backward, whereas HyperParallel-FSDP keeps the standard shard lifecycle and confines Muon communication to the optimizer step.

per-operator DTensor dispatch overhead and a dual mode in which verification and production share one plan. Data-Parallel Training. ZeRO established optimizerstate, gradient, and parameter sharding [21]; PyTorch FSDP productionized it with per-unit all-gather and reducescatter [27], and HSDP added the replica dimension [18]. Megatron-LM composes tensor, pipeline, and data parallelism manually [22]; MegaScale reports 10,000-accelerator production training [8]; TorchTitan packages the stack declaratively [14]. HyperParallel-FSDP’s contribution is orthogonal to the sharding policy itself: a two-tier realization in which intra-supernode collectives are per-parameter and zero-copy, the cross-supernode reduction is fused without copies, and the backward schedule contains no layer-local wait on the slow tier, and residual work settles once at the end of the backward pass. Distributed Muon. Muon was introduced by Jordan et al. [10] and scaled to LLM pretraining by Moonlight, which distributes it in bucket-based ZeRO-1 fashion [15]; NorMuon improves its second-moment normalization [13]. Among per-parameter realizations, DMuon [2] in the inspected revision integrates an owner-based Muon path with the FSDP2 lifecycle: each matrix is assigned a single owner rank holding the authoritative parameter and optimizer state, and owner publication, gradient reduction, and prefetch run on DMuon’s own lifecycle hooks rather than those of a stock FSDP2 stack. Owner-to-all parameter publication and allto-owner gradient reduction are organized as a two-stage intra-/inter-node hierarchy whose XOR owner-slot layout spreads cross-node contention, pipelined against forward and backward compute and published asynchronously; the owner-side Newton–Schulz iteration runs in Gram space with symmetry-aware kernels, shape-grouped batching, and autotuning; and owner assignment minimizes a measured makespan model. Tensor parallelism is accommodated by a second, nested ownership level detected from the parameters’ distributed-tensor placements. Fused expert tensors are a boundary case of the inspected revision: [𝐸, 𝑀, 𝑁 ] expert weights not selected by DMuon remain on the stock FSDP2 and original-optimizer path, while handing them to the current non-TP DMuon path would treat them as a twodimensional [𝐸, 𝑀𝑁 ]-like matrix rather than orthogonalizing each of the 𝐸 experts separately; we therefore report expert-weight coverage explicitly in the evaluation protocol (§7.4). MatrixFSDP [4] takes a third path under ZeRO-3 parameter sharding: for each 2-D weight one data-parallel rank owns the whole matrix and the remaining ranks hold empty shards, so the ordinary backward reduction delivers the full Muon input to the owner and the optimizer step issues no matrix collective; a global owner planner balances resident bytes and optimizer work, realized through owner-segment point-to-point collectives, owner-buffer pinning, and ownershard checkpoint resharding. Matrices already fragmented

Expert-Weight Coverage. Coverage is reported qualitatively because the compared integrations expose no common per-expert numel accounting. HyperParallel-FSDP treats a fused [𝐸, 𝑀, 𝑁 ] expert weight as 𝐸 independent [𝑀, 𝑁 ] matrices and orthogonalizes each expert separately (§6); the MoE comparison above classified the same 336 parameters as Muon on both sides, so the fused expert weights were exercised on the native path. MatrixFSDP is capability-limited as described above. DMuon was evaluated on the dense model here; per the inspected-revision analysis (§8, Appendix A), fused [𝐸, 𝑀, 𝑁 ] weights not selected by DMuon remain on the stock FSDP2 and original-optimizer path, and the current non-TP DMuon path would treat a selected fused expert weight as a two-dimensional [𝐸, 𝑀𝑁 ]-like matrix rather than orthogonalizing each of the 𝐸 experts separately.

Related Work

Distributed Tensor Abstractions and Auto-Parallelism. GSPMD introduced annotation-driven sharding with compiler propagation for XLA [25]; OneFlow’s SBP formalized placement algebra at the framework level [26]; MindSpore’s auto-parallelism derives strategies by search [6]; and PyTorch’s distributed tensor brought the model to eager execution via operator-level dispatch below the autograd engine [1, 19]. HyperParallel-FSDP shares the declarative surface but differs in the interception layer: by placing sharding semantics above autograd, it supports a precompiled-boundary execution with zero steady-state 19

by tensor parallelism are excluded from owner placement and handled by the surrounding TP path. Canzona [24] decouples logical optimizer assignment from physical parameter distribution for matrix-based optimizers generally (Muon, Shampoo, SOAP): on the data-parallel plane an 𝛼-balanced static partitioning assigns whole parameters to ranks so the optimizer step is communication-free, and on the tensorparallel plane an asynchronous micro-group pipeline batches fragmented updates behind compute, which is viable because tensor parallelism typically resides in the intra-node highbandwidth domain, in contrast to the inter-node data-parallel plane. TorchTitan’s FlexShard instead gives the optimizer a second layout system: per-parameter compute layouts are declared next to the DTensor storage layouts, and a general reshard planner lowers the storage-to-compute transition onto packed all-to-all collectives [14, 17]. FlexShard’s abstraction is general across optimizers; we instead trade that generality for an engine specialized to Muon’s fixed structure, keeping the algorithm on local tensors and the communication on hardware-friendly collectives. Relative to these systems (§6.1, Table 3; versions as inspected, runtime revisions pinned in §7.4), HyperParallel-FSDP differs along three axes: communication domains and the two-level (shard and replica) deduplication are derived from the distribution semantics of general multi-dimensional meshes rather than assigned per parallel strategy or over a fixed data-parallel mesh; replica dedup domains are aligned to machine boundaries, which is a configuration premise of the current realization (§6.5), and HyperParallel-FSDP performs deterministic recomputation across subgroups, so that under the canonical mesh-to-topology mapping (Figure 1) the evaluated configuration sends no optimizer traffic over the slow tier; and parameters whose matrix plane is not sharded (canonically expert-axis-sharded MoE weights) require no Muon-specific optimizer communication over the expert-sharding axis.

Atlas 900 A3 SuperPoD from 16 dies to 384 physical cards, the two tensor-runtime modes execute identical device work and loss trajectories over the evaluated horizon (§7.2); the fully-sharded layer reduces mean step time by 29.7% against the evaluated PyTorch FSDP2 configuration at 16 dies and by 25.5% against Megatron DDP with distributed optimizer at 128 ranks, with per-step losses tracking the baseline at Pearson 𝑟 > 0.999997 over 1,000 steps; and a 505B-parameter MoE trains at 421k tokens/s on 768 ranks with FSDP communication at 2.9% of step time (§7.3); and the layout-driven distributed Muon reduces profiler step time by 5.4–16.0% against DMuon, MatrixFSDP, and TorchTitan FlexShard/DistMuon integrations at 16 NPUs, where the owner-resident baselines additionally exhibit a cross-rank memory imbalance (OOM at a 40-layer configuration that HyperParallel-FSDP trains) and a recompute-exposed gradient-lifetime bug (§7.4). Several limitations remain. Pipeline parallelism, though available in the runtime, is not yet integrated into the declarative sharding plan; context parallelism is limited to the all-gather family (no ring attention); validation mode at LLM scale needs a lightweight reference runner; tensorparallel all-gather/reduce-scatter pipelining is future work; and expert-parallel all-to-all scheduling is being addressed separately. HyperParallel-FSDP is open source [7].

References [1] Jason Ansel, Edward Yang, Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesensky, Bin Bao, Peter Bell, David Berard, Evgeni Burovski, et al. 2024. PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2 (ASPLOS ’24). ACM, 929–947. doi:10.1145/3620665.3640366 [2] Vincent Chen, Starrick Liu, Regis Cheng, Dance Yang, Shalfun Li, Ryan Yu, Lucy Liang, Hang Su, Roy Gan, Hao Wang, and Qian Wang. 2026. DMuon: Efficient Distributed Muon Training with Near-Adam Overhead. arXiv preprint arXiv:2606.27153 (2026). [3] DeepSeek-AI. 2026. DeepSeek-V4: Towards Highly Efficient MillionToken Context Intelligence. arXiv preprint arXiv:2606.19348 (2026). [4] Ming Gao, Yanwu Xu, and Hao Zhang. 2026. MatrixFSDP: communication-free matrix optimizers under ZeRO-3 parameter sharding. arXiv preprint arXiv:2607.05895 (2026). [5] GLM-4.5 Team. 2025. GLM-4.5: Agentic, Reasoning, and Coding (ARC) Foundation Models. arXiv preprint arXiv:2508.06471 (2025). [6] Huawei. 2020. MindSpore: A Flexible and Efficient Deep Learning Framework. https://www.mindspore.cn. [7] HyperParallel Team. 2026. hyper-parallel: Distributed parallelism acceleration library for supernodes. https://atomgit.com/mindspore/ hyper-parallel. [8] Ziheng Jiang, Haibin Lin, Yinmin Zhong, Qi Huang, Yangrui Chen, Zhi Zhang, Yanghua Peng, Xiang Li, Cong Xie, et al. 2024. MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs. NSDI (2024).

Supernode-Scale Systems. CloudMatrix384 reports production serving on a 384-die supernode [28]. The companion HyperParallel-MoE system tackles MoE training on the same hardware class through multi-core interleaved scheduling [9]; HyperParallel-FSDP is complementary, addressing sharding-plan execution and validation, data-parallel communication, and the optimizer.

Conclusion

HyperParallel-FSDP makes declarative parallelization semantics explicit, verifiable, and free of steady-state dispatch: by intercepting above autograd, one sharding plan drives both a production mode with no per-operator DTensor dispatch and a validation mode with fail-fast checking, closed by gradient-equivalence testing. On top of this sharding-plan execution and validation layer, topology-aware fully-sharded communication and a layout-driven distributed Muon improve utilization of the two-tier supernode network. On

[9] Zewen Jin, Congkun Ai, Guangpeng Zhang, Hanbo Zhang, Haoran Wang, Shihan Xiao, Da Lei, Xuefeng Jin, Teng Su, and Cheng Li. 2026. HyperParallel-MoE: Multi-Core Interleaved Scheduling for Fast MoE Training on Ascend NPUs. arXiv preprint arXiv:2605.23764 (2026). 20

[10] Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. 2024. Muon: An optimizer for hidden layers in neural networks. https://github.com/KellerJordan/ Muon. [11] Kimi Team. 2025. Kimi K2: Open Agentic Intelligence. arXiv preprint arXiv:2507.20534 (2025).

distributed.fsdp.fully_shard.md; accessed August 2026. [21] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. 2020. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. In SC ’20: International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 1–16. [22] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. 2019. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. In arXiv preprint arXiv:1909.08053. [23] Mo Sun, Zihan Yang, Changyue Liao, Yingtao Li, Jie Zhang, Kaiqi Chen, Fei Wu, and Zeke Wang. 2026. Disdp: Disaggregating Compute, Network, and Storage for Model-Sharded Data-Parallel Training. In 2026 ACM/IEEE 53rd Annual International Symposium on Computer Architecture (ISCA). IEEE, 2442–2458. [24] Liangyu Wang, Siqi Zhang, Junjie Wang, Yiming Dong, Bo Zheng, Zihan Qiu, Shengkun Tang, Di Wang, Rui Men, and Dayiheng Liu. 2026. Canzona: A Unified, Asynchronous, and Load-Balanced Framework for Distributed Matrix-based Optimizers. arXiv preprint arXiv:2602.06079 (2026). [25] Yuanzhong Xu, HyoukJoong Lee, Dehao Chen, Blake Hechtman, Yanping Huang, Rahul Joshi, Maxim Krikun, Dmitry Lepikhin, Andy Ly, et al. 2021. GSPMD: General and Scalable Parallelization for ML Computation Graphs. arXiv preprint arXiv:2105.04663 (2021). [26] Jinhui Yuan, Xinqi Li, Cheng Cheng, Juncheng Liu, Ran Guo, Shenghang Cai, Chi Yao, Fei Yang, Xiaodong Yi, Chuan Wu, Haoran Zhang, and Jie Zhao. 2021. OneFlow: Redesign the Distributed Deep Learning Framework from Scratch. arXiv preprint arXiv:2110.15032 (2021). [27] Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, Alban Desmaison, et al. 2023. PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel. Proc. VLDB Endow. 16, 12 (2023), 3848– 3860. [28] Pengfei Zuo, Huimin Lin, Junbo Deng, Nan Zou, Xingkun Yang, Yingyu Diao, Weifeng Gao, Ke Xu, Zhangyu Chen, Shirui Lu, et al. 2025. Serving Large Language Models on Huawei CloudMatrix384. arXiv preprint arXiv:2506.12708 (2025).

[12] Kimi Team. 2026. Kimi K3: Open Frontier Intelligence. arXiv preprint arXiv:2607.24653 (2026). [13] Zichong Li, Liming Liu, Chen Liang, Weizhu Chen, and Tuo Zhao. 2025. NorMuon: Making Muon more efficient and scalable. arXiv preprint arXiv:2510.05491 (2025). [14] Wanchao Liang, Tianyu Liu, Less Wright, Will Constable, Andrew Gu, Chien-Chin Huang, Iris Zhang, Wei Feng, Howard Huang, Junjie Wang, Sanket Purandare, Gokul Nadathur, and Stratos Idreos. 2025. TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training. In The Thirteenth International Conference on Learning Representations (ICLR). https://openreview.net/forum?id= SFN6Wm7YBI. [15] Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, et al. 2025. Muon is Scalable for LLM Training. arXiv preprint arXiv:2502.16982 (2025). [16] NVIDIA. 2025. NeMo AutoModel: DTensor-native SPMD library for scalable and efficient training. https://github.com/NVIDIA-NeMo/ Automodel. [17] PyTorch Team. 2026. FlexShard: storage and compute layouts for distributed optimizers (TorchTitan). https: //github.com/pytorch/torchtitan/blob/496b11d43/torchtitan/ distributed/flex_shard/README.md; commit 496b11d43, matching the revision pinned in the evaluation; accessed August 2026. [18] PyTorch Team. 2026. Hybrid sharded data parallel (HSDP) in PyTorch FSDP2. https://docs.pytorch.org/docs/2.9/distributed.fsdp.fully_shard. html; accessed September 2026. [19] PyTorch Team. 2026. PyTorch DTensor: API and operator dispatcher. https://github.com/pytorch/pytorch/tree/v2.9.0/torch/ distributed/tensor; dispatch path is version-dependent. [20] PyTorch Team. 2026. PyTorch FSDP2: fully_shard API documentation. https://github.com/pytorch/pytorch/blob/v2.9.0/docs/source/

Full Design-Contrast Matrix for Distributed Muon

Table 8 extends Table 3 to all eleven dimensions we tracked while comparing the systems, including the five dimensions (programming view, scheduling, fused-expert handling, model adaptation, extensibility) that the main text omits for space. Statements about other systems describe the versions as inspected, with runtime revisions pinned in §7.4; “not described” marks capabilities we did not find in those revisions, not a claim of permanent absence. Table 8. Full design contrast of distributed Muon realizations (versions as inspected; runtime revisions pinned in §7.4): HyperParallel-FSDP’s Muon-specialized static plan against DMuon [2], MatrixFSDP [4], Canzona [24], and TorchTitan’s general FlexShard/DistMuon substrate [14, 17]. HyperParallel-FSDP Muon

DMuon

MatrixFSDP

Canzona

FlexShard DistMuon

Programming Local tensors; internal static view global plan

FSDP2-lifecycle-integrated, owner-based Muon path with its own owner-publication, gradient-reduction, and prefetch lifecycle

ZeRO-3 parameter shards; owner placement invisible to the algorithm

Conventional optimizer interface; logical task assignment decoupled from physical distribution

Storage + compute layouts

Domain Derivation

Communication domains and dedup derived from layout semantics (matrix-plane predicate)

Owner assignment over the data-parallel mesh; TP detected from placements, nested ownership

Owner-shaped ZeRO-3 shards: one rank owns each whole 2-D matrix, others hold empty shards

Logical optimizer assignReshard planner lowers userment decoupled from phys- declared storage-to-compute ical parameter distribution; layouts static whole-parameter partition (DP)

State Dedup Scope

Shard + replica two-level; momentum stays sharded; replica state reduced to 1/𝑅𝑠 within topology-aligned subgroups

One authoritative owner per matrix (parameter + momentum); non-owners keep zero-size placeholders

Whole matrix and optimizer state on the owner rank; empty shards elsewhere; ZeRO-3-scale memory

Whole-parameter optimizer state at the statically assigned rank

Owner-style state on the compute layout (Owned, BlockShard)

CrossTopology Strategy

Replica subgroups from the local device count (contiguous rank blocks, machine-aligned under the standard rank mapping; a configuration premise, §6.5); deterministic recomputation across subgroups; no slow-tier optimizer traffic under the canonical mapping (Fig. 1)

Two-stage intra-/inter-node hierarchy; XOR owner-slot layout spreads contention; spans both tiers

Optimizer step local; the backward reduction delivering the full gradient spans the DP mesh

DP optimizer step zerocommunication; TP reconstruction confined to the intra-node high-bandwidth domain

Packed all-to-all; topology expressed through the user-chosen mesh and compute layouts; no automatic machine-boundary policy described

Batching / Load Balance

Shape-fused batched NS; greedy owner balancing; ≤1 batch in flight

Shape-grouped batched GramNS with SYRK kernels and autotuning; measured makespan (MILP) owner assignment

Global owner planner balancing resident bytes and optimizer work (greedy / scope-greedy / cost-aware)

𝛼 -balanced greedy LPT (DP); micro-group balanced scheduling with greedy rollback (TP)

BucketConfig groups and orders parameters for packed redistribution and communication-compute overlap

Scheduling

Static plan replayed unchanged per step

Runtime hook pipeline: lookahead publication, asynchronous publish

Static global owner plan; block-local execution

Precomputed static partition Layout lowering + reshard map; asynchronous task planner execution with computecompute overlap

Dominant Collective

Fused all-gather + relay broadcast

Two-stage owner-to-all broad- None in the optimizer DP: none; TP: fused all-to-all Packed all-to-all cast / all-to-owner reduce step (the backward reduce- over micro-groups scatter lands the input on the owner); ownersegment P2P

Mesh Coverage

General multi-dimensional FSDP/HSDP data-parallel mesh meshes; no Muon-specific + nested TP ownership optimizer communication over the expert-sharding axis when the matrix plane is unsharded

ZeRO-3 data parallelism; TP-fragmented matrices excluded from owner placement

ZeRO-1 DP + TP; supports Muon, Shampoo, SOAP

Layouts on named DeviceMesh axes; DistMuon constrains flat matrix-batch compute to BlockShard on at most one non-unit mesh axis

Fused-Expert Handling

Fused [𝐸, 𝑀, 𝑁 ] treated as 𝐸 independent [𝑀, 𝑁 ] matrices; per-expert NS; never flattened to [𝐸, 𝑀𝑁 ]

In the inspected revision: unselected fused [𝐸, 𝑀, 𝑁 ] weights stay on the stock FSDP2 + original-optimizer path; the current non-TP path would treat them as a 2-D [𝐸, 𝑀𝑁 ] -like matrix, not 𝐸 per-expert [𝑀, 𝑁 ] matrices

Not described

Native 3-D Shard(0) batches as listed above

Model Adapta- Split/cat/restore in the model tion adapter (fused QKV, MLA, MoE FC1)

Per-layer slices of the host stack; no model-specific transform layer

Non-matrix tensors packed Parameter-granularity into tail owners and left on atomicity; no model-specific AdamW transform layer described

BlockShard equal blocks, Owned whole-matrix, or native 3-D Shard(0) batches

Extensibility

Ownership-strategy plug-in; Muon-specific

Matrix optimizers of the Unified across matrix-based Muon family under ZeRO- optimizers 3

General across optimizers

Per-stage hooks within Muon

Record · ID 1006819 · SHA-256 9b4cbde2b3d39d6b
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.