SS-HSDP: Efficient and Flexible Sharded Data Parallelism via Runtime-Update Decoupling

SS-HSDP: Efficient and Flexible Sharded Data Parallelism via Runtime–Update Decoupling

arXiv:2611.16467v1 [cs.DC] 12 Nov 2026

Chendong Xiang∗ , Pengle Zhang∗ , Jinliang Li , Zichen Liang Jintao Zhang , Fan Bao‡ , Jun Zhu‡ , Jianfei Chen‡ Tsinghua University, Shengshu Technology https://www.vidu.com/

Abstract Sharded data parallelism enables memory-efficient training of large models. While efficient training kernels favor packed, kernel-specific data layouts, optimizers may require variant data structures as optimization units. Existing sharded data-parallel frameworks couple these two choices, limiting runtime efficiency and optimization flexibility. To address this, we present SS-HSDP, a training framework that separates runtime parameters from update parameters, allowing computation layouts and optimization units to be organized independently. Through shared data views and hierarchical sharding, SS-HSDP supports this separation while avoiding redundant data movement. This design supports a broad range of optimizers, including matrix-based methods such as Muon, and performance optimizations such as fully zero-copy execution. Experiments on video and language model workloads with MXFP8 training show that SS-HSDP achieves higher Model FLOPs Utilization (MFU) than FSDP2, veScale-FSDP, and Megatron-FSDP. At 128 GPUs, MFU improves by up to 21.7% over the fastest evaluated baseline, with comparable or lower peak memory usage.

1

Introduction

ZeRO [32] enables training large models on memorylimited devices by partitioning weights, gradients, and optimizer states across data-parallel ranks. Under ZeRO3, weights are typically all-gathered before each layer’s computation, and gradients are reduce-scattered after its backward computation. These collectives can overlap with computation in other layers, potentially hiding most of the communication cost in compute-bound workloads. With relatively modest changes to model code, frameworks such as DeepSpeed [33] and PyTorch FSDP2 [31, 47] make this approach practical for largescale training. This approach has been successfully used to train both language models [14, 15] and video generation models [22, 38].

SS-HSDP

Figure 1 Parameter separation in SS-HSDP, and MiMo MXFP8 training performance with AdamW and framework capabilities.

Recent advances in model architectures, low-precision training and optimization impose new requirements on both runtime computation efficiency and optimizer flexibility. For MoE models [8, 12], redundant movement of their large weights and gradients can introduce substantial overhead. Low-precision weight formats such as MXFP [28] represent weights using quantized values and scale metadata, whose layout must match the requirements of high-performance kernels [9, 23]. Meanwhile, matrix-

based optimizers such as Muon [17] operate on matrices that may be submatrices of the weights used for computation. For example, when applying Muon per attention head [18], each head is updated independently, whereas forward and backward computation can use one packed QKV projection, with its values and scales arranged for the target GEMM kernel.

*Core contributors, co-first authorship. ‡Advisors.

1

These requirements expose a mismatch between how parameters are computed and how they are updated. Runtime computation requires parameter data to be arranged for optimized kernels, packed across multiple parameters and stored in kernel-specific layouts, whereas optimizers require the same data to be accessed in the units they optimize, such as individual attention heads. Tying update units to the parameter organization used for computation constrains these choices: supporting different update units may require reorganizing data or customizing the optimizer.

enable efficient computation and flexible optimizer updates without redundant data movement. • Implement A clean, high-performance implementation that supports diverse training scenarios and achieves higher MFU than FSDP2, veScale-FSDP, and Megatron-FSDP. We will release the code in the future.

Existing frameworks address only part of this mismatch. Megatron-FSDP [24] optimizes the runtime data path with flattened buffers, persistent low-precision weights, and direct gradient writes, but does not support structured optimizer updates. veScale-FSDP [41] aligns ragged shards with the update and quantization structure to support distributed Muon updates and block-scaled quantization, but does not pack or transform parameters into kernel-specific layouts. Our goal is to satisfy both requirements without trading one for the other.

2

Background and Motivation

2.1

Notation and ZeRO Training

Tensors and shards A tensor is a multidimensional array of numerical values; a matrix is a two-dimensional tensor. We use lowercase letters for tensors and corresponding uppercase letters for tensor lists, e.g., W = [w0 , . . . , wk ]. A tensor list is a collection of tensors with custom layouts. A logical tensor denotes the complete tensor, regardless of whether all, part, or none of its data are stored locally. Sharding partitions a tensor or tensor list across a group of ranks, each holding a local shard. The number of ranks in this sharding group is its sharding degree. When shards are replicated, ranks holding replicas of the same shard form a replica group. Superscripts identify the sharding level of a local shard, as defined in the ZeRO sharding paragraph below; the same notation applies to tensors and tensor lists. Table 1 summarizes the notation.

We present SS-HSDP, a framework that decouples runtime parameters for efficient computation from update parameters for flexible optimization. Figure 1 illustrates this separation and its performance benefits. Each update parameter defines a unit of optimization and its associated operations, including optimizer updates and weight casting or quantization. The weights of one or more update parameters are packed or transformed into a runtime parameter’s weight in the representation required by the computation kernel. With careful memory management and hierarchical sharding design, this separation needs no extra memory operations. Rather, it preserves update flexibility while enabling computation- and communicationfriendly data layouts that minimize redundant copies and exploit optimized kernels and communication operations.

Parameters and updates Following Megatron’s naming convention [25, 34], a parameter P has several associated tensors for computation and optimization. In mixedprecision training, a parameter has a weight w stored in a low-precision format such as BF16 or FP8, and an FP32 main weight wmain for optimization [21]. Gradients produced by backward computation are reduced across data-parallel ranks and accumulated over microbatches into a main gradient g. The optimizer uses g to update wmain and its optimizer states, such as momentum m. We consider FP32 main gradients and optimizer states unless otherwise stated. We denote the optimizer update by F and weight preparation by G:

Building on this design, SS-HSDP further reduces memory allocations and copies. On supported topologies, it transfers data using copy engines within nodes and hostdriven remote memory access (RMA) between nodes, reducing contention for streaming multiprocessors (SMs). As a result, SS-HSDP achieves up to 1.33× speedup and 27.3% memory saving over PyTorch FSDP2 on 128 GPUs. In summary, our contributions are:

F

(wmain , m) ←− (wmain , g, m), (w, s) = G(wmain ).

• Identify We identify a mismatch between kernelspecific runtime layouts and optimizer update units, and show how coupling them limits runtime efficiency and optimizer flexibility.

Depending on the optimizer, additional or different state tensors may be used. G casts or quantizes the main weight, producing optional scale metadata s for quantized formats.

• Design We design a unified framework that separates runtime parameters from update parameters, with memory management and hierarchical sharding that

ZeRO sharding ZeRO stage 1 shards main weights and optimizer states; stage 2 additionally shards main gradients; stage 3 further shards computation weights [32]. We 2

SS-HSDP

