veScale-FSDP: Scalable FSDP with RaggedShard

• veScale-FSDP is a scalable implementation of Fully Sharded Data Parallelism that leverages the novel RaggedShard scheme for block-aligned, structure-aware sharding across large GPU clusters.
• It integrates a dynamic structure-aware communication planner and unified buffer management to minimize padding and optimize throughput with lower memory overhead.
• The system supports advanced features like block-wise quantization and non-element-wise optimizers, enhancing scalability for large language models and neural network verification tasks.

veScale-FSDP is a modern, scalable implementation of Fully Sharded Data Parallelism (FSDP), designed to efficiently distribute model and optimizer states across large accelerator clusters while natively supporting both structure-aware tensor manipulations (such as block-wise quantization and non-element-wise optimizers) and highly scalable neural network verification workloads. The centerpiece of veScale-FSDP is the “RaggedShard” sharding scheme, which enables flexible, block-aligned sharding without intrusive code modifications, and a structure-aware planning algorithm for optimal communication layout. It delivers higher throughput and lower memory usage than prior FSDP or ZeRO systems, while seamlessly scaling to tens of thousands of GPUs (Wang et al., 25 Feb 2026). This system is also foundational for scalable formal verification of DNNs via precise, bitwise-identical sharding of weight tensors and efficient orchestration of bound-propagation algorithms (Vorobyov et al., 8 Jun 2026).

1. Motivation and Limitations of Standard FSDP

Fully Sharded Data Parallelism (FSDP), or ZeRO, is an established paradigm for training large models by partitioning model and optimizer states across devices. Traditional FSDP implementations rely on element-wise or row-wise partitioning schemes, which fragment tensors arbitrarily and thus fail to support structure-aware computations such as block-wise quantization or matrix-level optimizers. Furthermore, these schemes lead to inefficient communication (fragmented collectives with manual copy overheads), padding-induced memory bloat, and increased allocator fragmentation, all limiting their scalability and utility for emerging structured training and verification methods (Wang et al., 25 Feb 2026, Vorobyov et al., 8 Jun 2026).

2. RaggedShard: Flexible Block-Aligned Sharding

veScale-FSDP introduces RaggedShard, a sharding scheme that supports partitioning tensors into contiguous blocks (“atomic shards”) of user-specified shape (e.g., 32×32), with arbitrarily non-uniform assignment across devices. Given a 2D tensor $T\in\mathbb{R}^{R\times C}$ and block granularity $g=(r_b, c_b)$, RaggedShard views $T$ as a grid of atomic blocks $$B_{p,q}=T[p\cdot r_b:(p+1)r_b,\;q\cdot c_b:(q+1)c_b]$$. Each device $k$ is assigned a subset $\mathcal{B}_k$ of these blocks, and the union over all devices covers all blocks. This enables alignment of quantization and optimizer metadata to shard boundaries, eliminates cross-device metadata exchange, and allows entire block-structured operations to be performed locally on-device (Wang et al., 25 Feb 2026).

Key properties:

• 1D (row-wise) sharding: $c_b = C$, $r_b=1$; elements as rows, matching FSDP2.
• Element-wise: $r_b = c_b = 1$, matching ZeRO/FSDP1.
• Full-matrix: single device is assigned all blocks; enables block-gather for matrix optimizers.

3. Structure-Aware Communication Planning and Buffer Management

Efficient scaling depends on optimizing communication collectives. veScale-FSDP’s structure-aware planner frames packing of RaggedShard tensors into a per-device communication buffer as an NP-hard interval packing problem with modular block-alignment constraints. It uses a dynamic-programming plus heuristic algorithm to pack blocks, minimizing padding and maximizing collective efficiency. For each tensor $t$ with element count $g=(r_b, c_b)$0 and block size $g=(r_b, c_b)$1, optimal offsets $g=(r_b, c_b)$2 into the global buffer are selected to ensure block-aligned coverage, non-overlap, and minimum stride $g=(r_b, c_b)$3. The planner rapidly solves the layout problem (≤0.3s for hundreds of tensors in practice) and determines buffer alignment for collective operations (e.g., NCCL AllGather/ReduceScatter).

The “DBuffer” primitive then manages a global communication buffer per device, zero-copying RaggedShard intervals in and out, fusing per-tensor kernels for minimal allocator and synchronization overhead. This unified buffer ensures maximal compute/communication overlap and prevents fragmentation-induced slowdowns (Wang et al., 25 Feb 2026).

