veScale-FSDP: Flexible and High-Performance FSDP at Scale

veScale-FSDP is a redesigned, high-performance FSDP system that introduces a flexible RaggedShard format and structure-aware planning to overcome the limitations of existing systems in supporting block-wise quantization and non-element-wise optimizers, thereby achieving significantly higher throughput, lower memory usage, and efficient scaling to tens of thousands of GPUs.

The Gist

Imagine you are trying to build a massive skyscraper (a giant AI model) using a team of thousands of construction workers (GPUs).

In the past, the standard way to organize this team was FSDP (Fully Sharded Data Parallel). Think of FSDP as a strict foreman who says: "We are going to cut the blueprints into tiny, equal-sized square tiles. Every worker gets exactly one tile. If a tile needs to be glued to another, we have to stop, gather all the tiles back together, glue them, and then cut them up again."

This works okay for simple buildings, but modern AI models are like complex, futuristic architecture with weird shapes, curved walls, and special materials. The old "square tile" method causes two big problems:

1. The "Square Peg in a Round Hole" Problem: New, fancy construction techniques (like "block-wise quantization" or "matrix optimizers") need to work on specific chunks of the blueprint (like a whole 128x128 block). The old foreman keeps cutting these blocks in half, forcing workers to stop and reassemble them constantly.
2. The "Wasted Space" Problem: Because the tiles must be perfectly equal, workers often end up with empty space on their desks (padding) just to make the math work. This wastes memory and slows everyone down.

Enter veScale-FSDP.

The authors of this paper invented a new system that acts like a smart, flexible project manager instead of a rigid foreman. Here is how it works, using simple analogies:

1. The "Ragged Shard" (The Flexible Puzzle)

Instead of forcing every worker to hold a perfect square tile, RaggedShard lets workers hold irregular puzzle pieces.

• The Old Way: If a blueprint section is 100 units wide and you have 3 workers, the old system cuts it into 33, 33, and 34. But if the "fancy construction" needs a 32-unit block, the cut right through the middle of the block, ruining it.
• The veScale Way: RaggedShard says, "Worker A, you take the first 32-unit block. Worker B, you take the next 64-unit block. Worker C, you take the rest."
• The Result: The workers can now use the special construction techniques (like block-wise quantization) without ever having to stop and reassemble the pieces. The blocks stay intact, just like a puzzle piece that fits perfectly in your hand.

2. The "Planning Algorithm" (The Traffic Controller)

Now that everyone has different-sized puzzle pieces, how do we pass them around without causing a traffic jam?

• The Problem: If you just line up the pieces randomly, you might end up with a huge gap (padding) between pieces, or a piece might get split across two workers' desks, requiring a messy hand-off.
• The Solution: The paper introduces a Planning Algorithm. Imagine a traffic controller who looks at all the puzzle pieces and says: "Okay, put the big pieces here, the small pieces there, and arrange them so that when we pass them down the line, they fit together perfectly with zero gaps."
• The Magic: This algorithm solves a complex math problem (NP-hard) to find the most efficient arrangement in milliseconds, ensuring no space is wasted and no time is lost shuffling pieces around.

3. The "Distributed Buffer" (The Super-Highway)

Once the pieces are arranged, they need to move between workers.

• The Old Way: Workers often had to copy their pieces onto a temporary clipboard (memory) before passing them, which is slow and messy.
• The veScale Way: They built a Distributed Buffer (DBuffer). Think of this as a super-highway where the pieces are already parked in the right spots. When a worker needs a piece, they just point to it. No copying, no moving, no waiting. It's "zero-copy" access.

Why Does This Matter? (The Results)

Because of these three innovations, the team can build their skyscraper much faster and with fewer resources:

• Speed: They are 5% to 66% faster than previous systems. It's like going from a slow, stop-and-go commute to a high-speed train.
• Memory: They use 16% to 30% less memory. This is like fitting the same amount of furniture into a smaller apartment by arranging it perfectly, rather than needing a warehouse.
• Scale: They can now train models on tens of thousands of GPUs without the system crashing or running out of memory.

The Bottom Line

veScale-FSDP is a system upgrade that stops treating AI models like rigid, uniform blocks of data. Instead, it treats them like flexible, organic structures. By allowing workers to hold irregular pieces (RaggedShard), arranging them perfectly (Planning), and moving them instantly (DBuffer), it unlocks the ability to train the next generation of super-intelligent AI models that were previously too difficult or expensive to build.

It's the difference between trying to build a castle with only square Lego bricks versus having a set of custom-shaped bricks that fit together perfectly.

Technical Summary

Technical Summary: veScale-FSDP

1. Problem Statement

Fully Sharded Data Parallel (FSDP), also known as ZeRO, is a cornerstone for training large-scale models due to its flexibility and minimal intrusion on model code. However, existing FSDP systems (e.g., DeepSpeed ZeRO, PyTorch FSDP1/2, Megatron-FSDP) face critical limitations when scaling to modern, structure-aware training paradigms and massive GPU clusters (10K+ devices):

• Incompatibility with Structure-Aware Training: State-of-the-art models utilize non-element-wise optimizers (e.g., Shampoo, Muon) and block-wise quantization (e.g., DeepSeek-V3). These methods require atomic tensor blocks (e.g., 2D matrices or specific block sizes) to function correctly.
  • Current Limitation: Existing FSDP systems shard parameters either element-wise or row-wise. This creates sharding boundaries that misalign with the required block structures, forcing developers to either intrusively modify model/optimizer code to handle boundary checks and padding or suffer from complex, inefficient communication logic.
