Source-linked AI summary
veScale-FSDP: Flexible and High-Performance FSDP at Scale
Zezhou Wang, Youjie Li, Zhiqi Lin, Jiacheng Yang, Cong Xie, Guanyu Feng, Zheng Zhong, Ziyue Huang, Hongyu Zhu, Zhi Zhang, Yanghua Peng, Xin Liu
TL;DR
Existing FSDP systems use rigid sharding that conflicts with block-structured training and introduce efficiency costs at very large scale. veScale-FSDP combines RaggedShard with structure-aware planning and zero-copy communication. It supports block-wise quantization and non-element-wise optimizers while delivering higher throughput, lower memory usage, and scaling to tens of thousands of GPUs.
Problem
Existing element-wise or row-wise FSDP sharding cannot reliably preserve the block structures required by block-wise quantization and non-element-wise optimizers, while also incurring communication and memory overheads at scale.
Method
veScale-FSDP combines arbitrary-granularity RaggedShard tensors with a structure-aware planning algorithm and Distributed Buffer support for zero-copy access.
Results
5∼66% higher throughput and 16∼30% lower memory usage were achieved across dense and sparse LLMs while scaling efficiently to tens of thousands of GPUs.
Takeaways & Limitations
veScale-FSDP natively accommodates Muon and block-wise quantization methods such as 8-bit Adam within a scalable FSDP system.
Takeaways & Limitations
The planning problem is NP-hard, so veScale-FSDP uses a polynomial-time heuristic rather than exact large-scale optimization.
Abstract
from arXiv · showhide
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.
1 Introduction
FSDP offers memory-efficient large-scale training, but fixed sharding formats limit structure-aware methods and current systems incur substantial performance and memory overheads. veScale-FSDP addresses both issues with flexible sharding, structure-aware planning, and zero-copy communication.
- FSDP is widely used for large-model training because it efficiently distributes model and optimizer states while remaining decoupled from model architecture.
- Existing FSDP sharding often misaligns with block structures required by Shampoo, Muon, and block-wise quantization.
- Misaligned boundaries force model or system developers to add intrusive code changes, boundary checks, padding, and communication logic.
- Current systems incur fragmented or slow collectives, tensor-copy overhead, padding costs, and inefficient memory management, especially beyond 10K GPUs.
- RaggedShard supports arbitrary sharding granularity and custom block sizes while composing with PyTorch DTensor formats.
- veScale-FSDP combines structure-aware planning and Distributed Buffer primitives to improve communication efficiency, enable zero-copy access, and reduce memory fragmentation.
- 5∼66% higher throughput and 16∼30% lower memory usage were achieved across dense and sparse LLMs while scaling to tens of thousands of GPUs.
2 Background and Motivation
Structure-aware training requires sharding that preserves matrix and quantization-block structure, but standard DTensor and FSDP formats remain too rigid. Existing systems also introduce copy, padding, communication, and load-balance costs.
- Shampoo and Muon update parameters as original 2D matrices, requiring matrix gathering before updates and redistribution afterward.
- Block-wise quantization keeps each quantization block on one device to avoid metadata exchange and preserve efficiency.
- DTensor supports Shard(dim), Replicate, and Partial placements plus redistribution, but Shard cannot express block-wise or uneven sharding.
- FSDP2 exposes per-parameter Shard(0) DTensors but retains fixed even sharding and adds parameter-copy overhead from interleaved memory addresses.
- Megatron-FSDP avoids FSDP2 copy overhead through concatenated sharding, but padding can increase memory and communication costs.
- Fixed row-wise sharding may split quantization blocks, while padding inside buffers can break contiguity and create interleaved copies.
3 Overview
veScale-FSDP rearchitects the FSDP2 backend while preserving PyTorch’s native fully_shard API. Its RaggedShard format and planned layouts support complex models, structure-aware optimizers, and efficient communication.
- veScale-FSDP supports sophisticated large models, including sparse MoE structures, and structure-aware optimizers with non-element-wise operators.
- RaggedShard supports arbitrary sharding granularities over contiguous storage and arbitrary distributions across devices, preserving single-device semantics for complex operators.
- The system groups RaggedShard tensors and rearranges their layouts with a planning algorithm derived from an NP-hard optimization problem.
- A Distributed Buffer maps planned layouts to a zero-copy communication path with minimal overhead.
4 RaggedShard for Flexibility
RaggedShard extends DTensor with arbitrary block granularity and distribution, addressing the structural limitations of element-wise and evenly row-wise sharding. Block-wise RaggedShard aligns shard boundaries with quantization blocks and composes with existing parallelism placements.
- Element-wise sharding can break non-element-wise operations, complicate redistribution, and misalign block-wise FP8 quantization with shard boundaries.
- Row-wise even sharding enables non-element-wise computation and dimension redistribution but does not guarantee alignment with block boundaries.
- RaggedShard supports arbitrary sharding granularity in contiguous memory and arbitrary distributions of non-shardable blocks across devices.
- Block-wise RaggedShard uses customizable tensor-block shapes, enabling non-element-wise computation, efficient redistribution, and exact alignment with block-wise quantization.
- RaggedShard composes with replicated, partial-value, and evenly sharded DTensor placements for strategies such as Tensor Parallelism and Expert Parallelism.
- RaggedShard reuses DTensor-based checkpointing stacks and inherits communication-free sharded checkpointing optimizations.
5 Grouped RaggedShard for Performance
veScale-FSDP groups RaggedShard tensors through structure-aware communication planning and a Distributed Buffer. The approach preserves block and tensor contiguity, balances device buffers, and reduces communication and memory overhead.
- Naive grouping can split blocks, insert padding within tensors, and create unequal per-device buffers, reducing communication efficiency.
- The planner permutes tensors and pads between them to preserve block boundaries, tensor contiguity, and balanced per-device communication buffers.
- The layout problem minimizes uniform per-device buffer size subject to non-sharded-block, contiguous-memory, and balanced-load constraints.
- Because the optimization is NP-hard and large deployments may span hundreds of thousands of devices, veScale-FSDP uses a polynomial-time heuristic instead of impractical ILP solving.
- The heuristic-guided dynamic-programming algorithm considers default, block-size, and tensor-shape orders before placing tensors into the smallest feasible global buffer.
- Distributed Buffer: DBuffer provides global buffer semantics over an N-dimensional device topology and executes group-level operators for grouped tensors.
- Distributed Buffer: DBuffer enables zero-copy access before and after communication while fusing identical kernels across tensors to reduce blocking time.
6 Evaluation
The evaluation measures veScale-FSDP against established FSDP systems across end-to-end performance, scalability, optimizer support, planning quality, and component contributions. Across these studies, veScale-FSDP improves throughput and memory efficiency, scales to large GPU and model sizes, and supports structure-aware training.
- Evaluation setup: The evaluation compares veScale-FSDP with DeepSpeed ZeRO, PyTorch FSDP1 and FSDP2, and Megatron-FSDP on LLaMA-3-70B, GPT-OSS-120B, and an internal MoE model.The baselines use ZeRO-3 with mixed precision unless otherwise specified.
- End-to-end performance: 11∼66% higher throughput is achieved on MoE models, while LLaMA-3-70B is 5% faster than DeepSpeed, FSDP1, and FSDP2.The throughput gains are attributed to communication overlap, zero-copy collectives, and flexible sharding granularity that avoids padding overhead.
- End-to-end performance: 16–30% lower peak reserved memory is achieved across benchmarks through batched DBuffer allocation and deterministic memory management.The reported reductions include 12% less memory than FSDP2’s per-parameter eager allocation and 33% padding-related memory inflation for Megatron in MoE experiments.
- Scalability and composability: veScale-FSDP scales linearly with a 120M-token global batch up to 10K GPUs and delivers a 3.4× throughput gain from 1K to 8K GPUs at a 16M-token batch.Cross-node Expert Parallelism reduces FSDP communication time at larger scales, although token exchange and reduced kernel efficiency cause a performance drop at very large scales.
- Scalability and composability: 2.4T-parameter models can be trained on 1K GPUs without performance degradation, with slightly improved MFU as model size grows.The evaluation fixes the GPU count at 1K and scales models from 400B to 2.4T parameters.
- Planning quality: 1× and 16× row granularities keep padding overhead below 3% across all tested FSDP sizes for DeepSeek-V3 and GPT-OSS.With 128× rows, GPT-OSS shows step-like fluctuations with spikes up to 18% because it fuses all experts into one parameter tensor.
- Performance breakdown: Disabling DBuffer reduces throughput to 92.8%, while disabling the planning algorithm reduces it to 65.4%; disabling RaggedShard is reported as not meaningfully usable.RaggedShard is required to preserve 32 × 32 block-wise semantics without intrusive tensor changes or manually implemented collectives.
7 Lessons Learned
The deployment experience identifies three lessons: small-scale profiling can predict large-scale performance under matching network conditions, abstractions should build on DTensor, and model definitions should remain decoupled from system optimization.
- Lesson-1: Small-scale profiling can estimate large-scale FSDP performance when computation and communication behavior remain comparable across GPU counts.The extrapolation assumes similar network topology, collective algorithms and protocols, plus workloads large enough to saturate bandwidth.
- Lesson-2: RaggedShard is implemented as an optional DTensor placement, enabling composition with Tensor Parallelism, Expert Parallelism, and existing training tools.DTensor provides the underlying abstraction for integrating established parallelization strategies.
- Lesson-3: Decoupling model definition from system-level parallelization lets researchers modify architectures while retaining linear performance scaling.The lesson contrasts this design with frameworks that tightly couple parallelization optimizations to model code.
8 Conclusion
veScale-FSDP combines RaggedShard with structure-aware planning to provide flexible and high-performance large-scale training. Experiments report higher throughput, lower memory usage, and efficient scaling to tens of thousands of GPUs while integrating with techniques such as Muon.
- Conclusion: veScale-FSDP combines RaggedShard and a structure-aware planning algorithm to improve flexibility, performance, and GPU utilization.The system integrates emerging techniques such as Muon optimizers.
- Conclusion: 5∼66% higher throughput and 16∼30% lower memory usage are reported than existing systems while scaling efficiently to tens of thousands of GPUs.These are the paper’s aggregate reported performance and scalability outcomes.