4. Integration with Block-Wise Quantization and Non-Element-Wise Optimizers

RaggedShard natively supports block-wise INT8 quantization (such as 32×32-block 8-bit Adam) by ensuring each device holds complete quantization blocks, performs local scaling/zero-point calculation, and accumulates quantized optimizer state with no need for inter-device metadata exchange. For non-element-wise optimizers (e.g., Shampoo, Muon), which operate on full matrices or require matrix-level preconditioning via Newton–Schulz iterations, RaggedShard gathers the full tensor to the designated root device using DTensor’s .redistribute, performs the matrix operation, and redistributes the result, all with alignment to existing block structure—critical for supporting LLMs and Mixture-of-Experts (MoE) architectures at scale (Wang et al., 25 Feb 2026).

5. Performance Results and Scalability Metrics

veScale-FSDP demonstrates superior throughput (5–66% higher) and reduced peak memory usage (16–30% lower per GPU) compared to DeepSpeed ZeRO, FSDP1/2, and Megatron-FSDP across multiple LLM and MoE benchmarks. For example, on LLaMA-3-70B, throughput is $g=(r_b, c_b)$4 that of ZeRO, while for large MoE models it reaches $g=(r_b, c_b)$5 (Wang et al., 25 Feb 2026). Weak scaling benchmarks (1K → 8K GPUs) yield per-GPU MFU $g=(r_b, c_b)$6, and strong scaling is linear up to 10K GPUs. Model scaling experiments verify that up to 2.4T parameters can be trained without throughput loss. Operator-level performance modeling confirms that $g=(r_b, c_b)$7, with compute-communication overlap optimizing efficiency to hardware limits.

Performance Comparison Table

Model	DeepSpeed ZeRO	Megatron-FSDP	veScale-FSDP
LLaMA-3-70B	1.00×	1.04×	1.05×
GPT-OSS-120B	1.00×	0.98×	1.11×
Internal MoE	1.00×	0.94×	1.66×

6. Application to Neural Network Verification

veScale-FSDP enables scalable formal DNN verification under the $g=(r_b, c_b)$8-CROWN and auto_LiRPA frameworks. For each weight tensor $g=(r_b, c_b)$9 (linear: $T$0, conv: $T$1), the output-channel dimension is sharded, with each GPU rank holding $T$2. Per-layer AllGather operations reconstruct the full $T$3 only for bound-propagation steps (either IBP forward or CROWN backward); otherwise only the shard $T$4 is resident. This ensures memory-efficient, bitwise-equivalent bound propagation relative to single-GPU baselines, attaining up to 90% baseline memory saving and 34–39% peak memory reduction for wide MLPs at $T$5 (Vorobyov et al., 8 Jun 2026). In convolutional architectures (e.g., CIFAR-100 ResNet-large, VNN-COMP 2024), complete verification is achieved with identical lower bounds and memory scaling. The dominant memory term in $T$6-CROWN+Branch-and-Bound verification is shown to be per-neuron $T$7 tensors (order $T$8), which remain unsharded—a critical direction for future scaling.

7. Limitations and Prospects for Future Work

Current limitations include slight planning complexity spikes for highly heterogeneous block sizes (hundreds of unique $T$9), persistent root-device bottlenecks for matrix-iteration optimizers (pinned full-matrix preconditioning in distributed Muon), and predominant memory usage by unsharded per-neuron tensors in neural verification contexts. Further limitations arise when $$B_{p,q}=T[p\cdot r_b:(p+1)r_b,\;q\cdot c_b:(q+1)c_b]$$0, as unsharded $$B_{p,q}=T[p\cdot r_b:(p+1)r_b,\;q\cdot c_b:(q+1)c_b]$$1 matrices and $$B_{p,q}=T[p\cdot r_b:(p+1)r_b,\;q\cdot c_b:(q+1)c_b]$$2 tensors dominate total memory, leading to diminishing returns; this motivates sharding these activation- or relaxation-related tensors. Planned directions include sharding $$B_{p,q}=T[p\cdot r_b:(p+1)r_b,\;q\cdot c_b:(q+1)c_b]$$3 tensors, integrating domain-parallel branch-and-bound for neural verification, exposing RaggedShard semantics to PyTorch and other frameworks, exploring pipelined/parallel matrix optimizer support, and exploiting hierarchical memory and communication/storage tradeoffs (Wang et al., 25 Feb 2026, Vorobyov et al., 8 Jun 2026).