Figure 2 FSDP2-style mixed-precision training with two sharding ranks (rank 0 shown). Colored segments track one parameter’s weights, gradients, and optimizer states; white segments belong to other parameters. Solid arrows show the reference dataflow, while numbered dashed paths highlight alternative implementations. AG and RS denote all-gather and reduce-scatter, respectively. Compute stream … Compute stream … … … denote local shards at the weight, gradient-accumulation, and optimizer-update levels byAGsuperscripts z3, z2, and z1, stream AG stream with sharding degrees sw , sg , and so , respectively. For RS stream RS stream example, wz3 and W z3 denote local shards of a weight tenTime Time sor and a weight tensor list. Conventional ZeRO-3 uses the same sharding group for all three levels. Hybrid sharded Figure 3 ZeRO-3 Layer-wise(L) overlap. F and B denote forward and backward; AG and RS denote all-gather and reducedata parallelism (HSDP) uses smaller sharding groups and scatter. replicates shards across groups to limit repeated weight all-gathers. Outer ZeRO-1 further shards main weights and optimizer states across replica groups [24], reducand backward computation, G casts (wmain )z3 to BF16 ing persistent memory without enlarging the groups for wz3 . Copy-in packs these shards into the all-gather input repeated all-gathers. With gradients accumulated over W z3 . All-gather produces W in an interleaved layout: microbatches, outer communication is needed once per data are ordered by rank, leaving each parameter’s shards optimizer step. separated by other parameters’ data. Copy-out assembles each contiguous weight w for computation. Backward Training execution We call forward and backward comproduces per-parameter gradients, which are converted to putation and their associated communication runtime. UnFP32 and packed by destination rank into a reduce-scatter der ZeRO-3, each layer all-gathers weights before compubuffer. HSDP additionally all-reduces the results across tation and reduce-scatters gradients after backward. These replica groups. The reduced shards initialize g z2 on the collectives can overlap with computation in other layers. first microbatch and are added to it thereafter. After all Figure 3 illustrates this layer-wise overlap. After gradimicrobatches, F updates the local main-weight shards ents from all microbatches have been accumulated, the (wmain )z3 . update phase applies F to parameters across the model and prepares weights through G for the next training step. As indicated by the circled numbers in Figure 2, FSDP2’s default communication path incurs repeated weight 2 buffer copies ⃝ 3 ⃝, 4 and separate gradient accucasts ⃝, 2.2 FSDP2 Dataflow 5 Its parameter-wise sharding also requires admulation ⃝. Figure 2 shows PyTorch FSDP2’s BF16 mixed-precision 6 and ditional handling for structured optimizer updates ⃝ dataflow with FP32 gradient reduction and accumula1 while packed or transformed kernel-specific layouts ⃝, tion [31]. FSDP2 shards each parameter along dimension HSDP replicates main weights and optimizer states rather 0 by default and represents its shards as DTensors [30]. 7 We examine these limitathan further sharding them ⃝. Hybrid sharded data parallelism (HSDP) additionally tions below. replicates these shards in replica groups. Before forward

3

Table 1 Notation. Symbol P, P i

′

Meaning

w, s wmain , g m u, v w′ , s′ , g ′ W, S, G, M, U, V W ′ , S ′ , G′

Update and runtime parameters Index of an update parameter inside a runtime parameter Weight and quantization scales Main weight and main gradient Optimizer state, e.g., momentum Update temporaries Runtime weight, scales, and gradient Lists of corresponding tensors Lists of corresponding runtime tensors

F fj G Tw Tg

Optimizer update The jth operation in F Weight casting or quantization Weight transform (wi , si )i → (w′ , s′ ) Gradient transform g ′ → (gi )i

z3, z2, z1 sw , sg , so

ZeRO-3/2/1 shard superscripts Sharding degrees at z3, z2, z1

2.3

Existing frameworks address these requirements in different ways. Megatron-FSDP [24] uses Transformer Engine for kernel-specific low-precision weight transforma1 and retains the prepared weight shards across tions ⃝ 2 Its flattened buffers avoid weight layout microbatches ⃝. 3 while compatible kernels write gradients dicopies ⃝, 4 It also supports outer ZeRO-1 rectly into these buffers ⃝. 7 but does not support structured optimizer upsharding ⃝, 6 veScale-FSDP [41] uses structure-aware ragged dates ⃝. sharding to accommodate quantization blocks and dis6 and DBuffer to avoid weight tributed Muon updates ⃝, 3 However, neither framework decouples layout copies ⃝. update units from runtime layouts, limiting the combination of update flexibility and runtime efficiency. SS-HSDP decouples update units from runtime layouts, 6 while packing supporting structured optimizer updates ⃝ 1 and transforming weights for efficient computation ⃝. Shared buffers and direct gradient writes avoid redundant 3 ⃝. 4 Prepared weight shards weight and gradient copies ⃝ are reused across microbatches to avoid repeated conver2 while further sharding main weights and optision ⃝, mizer states across replica groups can offset the additional 7 Custom ZeroSM reduce-scatter accumemory cost ⃝. mulates directly into g without a separate post-reduction 5 Section 3 presents the parameter separation, memadd ⃝. ory management, and sharding design behind these optimizations.

Challenges and Analysis

In conventional BF16 training with AdamW, w, wmain , g, and the optimizer states have matching shapes, while both F and G operate elementwise. These tensors can therefore be sharded identically at arbitrary element boundaries, without considering matrix or block structure. Structured updates and block quantization change these requirements. Muon is a matrix-based optimizer [17]; applying it per attention head requires updating each head matrix independently [18], rather than treating a packed projection 6 as one update unit ⃝. Similarly, MXFP quantization in G computes a shared scale for a block of weight values [28]. These operations require complete matrices, heads, or quantization blocks; when sharded, metadata must identify the fragments needed for communication.

3

Design

In this section, we present the core design of SS-HSDP, including the separation of update and runtime parameters, data trees and shard chains for managing flattened buffers and hierarchical sharding, and parameter grouping.

3.1

Update and Runtime Parameters

To support runtime efficiency and update flexibility, SSHSDP divides the parameter into an update parameter and a runtime parameter, denoted by P and P ′ , respectively. An update parameter defines an independent unit of optimization, with the tensors and weight preparation operation G defined in Section 2.1. We express its optimizer update as an ordered sequence F = [f0 , f1 , . . . , fn ], with optional temporary tensors such as u and v. Each operation specifies the tensors it reads and writes. Although it contains the low-precision representation (w, s), an update parameter is not directly used by forward or backward kernels.

Runtime computation instead favors packing Q, K, and V projections into one GEMM, with weights and scales 1 Preparing arranged in kernel-specific layouts [9, 23] ⃝. these representations for every microbatch can be costly. Reusing the prepared weight shards avoids this repeated 2 but increases persistent memory usage. This work ⃝, additional memory cost can be offset by outer ZeRO7 For communication, compatible contiguous buffer 1 ⃝. layouts let collectives and computation kernels share stor3 ⃝. 4 However, age, avoiding weight and gradient copies ⃝ evenly sharding these buffers can split matrices, heads, or quantization blocks required by updates. The challenge here is to preserve runtime efficiency without restricting update flexibility or introducing redundant data movement.

A runtime parameter P ′ contains one or more update parameters P0 , . . . , Pk whose weight, scale, and maingradient tensors are denoted by wi , si , and gi , respectively. P ′ holds the runtime weight w′ , optional scales 4

Data tree

W W′

wL Linear

w′L

w′Q

Linear

Projection

W′ Full

Root (owns storage)

w0

w1

w2

Head0

Head1

Head2

Linear

Head0

Resolve root + offset

Shard chain W Linear w′L

w′Q

ranks 0, 2

w0

Full

Head0 Head1 Head2

2

Head0 0, 1, 2, 3

w0z3

ZeRO-3

Wz2 Linear ranks 0, 2

Linear

w0z2

ZeRO-2

view

Head0 Head1 Head2 rank 1

rank 3

Runtime

Each optimizer step

𝒯g: unpack (view-only)

Wmain ZeRO-1

w0z1

Linear

∅ rank 0

shard

∅ 2

1

rank 0

Linear

