--- title: 'veScale-FSDP: Scalable FSDP with RaggedShard' url: https://www.emergentmind.com/topics/vescale-fsdp type: topic --- # veScale-FSDP: Scalable FSDP with RaggedShard 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 [2602.22437]. 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 [2606.09377]. ## 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 [2602.22437], [2606.09377]. ## 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 [2602.22437]. 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 $e_t$ and block size $g_t$, optimal offsets $\ell_t, r_t$ into the global buffer are selected to ensure block-aligned coverage, non-overlap, and minimum stride $S$. 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 [2602.22437]. ## 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 [2602.22437]. ## 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 $1.05\times$ that of ZeRO, while for large MoE models it reaches $1.66\times$ [2602.22437]. Weak scaling benchmarks (1K → 8K GPUs) yield per-GPU MFU $\geq 92\%$, 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 $T_{\text{iter}}(m)\approx\max(T_{\text{comp}},T_{\text{comm}})$, 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 $\alpha, \beta$-CROWN and auto_LiRPA frameworks. For each weight tensor $W$ (linear: $[d_\text{out}, d_\text{in}]$, conv: $[C_\text{out}, C_\text{in}, k_H, k_W]$), the output-channel dimension is sharded, with each GPU rank holding $W^{(r)}$. Per-layer AllGather operations reconstruct the full $W$ only for bound-propagation steps (either IBP forward or CROWN backward); otherwise only the shard $W^{(r)}$ 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 $P=2$ [2606.09377]. 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 $\alpha, \beta$-CROWN+Branch-and-Bound verification is shown to be per-neuron $\alpha$ tensors (order $O(d h B)$), 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 $g_t$), 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 $P \gg 2$, as unsharded $A$ matrices and $\alpha$ tensors dominate total memory, leading to diminishing returns; this motivates sharding these activation- or relaxation-related tensors. Planned directions include sharding $\alpha/\beta$ 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 [2602.22437], [2606.09377]. Source: https://www.emergentmind.com/topics/vescale-fsdp