• Performance and Memory Inefficiency:
  • Communication Overhead: Current implementations suffer from fragmented AllGather operations, unaligned communication buffers, and excessive data copying (interleaved copy-in/copy-out) between sharded and full tensors.
  • Memory Fragmentation: Inefficient memory management leads to high peak memory usage, causing Out-of-Memory (OOM) errors or forcing over-provisioning of GPU resources.
  • Scalability Limits: These inefficiencies hinder scaling to tens of thousands of GPUs and trillion-parameter models.

2. Methodology

The authors introduce veScale-FSDP, a redesigned FSDP system that couples a flexible sharding format with a structure-aware planning algorithm and a high-performance buffer primitive.

A. RaggedShard: Flexible Sharding Format

To address the rigidity of element/row-wise sharding, veScale-FSDP introduces RaggedShard, a new DTensor placement format.

• Arbitrary Granularity: Unlike fixed formats, RaggedShard allows sharding at any granularity (e.g., specific block sizes, rows, or planes) and supports arbitrary distribution of these blocks across devices.
• Block-Awareness: It enables Block-wise RaggedShard, where tensors are partitioned into customizable 2D blocks. This ensures that quantization blocks or optimizer matrix blocks align perfectly with shard boundaries, eliminating the need for cross-device synchronization of block metadata.
• Composability: RaggedShard is designed to compose seamlessly with existing DTensor placements (e.g., Tensor Parallelism, Expert Parallelism) without breaking the abstraction.

B. Structure-Aware Planning Algorithm

Grouping RaggedShard tensors for collective communication is non-trivial. Naive grouping leads to:

1. Sharded Blocks: Breaking atomic blocks across devices.
2. Non-Contiguous Memory: Requiring padding within tensors, causing copy overhead.
3. Imbalanced Load: Causing network underutilization.

Solution: The authors formulate the grouping problem as an NP-hard optimization problem (reducible from the Partition problem).

• Heuristic Approach: Instead of solving the ILP (which is too slow for large scales), they propose a polynomial-time dynamic programming (DP) algorithm guided by permutation heuristics.
• Strategy: The algorithm permutes tensors and pads between them (rather than within) to balance load and align block boundaries. It exploits the regularity of Transformer models (e.g., consistent block sizes across layers) to find near-optimal solutions quickly (<0.3s runtime).

C. Distributed Buffer (DBuffer)

To maximize hardware efficiency, veScale-FSDP introduces DBuffer, a global buffer primitive.

• Zero-Copy Access: DBuffer provides a persistent address mapping for RaggedShard tensors, enabling zero-copy access to data before and after communication.
• Kernel Fusion: It fuses identical operations (e.g., scaling, zeroing) across multiple tensors into a single kernel launch, reducing blocking time.
• In-Place Operations: Supports in-place communication and computation to minimize memory footprint and fragmentation.
• Batched Allocation: Uses batched memory management to reduce fragmentation and ensure deterministic deallocation.

3. Key Contributions

1. RaggedShard Abstraction: A novel sharding format that natively supports arbitrary block sizes and uneven distributions, enabling structure-aware training (block-wise quantization, matrix optimizers) without intrusive code changes.
2. Optimization Algorithm: A practical, polynomial-time planning algorithm that solves the NP-hard tensor grouping problem, minimizing padding overhead and ensuring balanced communication loads.
3. DBuffer Primitive: A high-performance communication buffer that achieves zero-copy access, kernel fusion, and reduced memory fragmentation.
4. Production-Ready System: A fully functional FSDP implementation that replaces PyTorch FSDP2's backend while maintaining API compatibility (fully_shard), validated on production workloads.

4. Experimental Results

The system was evaluated on models ranging from 70B to 2.4T parameters across 128 to 10,000 GPUs.

• Throughput Improvement: veScale-FSDP achieves 5% to 66% higher throughput compared to DeepSpeed ZeRO, PyTorch FSDP1/2, and Megatron-FSDP.
  • Example: On MoE models, it is 11–66% faster due to optimized communication overlapping and elimination of padding overhead.
• Memory Efficiency: It reduces peak memory usage by 16% to 30%.
  • Mechanism: Deterministic batched memory management and elimination of interleaved copy overheads prevent the caching allocator from inflating peak memory.
• Scalability:
  • Successfully scales to 10,000 GPUs with near-linear weak and strong scaling.
  • Enables training of 2.4 Trillion parameter models on 1,000 GPUs without performance degradation.
• Support for Advanced Optimizers:
  • 8-bit Adam: Natively supports block-wise quantization without manual collectives.
  • Muon Optimizer: Enables distributed Newton-Schulz updates by allowing uneven sharding and root-rank computation, achieving 47.3% MFU on 256 GPUs.
• Planning Overhead: The planning algorithm adds negligible overhead (<0.3 seconds) during initialization.

5. Significance

veScale-FSDP represents a significant leap in distributed training infrastructure for Large Language Models (LLMs):

• Bridging the Gap: It resolves the fundamental conflict between the rigid sharding of current FSDP systems and the block-structured requirements of next-generation training techniques (quantization, advanced optimizers).
• Developer Productivity: By decoupling model definition from system optimization, it allows researchers to implement complex architectures and optimizers using standard PyTorch APIs without "hacking" system code.
• Industrial Viability: The system has been battle-tested in production at ByteDance, demonstrating that it can handle the extreme scale (10K+ GPUs) and memory constraints of modern AI training, setting a new standard for efficiency and flexibility in the field.
• Open Source: The core RaggedShard code is open-sourced, contributing to the broader PyTorch ecosystem (with planned integration into official PyTorch roadmaps).