Head0 Head1 Head2 rank 2

fi0

rank 1

u₂

6 Structured update Muon · Head0

rank 3

fi1

Newton–Schulz

FP32

Outer RS

G ZeRO-1 rank 0

rank 1

fi2

Local main-weight update

v₂

rank 3

Local momentum g, m → u (BF16) rank 1

u₁ P2P

full Head0 · rank 2

Head0 Head1 Head2 rank 2

7

after all microbatches

rank 2

Quantize main weight

0, 2 1, 3

ranks 1, 3

FP32 accumulation across microbatches

Each microbatch

rank 3

P2P if needed

ranks 1, 3

rank 2

rank 1

Head2

Head0 Head1 Head2

ranks 0, 2

0, 2 1, 3

ZeRO-1

rank 0

rank 2

Head0 Head1 Head2

Wz1

Linear

Update

Head0 Head1 Head2 FP8

Head1

ZeroSM Comm + accumulation

G′ ZeRO-2

K-major scale layout

ranks 1, 3 shared view

5

UE8M0

Head0 Head1 Head2

ranks 0, 2

Reuse across microbatches 1 𝒯w: Pack / transform + outer all-gather

rank 0

Head0

UE8M0

ranks 1, 3

Linear

Wz3 Linear

Head2

W ZeRO-1

W′

Projection

write

DeepGEMM scale layout FP8 Projection

Linear

Linear

4

forward / backward Direct grad

ZeroSM all-gather

3

W′ ZeRO-3

Shared-storage views

Head1

G′ Full

MXFP8 GEMM

Projection

FP32

P2P return

v₁

3

Weight

Gradient

Optimizer scratch

Other parameters

Rank cuts

Figure 4 SS-HSDP data structures and execution flow. Left: data-tree views and the Head0 shard chain. Right: packed runtime computation and structured updates. The example uses four ranks, with so = 4, sg = sw = 2.

s′ , and the gradient g ′ produced by backward computation. The transform operations Tw and Tg connect the two representations:

3.2

Data Tree and Shard Chain

SS-HSDP stores each tensor list in a contiguous buffer, grouping tensors with the same device and dtype for memory management and communication.

(w′ , s′ ) = Tw (w0 , s0 , . . . , wk , sk ), (g0 , . . . , gk ) = Tg (g ′ ).

Data Tree A data tree allows update and runtime tensors to share storage despite having different shapes and boundaries. Each tensor or tensor list is represented by a data node. A parent–child link defines a view at a byte offset within the parent; the child’s shape and dtype determine the view’s size. Only the root owns the allocation, and each child accesses this storage through its accumulated offset.

Tw prepares the layout required by computation kernels, while Tg maps the runtime gradient back to the main gradients for optimizer updates. EG1 We use MXFP8 Q projection with per-head Muon updates as an example. Each Q head has an update parameter Pi , while the Q projection is represented by one runtime parameter P ′ . For each head, Fi = [fi0 , fi1 , fi2 ], fi0 updates mi using gi and writes a BF16 intermediate to ui . fi1 applies Newton–Schulz to ui and writes the FP32 update direction to a separate scratch tensor vi . Then fi2 updates wimain using vi . Gi uses 32 × 32 block quantization in a k-major layout to produce FP8 wi and UE8M0 si . Then Tw packs the head weights wi and relayout the scales si into the layout required by DeepGEMM’s MXFP8 kernels for the whole Q projection. This avoids separate per-head GEMM launches. Tg maps g ′ back to the per-head gi .

In EG1, the head weights wi are consecutive views of W , while the runtime weight w′ views their combined region. The heads can therefore be updated independently and used together as one projection without moving their weight data. The weight-packing part of Tw requires neither a copy nor an additional allocation, since the packed weight is merely a view of the root allocation, as shown on the left of Figure 4; scale relayout for s′ is still performed to satisfy the kernel’s layout requirements. Shard Chain SS-HSDP partitions each flattened tensor list into contiguous chunks and further partitions each local chunk at subsequent levels, recording these partitions as a shard chain:

Packing and unpacking may introduce redundant copies and allocations when the tensors could instead share storage. The update operations must also remain valid when these buffers are sharded. We use data trees and shard chains to address these requirements.

W −→ W z3 −→ W z2 −→ W z1 . 5

When adjacent levels use the same sharding degree, their local shards cover the same data and can share storage through the data tree. A shard-chain link alone does not create a parent–child link in the data tree.

trees and shard chains. A runtime group schedules operations across one or more atom groups, handling coalesced communication, gradient accumulation at ZeRO-2, and the transforms Tw and Tg . These transforms can be configured at different sharding levels. An update group schedules optimizer updates F and weight preparation G at ZeRO-1, batching compatible operations across its atom groups. This organization supports layer-wise runtime execution and model-wide batching of compatible updates.

Each member tensor inherits this chain: its local fragment is the intersection of its range with the local buffer range, or empty if they do not overlap. wi −→ wiz3 −→ wiz2 −→ wiz1 . Each link records the sharding group and partition, mapping local fragments to their positions in the original tensor. The same construction applies to all tensor fields. Together with rank ownership, these mappings locate the fragments required for distributed updates.

3.4

We now follow EG1 through the workflow in Figure 4. The example uses four ranks, with ZeRO-1 degree four and ZeRO-2 and ZeRO-3 degree two. We place Tw at ZeRO-3 and Tg at ZeRO-2. During setup, SS-HSDP builds the groups, data trees, and shard chains, and initializes or loads the main weights and optimizer states. Before the first training step, G initialize weights and scales at ZeRO-1.

The shard chain can also include other partitions. This section focuses on hierarchical data-parallel sharding rather than its composition with tensor or expert parallelism. Distributed updates For each operation in F and for G, we declare its input and output tensors and an update unit size for each operand: the number of elements required together, from one element for elementwise operations to a complete matrix for matrix operations. SS-HSDP uses the shard mappings to locate the corresponding operand fragments at the ZeRO-1 level. When all required operands for a unit are local, the operation runs on local views. Otherwise, an execution rank receives the required fragments through point-to-point (P2P) communication, performs the operation, and returns the modified fragments to their owners. Only declared tensors participate in this exchange.

Runtime At the first forward pass after initialization or an optimizer update, the runtime group performs the outer all-gather on the FP8 weights and UE8M0 scales prepared by G at ZeRO-1. At ZeRO-3, Tw prepares packed runtime 1 These shards are weights and kernel-specific scales ⃝. 2 retained across microbatches ⃝. For forward and backward computation, all-gather pre3 pares the full runtime weights and scales as needed ⃝. The GEMMs consume views of these buffers, and backward writes directly into the flattened runtime gradient 4 Reduce-scatter accumulates the gradients in buffer G′ ⃝. 5 This path repeats for each microbatch FP32 at ZeRO-2 ⃝. without repeating weight quantization or scale relayout. Full weight buffers can be released after use, while the runtime weight shards and accumulated gradients remain available.

In EG1, fi0 and fi2 run locally; only fi1 requires the complete head intermediate ui . Main weights and momentum remain sharded, and shard boundaries need not coincide with head boundaries. The data tree manages allocations and shared views, while the shard chain preserves each tensor’s logical positions across sharding levels. Together, they preserve packed, evenly sharded runtime buffers without restricting optimizer update units.

3.3

Execution Flow

