Fully Sharded Data Parallel (FSDP), also known as Zero Redundancy Optimizer (ZeRO), is widely used for large-scale model training, because of its memory efficiency and minimal intrusion on model code. However, existing FSDP systems rely on fixed element-wise or row-wise sharding formats that conflict with block-structured computations. As a result, they struggle to support modern structure-aware training methods, including block-wise quantization and non-element-wise optimizers such as Shampoo and Muon. In addition, today's implementations incur communication and memory overheads that degrade efficiency at the scale of tens of thousands of GPUs. We introduce veScale-FSDP, a novel FSDP system that combines RaggedShard, a flexible sharding format, with a structure-aware planning algorithm to deliver both flexibility and performance. veScale-FSDP enables zero-copy FSDP communications and natively supports block-wise quantization and non-element-wise optimizers, achieving 5% to 66% higher throughput and 16% to 30% lower memory usage than existing FSDP systems, while scaling efficiently to tens of thousands of GPUs.
Figures & tables
Figure 1: Distributed Tensor (DTensor) for flexible communication and computation. The figure shows an example of DTensors executing a sharded matrix multiplication ( matmul ) on a device. A row-sharded ( Shard(0) ) DTensor is first redistributed to a replicated ( Replicate ) DTensor so that it can be multiplied with a column-sharded ( Shard(1) ) DTensor, producing a column-sharded ( Shard(1) ) DTensor. The blue part in each DTensor indicates the materialized local tensor on that device.
AllGather path
ReduceScatter path
AllGather
Copy-Out
ReduceScatter
Copy-In
Shard(0)
43.71 ms
5.22 ms
94.24 ms
12.37 ms
Shard(1)
44.35 ms
13.72 ms
95.36 ms
23.14 ms
Table 1: Interleaved copy overhead in FSDP2 for GPT-OSS-120B on 64 H800 GPUs. The table reports the time of interleaved Copy-Out relative to AllGather on the AllGather path, and the time of interleaved Copy-In relative to ReduceScatter on the ReduceScatter path. Shard(0) is the default sharding format for each parameter, while Shard(1) is used when Shard(0) incurs large padding.
Figure 2: The copy overhead in the FSDP2 design. Each parameter is represented as a Shard(0) DTensor and is evenly sharded across devices. Before AllGather, the input communication buffer on each device needs to copy in the local shards of all parameters. After AllGather, the output communication buffer stores each parameter in interleaved memory addresses rather than a contiguous address. As a result, FSDP2 must copy out each parameter from the output buffer to materialize the full tensor with a contiguous address needed for computation. The ReduceScatter path is a reversed one and incurs interleaved Copy-In overhead.
Figure 3: veScale-FSDP overview.
Figure 4: Flexibility comparison of different sharding formats.
Figure 5: Composability of RaggedShard with existing evenly sharded DTensor placements in 2D parallelism schemes such as FSDP × EP (Expert Parallelism).
Figure 6: Communication layouts for grouped RaggedShard DTensors. (a) shows a naive strategy that concatenates tensors directly into the communication buffer. This can split blocks across shard boundaries, introduce padding within tensors, and produce unequal per-device buffer sizes, leading to inefficient communication. (b) shows our structure-aware planner that carefully plans the communication layout – first permutes tensors and then inserts padding between tensors rather than within them. This preserves non-sharded blocks, keeps each tensor contiguous in memory, and balances per-device communication buffers, leading to efficient communication.
Figure 7: Distributed Buffer ( DBuffer ) for high-performance communication. A 2D DBuffer for parameter AllGather is shown; conversely, a 2D DBuffer redistributing from (Partial, Partial) to (Replicate, Shard) implements 2D gradient reduction with ReduceScatter and AllReduce.
Figure 8: FSDP training performance. Top row: normalized aggregate throughput (tokens/s). Bottom row: peak per-GPU memory (GB). We sweep FSDP (ZeRO-3) at 128/256 GPUs and HSDP with 2- and 4-way replication (2256, 4256 GPUs).
Figure 9: Scalability of veScale-FSDP. (a) Weak scaling to 8K GPUs. (b) Strong scaling to 8K GPUs. (c) Strong scaling at larger scale with normalized throughput and normalized GPU count. (d) Model scaling on 1K GPUs up to 2.4T parameters.
Figure 10: Training convergence with veScale-FSDP on 64 GPUs (FSDP size 64) for 8-bit Adam and distributed Muon.
Figure 11: Padding overhead (extra padding bytes over total parameter size) of RaggedShard communication. Lines compare the sharding granularities (1 × /16 × /128 × parameter row size) versus FSDP sharding size (number of GPUs).
veScale-FSDP Component
Normalized Throughput
Combined
100.0%
Disable DBuffer only
92.8%
Disable Planning Algorithm only
65.4%
Disable RaggedShard only
N/A
Table 2: Component ablation for 8-bit Adam, reported as normalized throughput after disabling each component independently. N/A indicates that the corresponding configuration is not meaningfully runnable without intrusive modifications to the model or optimizer code or manual management of custom collectives.