Suppose you want to train a model with a matrix optimizer such as Muon. The optimizer does not think of a weight matrix as a bag of independent numbers; it needs the original 2D matrix so it can apply a matrix-level transformation. Now suppose your distributed training system cuts that matrix at whatever element happens to fall on the next GPU boundary. The model still fits in memory, but the optimizer’s natural unit of work has been destroyed.
That is the tension this paper starts from. FSDP and ZeRO became foundational because they let training state—parameters, gradients, and optimizer state—be spread across many devices without forcing the model author to redesign the model around the parallelization scheme. For ordinary element-wise operations, an arbitrary shard boundary is often fine. Even row-wise sharding is enough for many matrix-shaped computations. The sharding scheme can mostly be treated as a systems detail.
Modern training methods make that assumption less safe. Shampoo and Muon operate on matrix structure. Block-wise quantization has a different but related requirement: every quantization block should ideally live entirely on one device, because its values share scale metadata. If a block straddles two devices, the implementation now has to exchange metadata or reconstruct the block before quantizing it. What looked like a harmless shard boundary turns into extra communication and special-case logic.
Existing systems can work around these mismatches, but the workarounds have a cost. Model or optimizer code can be rewritten so its boundaries happen to match the sharding scheme. The runtime can insert padding. It can add boundary checks, copies, or collectives. Each option makes the system either less transparent or less efficient.
veScale-FSDP changes the question. Instead of asking the computation to adapt to a fixed sharding format, it asks whether the sharding format can represent the structure the computation already needs. The goal is ambitious because flexibility alone is not enough: the paper is targeting production training at thousands to more than 10,000 GPUs, where a small copy, a little padding, or a badly aligned collective can become a persistent tax on every iteration.
To see exactly where those costs appear, we first need the one communication pattern that all of these FSDP variants are trying to preserve.
FSDP gets its memory savings by refusing to keep a complete copy of every model state on every device. Most of the time, each of the devices owns only a shard. The full tensor is reconstructed only when a layer actually needs it for computation.
For a parameter used by the next layer, the devices run an AllGather. Each device contributes the shard it owns, and the collective makes the complete parameter available for the forward pass. The backward pass also needs the materialized parameter while computing gradients. Once that work is done, keeping a full gradient everywhere would waste the memory FSDP was meant to save.
So the reverse path uses ReduceScatter. Contributions to the gradient are reduced across devices, and the result is immediately divided back into shards. The steady-state rhythm is therefore simple: keep state sharded, temporarily materialize what computation needs, then return to sharded ownership.
DeepSpeed ZeRO and PyTorch FSDP1 implement this rhythm using concatenated buffers whose shard boundaries can cut through individual parameters at the element level. FSDP2 changes the representation: each parameter is represented as a DTensor, typically with an even Shard(0) placement. That gives the runtime a cleaner per-parameter object for communication, computation, and checkpointing.
veScale-FSDP does not discard this training contract. The model author still expects ordinary layer computation between AllGather and ReduceScatter. The paper changes the representation and physical layout underneath that contract, because that is where fixed shard boundaries start conflicting with modern tensor structure.
Consider a matrix whose rows are meaningful units for computation and whose 2D sub-blocks are meaningful units for quantization. An element-wise shard can cut anywhere inside that matrix. Once it does, a device may hold a fragment with neither the original matrix shape nor a complete quantization block. Operations that depend on those structures now need reconstruction or special handling.
Moving to even row-wise sharding fixes part of the problem. A Shard(0) DTensor splits the tensor along its first dimension, so each local shard is made of complete rows. That is much friendlier to matrix operations, and DTensor can redistribute the tensor between placements using collectives such as All2All. But an even row boundary is still chosen by the number of devices, not by the block shape required by quantization. A 128-row quantization block, for example, can still be cut by an otherwise perfectly valid even split.
RaggedShard separates two decisions that fixed sharding schemes tend to conflate. For each tensor , the system can choose an atomic non-shardable block size , and it can distribute different numbers of those blocks to different devices. If one row is the meaningful unit, can represent a row. If a larger 2D region must remain intact, the granularity can represent that block instead.
That makes block-wise RaggedShard more general than a special-purpose format for one optimizer. Element-wise, row-wise, and custom block-wise layouts can all be expressed by changing the atomic unit. The important constraint is simple: a device boundary is allowed only between those units.
This solves the representational mismatch, but it creates a harder systems problem. Once tensors have different block sizes and uneven shard lengths, grouping them efficiently for collective communication is no longer a matter of slicing one flat buffer at equal offsets.
The older systems each avoid one problem by accepting another. DeepSpeed ZeRO and FSDP1 concatenate tensors and shard the resulting buffer at element boundaries. That is compact and straightforward for classic ZeRO-style execution, but the representation has no native notion of the matrix or quantization block that must remain intact.
FSDP2 improves the abstraction by representing each parameter as an evenly sharded DTensor. The price shows up in memory layout around collectives. Before an AllGather, local parameter shards must be copied into the communication buffer. After the AllGather, pieces belonging to one full parameter can occupy interleaved addresses in the output buffer, while computation expects that parameter to be contiguous. FSDP2 therefore performs a Copy-Out to materialize the contiguous tensor. The ReduceScatter direction has the corresponding Copy-In cost.
The paper measures this directly on GPT-OSS-120B with 64 H800 GPUs. With Shard(0), AllGather takes 43.71 ms and the associated Copy-Out another 5.22 ms; ReduceScatter takes 94.24 ms and Copy-In another 12.37 ms. With Shard(1), the copy costs are larger still: 13.72 ms after AllGather and 23.14 ms before ReduceScatter. In the paper’s end-to-end measurements, these interleaved copies can account for up to 14% of a training iteration.
Megatron-FSDP attacks the copy problem more directly. It uses a concatenated representation designed for zero-copy collectives, while arranging boundaries so tensors can still appear row-sharded for DTensor-compatible checkpointing. But satisfying those row boundaries requires padding. Padding is not free: those bytes occupy memory and travel through collectives even though they contain no useful model state.
So the target for veScale-FSDP is stricter than “support uneven shards.” A useful production layout has to preserve the chosen block granularity, keep each tensor contiguous, balance the amount communicated by each device, and let communication operate directly on the storage the tensors already use. RaggedShard provides the first requirement; the next question is how to define it precisely enough for the runtime to plan the other three.
RaggedShard makes the forbidden cut explicit. For each tensor , the system chooses an atomic block size . A shard boundary may fall between these blocks, but it may not split one of them. If the tensor contains elements, then the number of atomic blocks is
The second change is just as important: those blocks do not have to be divided evenly across devices. One device can own more blocks than another. That uneven distribution is what makes the shard “ragged.” It lets the runtime respect structural boundaries even when an equal split would land in the wrong place.
The same abstraction covers several useful cases. If is one element, RaggedShard behaves like fine-grained element-wise sharding. If the atomic unit is one row, different devices can own different numbers of complete rows. If training requires a larger block to stay intact, the atomic unit can be that custom block instead. The format is therefore defined by a constraint on where cuts may occur, rather than by one hard-coded tensor dimension.
This matters for software composition. RaggedShard is implemented as a DTensor placement, so the tensor still participates in PyTorch’s distributed-tensor machinery. The model or optimizer can express the structure it needs without replacing the surrounding distributed abstraction with a custom data structure.
There is a catch. FSDP collectives are fastest when many tensors are grouped into large communication buffers. Once every tensor can have a different and an uneven number of blocks per device, simply concatenating those tensors can violate the very constraints RaggedShard was introduced to preserve.
Take three RaggedShard tensors and place them back-to-back in a communication buffer. The first device boundary may land halfway through one tensor’s -sized block. The representation said that block was indivisible, but the buffer layout has just split it across devices. Block-wise quantization would now need extra communication to recover the missing part.
Trying to fix that boundary with padding can create a second problem. If the padding is inserted inside a tensor, its logical contents no longer occupy one contiguous memory interval. After an AllGather, computation cannot simply view that region as the full tensor; some form of rearrangement or copying is needed. That recreates the same class of overhead veScale-FSDP is trying to remove.
There is also a collective-level constraint. Fast collective implementations expect symmetric communication: each participating device should contribute and receive compatible amounts of data. Different tensor sizes, different values, and arbitrary padding can leave one device with a larger local buffer than another, wasting bandwidth or forcing yet more padding.
The paper therefore changes the order of operations. It allows tensors to be permuted, then inserts any necessary padding between tensors. This leaves each tensor contiguous and gives the planner more freedom to make device boundaries coincide with atomic-block boundaries. The buffer is then divided into equal-size local shards.
So RaggedShard alone is only a language for saying what must remain intact. A separate planner has to find a global arrangement that satisfies all those local tensor constraints while keeping communication balanced and padding small.
The planner can be stated as a packing problem over one global communication buffer. Let be the RaggedShard tensors in a communication group, and let the group span devices. Every device must own the same-size local piece of that buffer; call that size . The whole buffer therefore has size .
Each tensor must occupy one contiguous interval . If is its size, contiguity and size give
and fitting into the global buffer requires . Intervals belonging to different tensors are also forbidden from overlapping. Padding is allowed in the gaps, which is exactly why the planner can preserve contiguity while moving tensor boundaries around.
The RaggedShard constraint appears at every device boundary . If a boundary falls inside tensor , its offset from the beginning of that tensor must coincide with an atomic-block boundary:
If lies outside , there is no constraint from that tensor. This condition is the mathematical version of “never split one -sized block across devices.”
The objective is to minimize . A larger means a larger global buffer and more padding or communication volume, so the planner wants the smallest equal per-device shard that satisfies tensor contiguity, non-overlap, and every block-alignment constraint simultaneously.
That formulation is precise enough to optimize. Unfortunately, the freedom to order tensors and place padding turns out to make the exact problem computationally hard.
The paper proves that minimizing under these placement constraints is NP-hard, using a reduction from the classic Partition problem. That result changes what a practical implementation should optimize for. A planner that insists on the exact global optimum may spend far too long solving the layout before training even begins.
One option is to encode the constraints as an Integer Linear Program. In principle, that gives a general exact optimization route. In practice, the authors report that off-the-shelf ILP solvers can take tens of minutes on realistic cases and may hit system timeouts. An initialization optimization that takes that long is a poor trade when the purpose of the layout is to make the subsequent training loop faster.
The desired replacement therefore has two properties. It must keep the hard semantic constraints—tensors remain contiguous, device shards remain balanced, and no block is split—and it must find a good in polynomial time. The authors choose a heuristic dynamic-programming approach that exploits the regular structure of transformer parameters.
Before looking at that algorithm, there is one source-fidelity limitation to state clearly: the paper says the problem is NP-hard by reduction from Partition, but the supplied text does not give the full construction of that reduction.
What can safely be said about the reduction is the structural part. The planner is choosing how indivisible blocks are arranged around equal-size device boundaries while also preserving whole contiguous tensors. Changing the tensor order or the padding changes which collections of blocks fit before each boundary. That is a discrete partitioning decision, rather than a smooth optimization over continuous positions.
The paper states that an instance of Partition can be reduced to this layout problem, which is enough to establish NP-hardness. It does not give the explicit mapping in the text available here—such as the exact tensors, block sizes, and device configuration constructed from a Partition instance. Filling in those details ourselves would turn a paper walkthrough into a new proof, so we will not do that.
Fortunately, none of the system design depends on knowing the omitted construction. What matters operationally is that arbitrary tensor permutations are too expensive to search exhaustively. The implemented planner narrows the search using properties that are common in transformer models, then solves the remaining placement problem with dynamic programming and aligned candidate shard sizes.
The planner first fixes a tensor order. Searching every permutation would recover the combinatorial difficulty we were trying to avoid, so the authors exploit a practical property of transformers: linear weights dominate parameter count, and similar layers tend to produce repeated tensor shapes and block granularities. They tested the default order, sorting by block size, and sorting by tensor shape; the paper says these choices are optimal or near-optimal in their statistics, and uses the default order for simplicity and debugging.
For a candidate per-device size , the dynamic program asks a feasibility question. Define as the minimum number of device-local shards needed to place every tensor before plus the first atomic blocks of . If the final state needs at most shards, then that is feasible.
A naive DP could still touch every block position of a huge tensor. The useful observation is monotonicity:
As advances, the required shard count can only stay the same or increase. Since there are only possible shard-count values, the algorithm groups ranges of that share the same value and skips the intermediate states. This is what makes the DP practical even when individual tensors contain many atomic blocks.
The outer search also avoids arbitrary values of . The planner starts from the collective-preferred alignment and progressively incorporates block granularities from through least-common-multiple constraints. Its case analysis shows that when a tensor fully contains a local shard, feasible sizes must lie on multiples of an alignment value . Feasibility is monotone over those aligned candidates, so the algorithm can binary-search for the smallest feasible instead of scanning every size.
To avoid enumerating every possible set of tensors that induces the fully-contained-shard case, the paper sorts tensors by element count and considers prefixes, yielding a 2-approximation. The reported planning runtime is below 0.3 seconds across the evaluated cases. Once this one-time plan exists, the next problem is execution: how do tensors and collectives use the planned memory without copying data into temporary communication buffers?
The planner gives every tensor a stable interval in the grouped buffer. DBuffer makes that interval the tensor’s actual backing storage rather than a staging area that data must be copied into before a collective.
A DBuffer represents a distributed global buffer over the device topology, with a sharding specification describing how that buffer is laid out. Each RaggedShard tensor is mapped to its planned slice. Because the mapping persists, the tensor can point directly at that memory before communication and directly at the corresponding result afterward. The Copy-In and Copy-Out steps that appeared in FSDP2 are unnecessary on this path.
This also makes in-place redistribution possible. An AllGather can transform the distributed buffer from a sharded form into a replicated form without first constructing separate per-parameter communication tensors. The same abstraction extends to multidimensional device topologies; the paper describes gradient reduction by redistributing a DBuffer between placements using ReduceScatter and AllReduce.
DBuffer also groups computation that would otherwise be fragmented across tensors. Before a collective, parameters may each require small operations such as add, scale, zero, or copy. Launching those kernels separately creates overhead and can delay communication. When several tensors need the same operation, DBuffer can fuse the work at the group level.
Memory management is handled at the same granularity. Batched allocations reduce fragmentation, and explicit stream dependencies make deallocation timing more predictable. So DBuffer is doing more than removing a memcpy: it gives the planner’s layout a persistent runtime representation, then uses that representation to reduce copies, kernel-launch fragmentation, and allocator fragmentation together.
A system like this is much easier to adopt if structure-aware sharding does not force the rest of the distributed stack to be rewritten. The authors therefore implement RaggedShard as another DTensor placement, alongside familiar placements such as Replicate, Partial, and Shard(d).
That matters when FSDP is combined with other kinds of parallelism. Tensor Parallelism may already shard a tensor by rows or columns, and Expert Parallelism may shard along the expert dimension. RaggedShard then has to apply its own uneven split without accidentally cutting across the dimension that another placement already owns.
The paper handles Shard(0) with a dedicated StridedRaggedShard placement that carries the reordering and stride metadata needed to reconstruct the full tensor correctly. For Shard(d) where , the RaggedShard granularity is adjusted using a least-common-multiple with the relevant tensor stride. The practical purpose is simple: a ragged FSDP boundary should never cut through the unit established by the other sharding dimension.
Because the result is still a DTensor, checkpointing can reuse the existing DTensor-based stack, including PyTorch Distributed Checkpoint and its communication-free sharded checkpointing path. This is a useful engineering consequence of extending an existing abstraction instead of creating an isolated parallel tensor type.
The user-facing API stays similarly conservative. veScale-FSDP is implemented as a replacement backend for FSDP2 while preserving PyTorch’s fully_shard interface. Model code can keep the same high-level FSDP contract while RaggedShard, the planner, and DBuffer change how the states are represented and moved underneath it.
That leaves the empirical question: after adding this extra flexibility and planning machinery, does the system actually train models faster and with less memory than the existing FSDP implementations?
The paper’s main end-to-end comparison uses 1,024 GPUs and three representative workloads: an internal 160B Mixture-of-Experts model, GPT-OSS-120B, and LLaMA-3-70B. The baselines are DeepSpeed ZeRO, PyTorch FSDP1, PyTorch FSDP2, and Megatron-FSDP, all configured for ZeRO-3-style sharding with mixed precision.
The biggest throughput differences appear on the MoE workloads. veScale-FSDP is reported to be 11–66% faster than the baselines there. That is consistent with the design pressure we have been following: sparse expert computation leaves less per-GPU compute available to hide AllGather and ReduceScatter, so padding, copies, and poorly utilized communication become more visible in the iteration time.
On LLaMA-3-70B, where the workload is dense, the margin is smaller. The paper reports roughly a 5% throughput advantage over DeepSpeed, FSDP1, and FSDP2, with veScale-FSDP slightly ahead of Megatron-FSDP. The claimed sources are the same mechanisms discussed earlier: better communication overlap, zero-copy DBuffer collectives, flexible granularity that avoids unnecessary padding, and aligned communication buffers.
Memory moves in the same direction. Across the benchmarks, veScale-FSDP reports 16–30% lower peak reserved memory. The authors attribute this to deterministic, batched DBuffer allocation and explicit stream-dependency management, which reduce fragmentation and make buffers reusable sooner. Relative to FSDP2’s per-parameter eager allocation, the paper reports a further 12% reduction from the batched policy in its analysis.
Lower reserved memory is useful even when a job technically fits. Near the limit, PyTorch’s caching allocator can be forced to issue device frees that synchronize with the driver, stalling training. Reducing fragmentation therefore affects both capacity and iteration efficiency.
These percentages belong to the paper’s specific models, GPU counts, and configurations; they should not be read as universal constants. The stronger claim is architectural: the costs that RaggedShard, planning, and DBuffer were designed to remove are large enough to show up in end-to-end training, particularly for communication-sensitive MoE models.
The next experiments test whether the same design survives when scale itself changes. For weak scaling, the paper trains an internal 800B-parameter MoE model from roughly 1K to 8K GPUs while increasing the global workload so each GPU keeps a similar amount of work. Throughput stays near the linear-scaling reference across the tested input sizes.
Strong scaling is harder because the global batch stays fixed while more GPUs divide the same work. With a 120M-token global batch, the paper reports linear scaling up to 10K GPUs. With a smaller 16M-token global batch, moving from 1K to 8K GPUs produces a 3.4× throughput increase rather than 8×. The reason is expected: as each GPU receives fewer tokens, parameter AllGather and gradient ReduceScatter take a larger fraction of the iteration. The authors use cross-node Expert Parallelism at very large scale to reduce FSDP communication, trading some additional token-exchange and compute inefficiency for lower collective cost.
The system also scales model size at fixed hardware. On 1K GPUs, the evaluated internal MoE models grow from 400B to 2.4T parameters without reported performance degradation. Model FLOPS Utilization, or MFU, slightly improves as the models get larger because the added computation gives the GPUs more useful arithmetic relative to overhead.
The more direct test of RaggedShard’s flexibility is 8-bit Adam. The implementation quantizes optimizer statistics in 32×32 blocks and uses a 32-row RaggedShard granularity for matrix parameters. Because every quantization block stays wholly inside one local shard, each device can quantize its own state without exchanging scaling-factor metadata. The paper says this required only a few lines of veScale-FSDP-specific code.
Muon stresses a different structural requirement: its matrix-sign preconditioner needs the complete original 2D parameter matrix. For a 2D weight , the optimizer obtains its gradient , applies the momentum update to produce , remembers its original placement , and chooses a load-balanced root . A DTensor redistribution makes the full tensor materialize only on ; other ranks effectively have no matrix update to compute. The root runs the Newton–Schulz iteration, then is redistributed back to and applied to .
That implementation can overlap redistribution with computation because it is expressed through the DTensor interface rather than hand-written per-parameter communication. The optimized Muon setup in the paper reaches 47.3% MFU on 256 Hopper GPUs. These optimizer experiments are separate from the large-scale scaling runs, but together they test both promises of RaggedShard: preserve unusual tensor structure and remain usable inside a production distributed runtime.
The ablations make the division of labor between the three components unusually clear. RaggedShard is the semantic layer. In the 8-bit-Adam experiment, removing it is marked N/A rather than assigned a slower throughput number, because preserving 32×32 block semantics would then require intrusive changes to model and optimizer tensors or custom collectives for metadata and state gathering. In other words, RaggedShard is what makes the workload naturally expressible in the first place.
The planner is what turns that flexibility into an efficient communication layout. For 1× and 16× row granularities on the tested DeepSeek-V3-671B and GPT-OSS-120B configurations, padding stays below 3% across the evaluated FSDP sizes. When the granularity grows to 128× rows, the constraints become harder: GPT-OSS sees step-like padding spikes up to 18%. The paper attributes this to the interaction between large block granularity, collective alignment, and GPT-OSS’s choice to fuse all experts into a single parameter tensor. The effective shard size has to jump to compatible least-common-multiple boundaries.
The ablation shows how expensive it is to ignore that planning. With the full system normalized to 100% throughput, disabling the planning algorithm drops the result to 65.4%. Without block-preserving placement, the system falls back to redistribution to assemble optimizer states before per-block quantization, introducing substantial extra communication. The planner itself is cheap by comparison: its one-time runtime remains below 0.3 seconds in the reported experiments.
DBuffer supplies a smaller but still measurable additional gain. Disabling only DBuffer reduces normalized throughput to 92.8%, a 7.2% loss attributed to Copy-In and Copy-Out around collectives. Its batched memory policy also contributes to the lower peak reserved memory seen in the end-to-end results.
The authors close with three lessons from deploying veScale-FSDP on workloads above 10K GPUs. First, smaller-scale profiling can predict large-scale FSDP performance surprisingly well when the profiling setup exercises comparable network topology, collective algorithms, and bandwidth saturation. Their production practice is to profile around 64 GPUs and extrapolate, while using hierarchical parallelism to keep collective groups from becoming unmanageably large.
Second, extending DTensor pays off beyond code reuse. By making RaggedShard a placement inside an established abstraction, veScale-FSDP can compose with Tensor Parallelism, Expert Parallelism, and distributed checkpointing instead of rebuilding those capabilities. Third, separating model definition from system-level parallelization lets researchers change architectures without rewriting performance machinery tied to a specific model implementation.
Taken together, the design is easier to understand as three layers than as one FSDP optimization. RaggedShard expresses the structure that must survive sharding. The planner finds a compact, aligned way to place that structure across devices. DBuffer makes the resulting layout the actual zero-copy communication storage. The paper’s reported 5–66% throughput gains and 16–30% memory reductions are the empirical consequence of those layers working together, while the optimizer examples show why the extra flexibility matters in the first place.