Update After all microbatches, Tg maps the accumulated gradients to the update representation at ZeRO-2, and the outer reduce-scatter produces the ZeRO-1 main7 For Head0, the update spans ranks 2 gradient shards ⃝. 6 Each rank first executes f0 on its local fragments and 1 ⃝. to update momentum and produce BF16 u. Rank 2 then gathers the missing fragments of u from rank 1, executes Newton–Schulz f1 , and returns the corresponding FP32 v fragments to rank 1. Both ranks execute f2 locally to update their main-weight fragments. Main weights and momentum remain sharded throughout this matrix stage. Applying G to the updated main weights begins the next cycle.

Parameter Groups

SS-HSDP groups compatible parameters for shared storage and batched execution. Parameters with different dtypes, quantization formats, or optimizers may require separate tensor lists. The group boundaries for storage, runtime, and updates need not coincide. An atom group contains compatible runtime parameters and their update parameters. It packs matching tensor fields into separate tensor lists and constructs their data 6

Table 2 Representative kernel latencies (ms). Parentheses show latency increases over standalone execution.

The design above provides storage and sharding abstractions for decoupled runtime and update parameters. The next section builds on these abstractions to batch updates and transforms, eliminate redundant copies, and reduce contention between communication and computation.

4

Kernel FlashAttention forward DeepGEMM Grouped GEMM

5.774 (+22.7%) 9.087 (+51.2%)

Direct gradient write The same flattened layout supports direct reduce-scatter, but gradient production and accumulation can still introduce additional memory passes. PyTorch autograd produces per-parameter gradient tensors, which FSDP2 then copies into a packed reduce-scatter input buffer. As shown in Figure 5(b), SS-HSDP supports direct gradient writes from compatible backward kernels, such as those for linear layers, into views of the flattened reducescatter input buffer. This avoids separate per-parameter gradient tensors and the subsequent packing copies. The first write must initialize the entire target region, including zero-contribution entries; subsequent contributions accumulate. Otherwise, the buffer must be cleared before accumulating contributions. The remaining addition pass for accumulating reduced gradients across microbatches is fused into the final stage of reduce-scatter, as described next.

Batched Update and Transform

Within each runtime or update group, SS-HSDP batches compatible operations of the same type across parameters, as illustrated in Figure 5(a). With independent runtime and update groups, Tw and Tg can be batched within each layer, while F and G can be batched across the model. Within each stage, compatible local and P2P-dependent units are batched separately. In our fused Muon implementation, stages f0 and f2 use separate multi-tensor kernels, while f1 groups matrices with matching shapes and layouts for batched GEMMs. We similarly batch MXFP scale-layout transforms across tensors. Batching reduces kernel-launch overhead across parameters, particularly for small per-head updates and scale-layout transforms.

4.2

NCCL overlap

4.704 6.010

Optimization

The design above enables several optimizations, summarized in Figure 5: batching reduces kernel-launch overhead, shared views and direct gradient writes avoid redundant memory copies, and ZeroSM collectives reduce communication-induced SM contention.

4.1

Standalone

4.3

ZeroSM Communication

As shown in Figure 3, communication can overlap with computation to hide communication latency. In our profiling, NCCL communication kernels occupy 10–20% of the SMs. This creates SM contention between communication and computation, as illustrated in Figure 5(c). When ranks progress at different rates, communication kernels may remain resident while waiting for peer data or synchronization signals, occupying SM resources even without transferring data. Table 2 shows two profiled examples of slower compute kernels under NCCL overlap.

Fully Zero Copy

SS-HSDP arranges ZeRO-1 shards in the flattened buffer to match the group-rank order of hierarchical collectives, avoiding buffer reordering. In Figure 4, the four shards have owners [0, 2, 1, 3] in buffer order. Outer all-gathers within {0, 2} and {1, 3} produce two contiguous halves; runtime all-gathers within {0, 1} and {2, 3} then construct the full buffer directly. Gradient reduce-scatter follows the same layout in reverse, without reordering.

A fixed SM budget for computation kernels can mitigate contention, but also limits their parallelism after communication finishes. We instead use ZeroSM transfers to avoid reserving SM resources for communication.

Zero-copy weights When W ′ is a shared-storage view of W , using the same shard boundaries and rank ownership allows their local shards to share storage as well. In particular, W ′z1 and W ′z3 can directly share storage with W z1 and W z3 , respectively. Individual tensors access these tensor lists through sliced views. All-gather therefore reads directly from the existing local shard buffers, and forward and backward kernels consume views of the gathered buffer. No intermediate packing buffer or copy-in/copyout is needed for weights. Required scale relayouts are still performed.

ZeroSM P2P Within the NVLink domain, GPU copy engines perform peer transfers without using SMs. Our implementation uses registered symmetric allocations for communication buffers. Data-tree views address regions of these registered root buffers, avoiding separate allocations and registrations for individual views. Across nodes, we use NCCL Host remote memory access (RMA) [27]. CPU proxy threads coordinate NIC operations, while 7

Node 0

f2 rank 1 rank 2 P2P f1 rank 0 local

rank 3

Wgrad GEMM

Compute Time

Node 1

0 1 2 3 A B C D Inter-node exchange

Reduce

First Later write Accum

A C

B D

A C

B D

G′: BF16 / FP32 rank 0

Node 0

P2P

Partial sums

Local

A{0,1}

A{0} A{1}

FP32 reduce

A{2} A{3}

FP32 reduce

FP32 reduce + accumulate

Intra-node broadcast

f0 g′

Buffer Rank 0 Unit

1

2

Batch

(a) Batched update

3

G′: RS input

Cut

BF16 / FP32

(b) Direct write

Only NCCL ZeroSM compute + reduce

A A A A B B B B C C C C D D D D Full buffer on every rank

(c) SM contention

(d) ZeroSM AG

Node 1

P2P rank 2

A{2,3}

Inter-node P2P

FP32 / BF16 (SR)

A{0,1,2,3} G′z2

rank 0 · FP32 Persistent ZeRO-2

(e) ZeroSM RS

Figure 5 Optimizations enabled by runtime–update separation. (a) Compatible update units are batched per stage and execution rank. (b) Backward kernels write or accumulate directly into runtime-gradient views. (c) SM contention and explicit reduction overhead, shown schematically. (d) ZeroSM all-gather exchanges shards between nodes before broadcasting within each node. (e) Reduce-scatter for one output partition A; subscripts identify contributing ranks. SS-HSDP

Figure 6 MXFP8 training MFU and peak memory at fixed per-GPU workloads. We use sw = sg = 32; for SS-HSDP and Megatron-FSDP, so equals the total GPU count.

RDMA moves the payload without SM communication kernels. Together, these mechanisms support intra-node and inter-node ZeroSM P2P transfers between registered source and destination buffer regions.

the available SMs and HBM bandwidth for local arithmetic without waiting for remote peers inside the kernel. Figure 5(c) illustrates how reduced SM contention can outweigh this explicit reduction cost.

All-Gather To avoid transferring replicated shards across nodes, we perform inter-node exchange before intra-node broadcast (Figure 5(d)). Compared with exchanging each GPU’s buffer after intra-node broadcast, this reduces inter-node payload by a factor of N , where N is the number of participating GPUs per node.

To reduce inter-node traffic, we first aggregate contributions within each node (Figure 5(e)). The input buffer G′ holds BF16 or FP32 gradients. Intra-node P2P exchange brings slices to local reducers, which sum them in FP32, producing partials such as A{0,1} and A{2,3} . The partial sums remain FP32 or, for BF16 inputs, are stored in BF16 using stochastic rounding (SR). Inter-node P2P then transfers the partials to their destination ranks. The final FP32 reduction accumulates directly into the persistent buffer G′z2 , eliminating a separate result buffer and addition pass.

Reduce-Scatter Unlike all-gather, reduce-scatter still requires arithmetic on SMs. We use ZeroSM P2P for communication and launch reduction kernels only after their inputs are ready. These bandwidth-bound kernels can use 8

Table 3 Muon versus AdamW on 128 GPUs with BF16 computation. Peak memory (AdamW/Muon, GiB) includes PyTorch reserved memory and external communication buffers.

This follows the separation of BF16 communication and FP32 summation in DeepSeek-V4 [11], but applies SR to node-local partial sums. For a finite, in-range sum s, SR satisfies E[SR(s) | s] = s: the rounding error has zero expectation, consistent with unbiased gradient quantization in low-bitwidth training [3]. We write intra-node/inter-node for the dtypes of G′ and stored nodelocal partial sums, the latter also serving as the inter-node payload. The modes are FP32 (FP32/FP32), BF16+FP32 (BF16/FP32), and BF16SR (BF16/BF16). Section 6.1 evaluates their effect on training loss.

5

Experiments

5.1

End-to-End Performance

MFU (%) Model

sw

AdamW Muon Overhead Peak mem. (A/M)

Internal MoE 128 64 32

30.34 30.28 30.37

29.60 29.40 29.25

+2.5% +3.0% +3.9%

85.8/82.1 88.8/85.1 92.7/89.0

Qwen3

27.41 27.33 27.14

26.34 26.25 26.23

+4.1% +4.1% +3.5%

60.1/59.2 62.8/61.9 66.4/65.5

128 64 32

faster expert GEMMs. Long-context attention and normalization times remain the same when expert GEMMs use FP8.

Setup We compare SS-HSDP with PyTorch FSDP2, veScale-FSDP* , and Megatron-FSDP on representative compute-bound training workloads from video and language models: two internal video models, Qwen3-30BA3B [44], and a 12-layer MiMo [43] variant. We scale from 32 to 128 GPUs with fixed per-GPU workloads. Detailed settings and measurement procedures are provided in Appendix B.1.

5.2

Distributed Muon

Using distributed Muon instead of AdamW increases median step time by only 2.5–4.1% across the tested models and sharding configurations (Table 3). We compare paired runs at 128 GPUs with the same model, token count, sharding layout. Muon updates 88.1% of parameters in Internal MoE and 97.9% in Qwen3; AdamW updates the remainder. The implementation keeps main weights and momentum sharded: elementwise stages run locally, while matrix stages exchange only the required intermediate fragments. These results show that Muon can be applied to most model parameters with modest end-to-end training overhead.

Results Figure 6 reports the MXFP8 results. The upper panels show model FLOPs utilization (MFU) [7], defined as achieved model FLOPs per second divided by peak hardware FLOPs per second. The lower panels show peak memory usage per GPU. SS-HSDP achieves the highest MFU in all tested MXFP8 configurations. At 128 GPUs, the largest relative MFU gain over the fastest baseline is 21.7% on MiMo. MiMo’s large expert-weight volume makes weight quantization and copying costly. Weight reuse avoids repeated quantization across microbatches, while zero-copy execution eliminates redundant weight and gradient copies. These optimizations contribute to its larger performance gain.

6

Ablation and Analysis

6.1

Precision Ablation

To check that BF16SR in Section 4.3 does not change training convergence, we ablate communication precision and train a 12-layer, 8.1B-A1.3B Qwen3 MoE variant [44] on FineWeb-Edu [29] for 100B tokens on 128 GPUs with (sw , sg , so ) = (8, 8, 128).

Although retaining quantized weight shards increases persistent memory usage, outer ZeRO-1 sharding reduces the memory occupied by main weights and optimizer states. At 128 GPUs, peak memory is comparable to or lower than the lowest-memory baseline on all four models. On MiMo, SS-HSDP uses 125.0 GiB per GPU, 6.2% less than the lowest-memory baseline.

We compare intra/inter-node reduction input precisions: FP32, BF16+FP32, and BF16SR. All experiments use BF16 runtime weights and FP32 optimizer state. The three experiments use identical model, data order, optimizer settings, and seed; only the reduction precision changes. Figure 7 uses FP32 as the loss-difference baseline.

FP8 versus BF16 MXFP8 improves SS-HSDP’s throughput over BF16 by 16.3–18.5% on MiMo and about 7% on Internal MoE, while Qwen3 throughput changes by less than 1% (Appendix B.2). Qwen3 profiles show that quantization overhead can offset the savings from

All three experiments converge with training-loss differences on the order of 10−3 . The final-10k-step mean loss ranks BF16SR < BF16+FP32 < FP32, with differences of −1.10×10−3 and −8.39×10−4 from FP32, both below 0.05% in magnitude and stable over training.

* Results from our reproduction of veScale-FSDP.

9

Loss

FP32 Last 10k mean ≈ 2.25

7 2

Δ L = L − LFP3 2 (× 1 0 − 3 )

BF16+FP32 2.5

12

BF16SR

6.3 Outer ZeRO-1 Sharding and Weight Reuse

80–95k steps (zoom)

Outer ZeRO-1 sharding reduces persistent memory by further partitioning FP32 main weights and optimizer states, without enlarging the groups for repeated weight allgathers. In the main experiment, SS-HSDP and MegatronFSDP fix sw = 32 and set so to the GPU count, whereas FSDP2 and veScale-FSDP replicate 32-way optimizer shards. From 32 to 64 to 128 GPUs, SS-HSDP’s peak per-GPU memory falls by about 3.9 GiB per doubling on Internal MoE. On MiMo, it falls by about 8.6 GiB per doubling (Figure 6). Outer ZeRO-1 sharding yields greater memory savings for larger models. Weight reuse avoids repeated quantization and scale relayout by retaining low-precision weight shards and scales across microbatches. This requires additional persistent storage, but at 128 GPUs, SS-HSDP’s peak memory remains comparable to or lower than the lowest-memory baseline on all tested MXFP8 workloads in Section 5.

2.1

Stable-stage diff (steps ≥ 2k)

2 0

−3 0

20

40

60

80

95

Training steps (k) Raw

EMA (500)

Figure 7 Convergence and loss differences. All runs converge without sustained divergence. Final-10k-step mean loss ranks BF16SR < BF16+FP32 < FP32. Top: training loss. Bottom: ∆L = L − LFP32 . Solid lines show EMA (span 500). SS-HSDP · ZeroSM

SS-HSDP · SM-based

7

Large-model pretraining practices Training language models [11, 14, 44] and video models [1, 37, 38, 42] requires flexibility in sharding, runtime numerical representations, and optimizer update granularity. DeepSeekV4 uses different ZeRO partitions for dense and expert parameters to support Muon, while DeepSeek-V3 [10] and MiMo-V2-Flash [43] combine FP8 computation with higher precision for selected operations and parameters. Kimi K3 [18] and GLM-5 [13] use attention heads as Muon update units.

Figure 8 The benefit of ZeroSM depends on the workload. Throughput relative to default FSDP2 for MiMo (12 layers), using median step times. The dashed line marks default FSDP2; the x-axis (tokens per GPU per microbatch) is logarithmic.

6.2

Related Work

ZeroSM Communication across Workloads

We vary per-GPU workload size for MiMo at 64 GPUs with sharding degree 64 under BF16 computation and four microbatches per optimizer step, comparing ZeroSM, SM-based communication, and default FSDP2 (Appendix B.3).

Runtime optimizations for sharded training Runtime optimizations target communication volume, memory traffic, and GPU resource contention during forward and backward execution. MiCS [46] and ZeRO++ [39] improve communication locality through smaller sharding groups and intra-node backward weight all-gathers, respectively. TorchTitan [19] gathers weights in FP8 for subsequent low-precision computation. Collective bucketing and reordering [45] and topology-aware partitioning [2] improve overlap; DeepCompile [36] prefetches parameters to overlap all-gathers with computation and retains gathered weights to avoid repeated all-gathers. Computation–collective fusion avoids intermediate memory accesses [16]. Host-driven NCCLX and NCCL paths offload data movement to reduce SM usage [26, 35].

ZeroSM’s advantage emerges at larger workloads Figure 8 shows a crossover. At the smallest workload, ZeroSM delivers 33.5% lower throughput than SM-based SS-HSDP; at the largest, it gains 8.8%. The trend is consistent with a latency-versus-contention trade-off: communication-bound loads expose collective latency on the critical path, whereas compute-bound loads hide it through overlap and free SM capacity for model kernels. SM-based communication also outperforms FSDP2 SMbased communication retains zero-copy all-gather weight views and direct-gradient writes. Across the five tested workloads, it achieves 1.106–1.566× the throughput of default FSDP2, including 1.561× at the smallest workload. Thus, the throughput advantage over FSDP2 persists without ZeroSM.

Distributed parameter updates Distributed optimizer implementations coordinate matrix inputs, update computation, and persistent states across ranks. Distributed Muon [20] gathers matrix inputs for orthogonalization while retaining sharded optimizer states. Canzona [40] 10

and DMuon [6] assign matrix updates to designated ranks and balance optimizer workloads. Canzona further pipelines input reconstruction, update computation, and result redistribution under tensor parallelism. LinS [5] and InternEvo [4] independently configure parameter, gradient, and optimizer-state sharding to trade memory for communication.

8

[6] Vincent Chen, Starrick Liu, Regis Cheng, Dance Yang, Shalfun Li, Ryan Yu, Lucy Liang, Hang Su, Roy Gan, Hao Wang, and Qian Wang. DMuon: Efficient distributed Muon training with near-Adam overhead, 2026. [7] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, et al. PaLM: Scaling language modeling with pathways, 2022. [8] Damai Dai, Chengqi Deng, Chenggang Zhao, R.X. Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y. Wu, Zhenda Xie, Y.K. Li, Panpan Huang, Fuli Luo, Chong Ruan, Zhifang Sui, and Wenfeng Liang. DeepSeekMoE: Towards ultimate expert specialization in mixture-of-experts language models. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1280–1297, Bangkok, Thailand, August 2024. Association for Computational Linguistics.

Conclusion

SS-HSDP shows that flexible optimizer updates and efficient runtime execution can be supported within a unified sharded training framework. Shared storage and hierarchical sharding allow structured updates to coexist with packed computation and communication layouts while avoiding redundant data copies. Our evaluation demonstrates that this flexibility can be achieved with improved training efficiency and competitive memory usage. These findings support separating optimizer update requirements from runtime data organization as a practical approach to sharded training.

[9] DeepSeek-AI. DeepGEMM: Scaling factor format. Official DeepGEMM repository documentation, n.d. Accessed 2026-10-22. [10] DeepSeek-AI et al. DeepSeek-V3 technical report, 2024.

References

[11] DeepSeek-AI et al. DeepSeek-V4: Towards highly efficient million-token context intelligence, 2026.

[1] Fan Bao, Chendong Xiang, Gang Yue, Guande He, Hongzhou Zhu, Kaiwen Zheng, Min Zhao, Shilong Liu, Yaole Wang, and Jun Zhu. Vidu: A highly consistent, dynamic and skilled text-to-video generator with diffusion models, 2024.

[12] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1–39, 2022. [13] GLM-5-Team et al. GLM-5: from vibe coding to agentic engineering, 2026.

[2] Chang Chen, Xiuhong Li, Qianchao Zhu, Jiangfei Duan, Peng Sun, Xingcheng Zhang, and Chao Yang. Centauri: Enabling efficient scheduling for communicationcomputation overlap in large model training via communication partitioning. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3, 2024.

[14] Aaron Grattafiori et al. The Llama 3 herd of models, 2024. [15] Dirk Groeneveld et al. OLMo: Accelerating the science of language models, 2024. [16] Abhinav Jangda, Jun Huang, Guodong Liu, Amir Hossein Nodehi Sabet, Saeed Maleki, Youshan Miao, Madanlal Musuvathi, Todd Mytkowicz, and Olli Saarikivi. Breaking the computation and communication abstraction barrier in distributed machine learning workloads. In Proceedings of the 27th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, 2022.

[3] Jianfei Chen, Yu Gai, Zhewei Yao, Michael W. Mahoney, and Joseph E. Gonzalez. A statistical framework for lowbitwidth training of deep neural networks. In Advances in Neural Information Processing Systems, volume 33, pages 883–894, 2020. [4] Qiaoling Chen, Diandian Gu, Guoteng Wang, Xun Chen, Yingtong Xiong, Ting Huang, Qinghao Hu, Xin Jin, Yonggang Wen, Tianwei Zhang, and Peng Sun. InternEvo: Efficient long-sequence large language model training via hybrid parallelism and redundant sharding, 2024.

[17] Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. [18] Kimi Team et al. Kimi K3: Open frontier intelligence, 2026.

[5] Qiaoling Chen, Qinghao Hu, Guoteng Wang, Yingtong Xiong, Ting Huang, Xun Chen, Yang Gao, Hang Yan, Yonggang Wen, Tianwei Zhang, and Peng Sun. Lins: Reducing communication overhead of ZeRO for efficient LLM training. In 2024 IEEE/ACM 32nd International Symposium on Quality of Service (IWQoS), pages 1–10, 2024.

[19] Wanchao Liang, Tianyu Liu, Less Wright, Will Constable, Andrew Gu, Chien-Chin Huang, Iris Zhang, Wei Feng, Howard Huang, Junjie Wang, Sanket Jayant Purandare, Gokul Nadathur, and Stratos Idreos. TorchTitan: One-stop PyTorch native solution for production ready

11

LLM pretraining. In International Conference on Learning Representations, 2025.

[34] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. MegatronLM: Training multi-billion parameter language models using model parallelism, 2019.

[20] Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, Yanru Chen, Huabin Zheng, Yibo Liu, Shaowei Liu, Bohong Yin, Weiran He, Han Zhu, Yuzhi Wang, Jianzhou Wang, Mengnan Dong, Zheng Zhang, Yongsheng Kang, Hao Zhang, Xinran Xu, Yutao Zhang, Yuxin Wu, Xinyu Zhou, and Zhilin Yang. Muon is scalable for LLM training, 2025.

[35] Min Si, Pavan Balaji, Yongzhou Chen, et al. Collective communication for 100k+ GPUs, 2025. [36] Masahiro Tanaka, Du Li, Umesh Chand, Ali Zafar, Haiying Shen, and Olatunji Ruwase. DeepCompile: A compilerdriven approach to optimizing distributed deep learning training, 2025. Version 2, 19 February 2026.

[21] Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, and Hao Wu. Mixed precision training. In International Conference on Learning Representations, 2018.

[37] Team Seedance et al. Seedance 2.0: Advancing video generation for world complexity, 2026. [38] Wan Team. Wan: Open and advanced large-scale video generative models, 2025.

[22] NVIDIA. Cosmos world foundation model platform for physical AI, 2025.

[39] Guanhua Wang, Heyang Qin, Sam Jacobs, Xiaoxia (Shirley) Wu, Connor Holmes, Zhewei Yao, Samyam Rajbhandari, Olatunji Ruwase, Feng Yan, Lei Yang, and Yuxiong He. ZeRO++: Extremely efficient collective communication for large model training. In International Conference on Learning Representations, pages 50035– 50053, 2024.

[23] NVIDIA. cuBLAS: Block scaling factors layout. cuBLAS 13.4 Documentation, n.d. Accessed 2026-10-22; layouts and supported scaling modes are version-specific. [24] NVIDIA. Megatron-FSDP. Megatron Core 0.18.2 Documentation, n.d. Accessed 2026-10-22.

[40] Liangyu Wang, Siqi Zhang, Junjie Wang, Yiming Dong, Bo Zheng, Zihan Qiu, Shengkun Tang, Di Wang, Rui Men, and Dayiheng Liu. Canzona: A unified, asynchronous, and load-balanced framework for distributed matrix-based optimizers, 2026.

[25] NVIDIA. Megatron-LM and Megatron Core. GitHub repository, n.d. Accessed 2026-10-28. [26] NVIDIA. NCCL release 2.30.7. NCCL Release Notes, n.d. Zero-SM AllGather and All2all; accessed 2026-10-30.

[41] Zezhou Wang, Youjie Li, Zhiqi Lin, Jiacheng Yang, Cong Xie, Guanyu Feng, Zheng Zhong, Ziyue Huang, Hongyu Zhu, Zhi Zhang, Yanghua Peng, and Xin Liu. veScaleFSDP: Flexible and high-performance FSDP at scale. In Proceedings of Machine Learning and Systems, volume 8, pages 763–776. MLSys, 2026. Industry Track.

[27] NVIDIA. NCCL user guide: User buffer registration. NCCL 2.31.2 Documentation, n.d. Accessed 2026-10-22. [28] Open Compute Project. OCP microscaling formats (MX) specification. Technical Report Version 1.0, Open Compute Project, September 2023. [29] Guilherme Penedo, Hynek Kydlíček, Loubna Ben Allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro von Werra, and Thomas Wolf. The FineWeb datasets: Decanting the web for the finest text data at scale, 2024.

[42] Bing Wu, Chang Zou, Changlin Li, et al. HunyuanVideo 1.5 technical report, 2025. [43] Xiaomi LLM-Core Team, Bangjun Xiao, Bingquan Xia, Bo Yang, et al. MiMo-V2-Flash technical report, 2026.

[30] PyTorch. PyTorch DTensor: Distributed tensor. PyTorch Documentation, n.d. Version 2.14; accessed 2026-10-28.

[44] An Yang et al. Qwen3 technical report, 2025.

[31] PyTorch. PyTorch FSDP2: torch.distributed.fsdp.fully_shard. PyTorch Documentation, n.d. Version 2.14; accessed 2026-10-22.

[45] Ruisi Zhang, Tianyu Liu, Will Feng, Andrew Gu, Sanket Purandare, Wanchao Liang, and Francisco Massa. SimpleFSDP: Simpler fully sharded data parallel with torch.compile, 2024.

[32] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. ZeRO: Memory optimizations toward training trillion parameter models, 2019.

[46] Zhen Zhang, Shuai Zheng, Yida Wang, Justin Chiu, George Karypis, Trishul Chilimbi, Mu Li, and Xin Jin. MiCS: Near-linear scaling for training gigantic model on public cloud. Proceedings of the VLDB Endowment, 16(1):37–50, 2022.

[33] Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. DeepSpeed: System optimizations enable training deep learning models with over 100 billion parameters. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, KDD ’20, pages 3505–3506. Association for Computing Machinery, 2020.

[47] Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, ChienChin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, Alban Desmaison, Can Balioglu, Pritam Damania, Bernard Nguyen, Geeta Chauhan,

12

B

Yuchen Hao, Ajit Mathews, and Shen Li. PyTorch FSDP: Experiences on scaling fully sharded data parallel, 2023.

A

Evaluation and Artifact Details

Implementation size. The inspected source snapshot contains 85 Python and native source files, comprising 21,624 physical lines. A source-oriented count excluding blank lines, comments, and Python docstrings gives approximately 17,606 lines: 14,217 in Python and 3,389 in native code. Preprocessor directives and structural delimiters are retained. These totals cover shared infrastructure and specialized backends. The artifact preparation will package the counting script, file manifest, and immutable source identifier alongside this breakdown, so the same rule can be applied to extension examples.

Implementation Details

Mapping paper terms to the source. A logical parameter is represented by Parameter; the kernel-facing runtime tensor is carried by RuntimeParameter. The source identifiers zero3, zero2, and zero1 denote the retained weight, accumulation, and optimizer shard levels, respectively. Their paper symbols are sw , sg , and so . The copy-engine/RMA implementation resides under zero_sm; ZeroSmGrad selects its gradient lifecycle and precision modes. These identifiers are implementation names rather than additional parameter or sharding concepts.

Integration and extension measurements. We will provide runnable examples showing how to convert a model, define update parameters, construct runtime parameters, and configure sharding. The four-head query example packs WQ = [W0 ; W1 ; W2 ; W3 ], with Wi ∈ Rdh ×d , into a BF16 runtime tensor and retains FP32 main weights for per-head Muon. Its execution settings follow Table 4. Each example includes initialization, gradient clearing, first/last microbatch flags, normalization, and update. Engineering measurements count added conversion code, optimizer declarations, backend-specific code, and integration touchpoints for AdamW, expert-wise Muon, and per-head Muon. Shared utilities are reported separately. Correctness checks compare the declared logical updates with a local reference.

Packing and compatibility. TensorList stores member shapes and offsets and derives the fragments at each shard level. An Atom groups parameters with matching tensor field names, dtypes, devices, optimizer-state layouts, and update buffers. AdamW and Muon parameters with different state schemas therefore use separate Atom objects. Runtime and update groups batch compatible operations across these objects. Each input or output field declares the number of elements per update unit. Fields may use different numbers of elements per unit, but must describe the same number of units. Allocations and update execution. The data tree uses DataNode objects whose root owns storage. Finalization resolves child offsets and forwards allocation callbacks to that root. These callbacks attach or detach the runtime tensor’s view while retaining its metadata. Update intermediates such as ui and vi have update-scoped allocations; executors retaining temporary views are released or rebuilt with those allocations. Stages run in order within one update group, committing intermediates and scalar state for the next stage. Independent groups retain separate schedules.

Reproduction records. Machine-readable configurations will record hardware, software versions, source revisions, topology, rank mapping, model and optimizer settings, sharding degrees, grouping policy, precision, seeds, and timing boundaries. Raw per-step times and memory observations accompany aggregate statistics. Numerical cases preserve the original inputs, communicated representations, and rounding seeds where feasible, allowing input and partial quantization to be examined independently. Diagnostics document failed configurations and the selection of supported replacements.

Execution conditions. Table 4 distinguishes metadata checks from execution conditions whose correctness depends on the producer and scheduler. The direct gradient write mechanism and BF16 communication path are described in Section 4; the table records the corresponding storage, hardware, and restart conditions. In particular, explicit SUM normalization and valid gradients for every updated logical parameter connect the runtime’s accumulated values to the intended optimizer step.

Artifact preparation. Release preparation covers licensing, dependency specifications, runnable examples, multirank correctness tests, benchmark scripts, and backend configuration guidance. A manifest links each paper experiment to its configuration and output records. Together, these materials make extension size, reproduced behavior, and measured performance traceable to the same implementation snapshot.

13

Table 4 Execution conditions, their rationale, and the consequence of violating them. Checks are identified where implemented. Condition

Rationale

Consequence if violated

Direct writes cover every element on the first write of the current buffer lifetime; later writes accumulate. Compatible field schemas within each Atom; transforms preserve their declared complete units. Ordered stages within each update group; intermediates are produced before use and scalar updates agree across fragments. Nested sharding degrees; equal shards for copy-engine/RMA collectives; balanced node-major ranks and stable registered bases/offsets. Events or stream order connect producers, consumers, and reuse; shared RS workspace calls are serialized. BF16 partials use supported hardware SR; seed/counter state is restored on the same logical rank. BF16 RS uses full-tensor, layout-preserving Cast transforms, matching packed offsets, and separate BF16 input and FP32 accumulation storage. The specialized copy-engine/RMA path runs without CUDA Graph replay. All contributing shared-parameter forwards are registered before their final-microbatch backwards. Every updated trainable logical parameter has a valid final gradient; new accumulation rounds clear prior gradients explicitly. Loss, replica, and microbatch normalization are explicit after final synchronization.

Overwrite mode replaces clearing; zero-contribution slices receive explicit zeros.

Unwritten slices retain stale data, or repeated overwrites discard contributions.

Packing connects matching fields, dtypes, devices, optimizer states, and update buffers.

Planning rejects incompatible metadata; incompatible fields use separate Atom objects.

Each stage consumes committed outputs; independent groups have separate schedules.

Reads of temporary buffers before production are rejected; inconsistent scalar updates violate stage semantics. Invalid geometry and detected binding changes are rejected; lifetime violations risk invalid accesses.

B.1

Collective layouts and registered addresses describe the same rank-to-data mapping.

One arena region serves successive nonoverlapping lifetimes.

Concurrent reuse races with unfinished transfers or reductions.

The SR instruction and element-to-counter mapping define the numerical path.

Unsupported hardware or compiled images are rejected before communication; missing RNG state changes resumed rounding. Incompatible transform/layout metadata and detected storage overlap are rejected.

Reduction consumes the input directly while accumulation shards persist.

Host protocol counters and Python SR counters have no implemented replay semantics. Backward-use counts determine eligibility for final synchronization. Updates consume gradient shards, without unused-parameter discovery or missing-gradient skipping. Gradient collectives use SUM semantics.

Performance Settings

Internal MoE Internal Dense Qwen3 MiMo (12L)

Layers

Tokens

MBS

GA

33.8 19.2 30.5 73.5

8 24 48 12

72,186 72,186 24,576 24,576

1 1 4 4

4 4 4 4

Missing or stale contributions enter the update.

Omitting the intended scale changes the effective gradient.

Software. All end-to-end runs use PyTorch 2.14.0+cu132, CUDA 13.2, NCCL 2.30.7, and DeepGEMM 2.6.1b2. The PyTorch FSDP2 baseline is the implementation in PyTorch 2.14; veScale-FSDP runs from our reproduction; Megatron-FSDP is the implementation in Megatron Core 0.18.2. Baselines use the same model partition, precision settings, and per-GPU workload as SS-HSDP.

Table 5 Performance workloads. Parameter counts include all experts. MBS denotes samples per GPU per microbatch; GA is the number of microbatches per optimizer step. Params (B)

Sequential forward/backward/forward/backward reuse can synchronize too early.

across all GPUs. MXFP8 uses DeepGEMM for expert GEMMs, or the video FFN of Internal Dense; remaining GEMMs use BF16, with FP32 optimizer states. SS-HSDP and Megatron-FSDP retain quantized weights, whereas FSDP2 and veScale-FSDP quantize on demand.

Table 5 gives the workloads used in Section 5.1. All runs use 32, 64, or 128 GPUs (eight per node), AdamW, full activation recomputation, and gradient accumulation (GA) of four microbatches per optimizer step. Inputs are synthetic. Per-GPU microbatch size (MBS) and input length remain fixed as GPU count increases.

Model

Graph capture/replay is unsupported.

Timing uses six-step medians after two warmup steps. Model FLOPs exclude recomputation. We report per-GPU model FLOPs utilization (MFU), i.e., achieved model FLOP/s as a fraction of the GPU’s BF16 peak FLOP/s. Memory is peak PyTorch reserved memory plus external communication buffers, maximized over measured steps and ranks.

Weight and gradient sharding degrees are fixed at 32; SS-HSDP and Megatron-FSDP shard optimizer states 14

SS-HSDP

Figure 9 BF16 training performance. Workloads and sharding layouts match Figure 6, with BF16 GEMMs. Top: per-GPU model FLOPs utilization (MFU). Bottom: peak reserved memory plus external communication buffers. Each model uses its own y-axis ranges.

B.2

BF16 Training Performance

The sweep has 15 configurations. Each point is the median of zero-based steps 5–14 (ten timed steps after five warmup steps), from one run per configuration; these are within-run samples, not independent repetitions. All accepted runs passed configuration, host-identity, and GPUclean checks, with no OOM or allocation retry.

Figure 9 complements the MXFP8 results in Section 5.1 with the same workloads, GPU counts, and measurement procedure. Table 6 compares SS-HSDP’s MXFP8 and BF16 throughput, using ratios of median step times. The benefit depends on the workload: it is largest on MiMo, smaller on the internal models, and negligible on Qwen3.

ZeroSM uses BF16 runtime gradients with FP32 internode reduction and accumulation; SM-based communication uses FP32 gradients and NCCL collectives. Both retain SS-HSDP’s shared weight views and direct gradient writes, but differ in gradient storage, symmetric arenas, hierarchical reduction, and accumulation. ZeroSM offloads payload movement while intra-node and final reductions still use SMs, so the comparison does not isolate transport latency or SM occupancy. Both SS-HSDP modes use (sw , sg , so ) = (64, 64, 64), and FSDP2 uses sharding degree 64; the layout was fixed rather than varied, so gains are measured at this sharding degree only. SS-HSDP performs the gradient-to-optimizer reduction only on the final microbatch, and no deferred-replica-all-reduce variant is evaluated.

Table 6 SS-HSDP MXFP8 throughput relative to its own BF16 runs. Model

32 GPUs

64 GPUs

128 GPUs

Internal MoE Internal Dense Qwen3 MiMo (12L)

1.072× 1.031× 0.995× 1.185×

1.071× 1.032× 0.994× 1.163×

1.072× 1.027× 1.002× 1.165×

B.3

Communication-Backend Workload Sweep

Section 6.2 uses a separate BF16 sweep, not the MXFP8 scaling experiment. MiMo (12 layers, 73.5B parameters) runs on eight B200 nodes (eight GPUs per node) for each point, with all three configurations on the same hosts and input shape; otherwise the setup follows Appendix B.1, with DeepGEMM BF16 and seed 42. Context and expert parallelism are disabled; profiling and checkpointing are disabled during timing. MiMo uses MBS 1 at 2,048, 8,192, and 24,576, then MBS 4 at 24,576 and 36,864 (98,304 and 147,456 tokens per GPU per microbatch). Token count is a workload descriptor, not an arithmeticintensity classification. 15

Related documents

• Heterogeneous Structured Pruning Under Resource Constraints for Hardware-Efficient Large Language Models#1327031
• SAMM: Sharded Automated Market Maker#1199632
• FasterPy: An LLM-based Code Execution Efficiency Optimization Framework#1116186
• The Endgame of FSDP Communication: Layout-Aware Collectives on Copy Engines, RDMA, and NVSwitch#1489020
• The Endgame of FSDP Communication: Layout-Aware Collectives on Copy Engines, RDMA, and NVSwitch#1482976