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

TL;DR

veScale-FSDP is a redesigned FSDP/ZeRO system from ByteDance built around a new sharding format, RaggedShard, that supports arbitrary sharding granularity and arbitrary per-device block counts — letting non-element-wise optimizers (Muon, Shampoo) and block-wise quantized training (DeepSeek-V3 style FP8, 8-bit Adam) run without intrusive model surgery. A planning algorithm rearranges RaggedShard tensors into a shared DBuffer primitive for zero-copy, aligned, balanced collectives. The system reports up to 566% higher throughput and 1630% lower memory than existing FSDP implementations on internal MoE training, and scales to tens of thousands of GPUs. Code is upstreamed in the volcengine/veScale repo, swapping in behind PyTorch’s fully_shard API.

Key claims

Existing element-wise (DeepSpeed, FSDP1) and row-wise (FSDP2, Megatron-FSDP) sharding formats cannot align shard boundaries with the atomic blocks required by Muon/Shampoo (full 2D matrix locality) or by block-wise FP8 quantization, forcing intrusive model/optimizer rewrites [§1, §2.3, §4].
RaggedShard extends PyTorch DTensor with a placement that supports (i) customizable atomic block shape (contiguous elements/rows/2D blocks) and (ii) arbitrary distribution of blocks across devices; block-wise RaggedShard generalizes all prior FSDP sharding formats as a special case [§4, Fig. 4].
RaggedShard composes with existing DTensor Shard(dim) placements by introducing StridedRaggedShard for Shard(0) (reorder + stride metadata) and adapting granularity to LCM(stride(dim), user-granularity) for Shard(dim>0), preserving TP and EP interop [§4, Fig. 5].
Efficient grouped communication requires three constraints jointly — non-sharded block, contiguous tensor memory, balanced load — formulated as an NP-hard optimization (reducible from Partition); ILP solutions take tens of minutes at scale and time out [§5, Eq. block].
The proposed polynomial-time DP heuristic exploits transformer regularity (linear weights dominate, block sizes are consistent across layers) and a case analysis over how tensors align with shard boundaries, achieving an O(N log V)-style 2-approximation with negligible plan overhead [§5, Algorithm 1].
The Distributed Buffer (DBuffer) primitive backs grouped RaggedShard DTensors with slices of a single global buffer over an n-dim device topology, fusing identical per-tensor pre-communication kernels (add/scale/zero/copy), enabling zero-copy AllGather/ReduceScatter via persistent address mapping, and using in-place comm+compute [§5, Fig. 7].
On 1024 GPUs against DeepSpeed, FSDP1, FSDP2, Megatron-FSDP: veScale-FSDP is 1166% faster than all baselines on internal MoE and 5% faster than DeepSpeed/FSDP1/FSDP2 on LLaMA-3-70B, while reducing peak reserved memory by 16–30% [§6.1, Fig. 8].
FSDP2’s per-parameter even-sharding pays a 14% iteration-time penalty from interleaved AllGather copy-out and ReduceScatter copy-in; Megatron-FSDP’s Stride(0) constraint inflates the comm buffer by 33% on MoE due to padding inflation [§6.1, Table 1].
GPT-OSS-120B exposes a hard memory cliff for FSDP2: it trains at 128 devices but OOMs at 256, because per-parameter padding to enforce even splits over 128 experts × 256 devices effectively doubles the AllGather buffer [§6.1].
RaggedShard’s block-wise granularity natively supports 8-bit Adam (block-quantized optimizer state) and Muon (atomic 2D matrix locality for orthogonalization) without intrusive model code [§6.3].

Method

veScale-FSDP keeps PyTorch’s fully_shard API but replaces the FSDP2 backend in three layers. (1) RaggedShard: a new DTensor placement carrying a per-tensor atomic block shape plus a per-device block count, so a tensor can be split into ragged 1D or 2D blocks distributed asymmetrically across devices. It composes with TP/EP via StridedRaggedShard (for Shard(0) coexistence) and LCM-based granularity adaptation (for higher-dim shards). (2) Planner: given a fully_shard wrapping, group all RaggedShard DTensors into a contiguous communication buffer that minimizes per-device buffer size subject to non-sharded-block / contiguous-tensor / balanced-load constraints. The exact problem is NP-hard, so the planner explores a small set of heuristic permutations (default, sorted-by-block-size, sorted-by-shape) and runs a DP layout pass with O(NV log V)-style complexity; it monotonicity-batches contiguous index ranges per tensor and binary-searches over shard-size multiples when “Case-3” tensors (fully contain at least one shard) are present. (3) DBuffer: a 1- to n-D shared buffer primitive that owns a persistent address mapping for each RaggedShard tensor. Pre/post-comm work (zero, scale, add, copy) is fused at group level, AllGather/ReduceScatter operates in-place against aligned offsets, and the underlying caching allocator sees one batched allocation instead of many per-parameter ones — eliminating the FSDP2-style record_stream non-determinism and the Megatron-FSDP padding inflation in one stroke.

Results

Throughput (1024 GPUs, weak scaling): 1166% over best baseline on internal MoE; 5% over DeepSpeed/FSDP1/FSDP2 and a small lead over Megatron-FSDP on LLaMA-3-70B; “linear scalability” beyond.
Peak reserved memory: 16–30% reduction across LLaMA-3-70B / GPT-OSS-120B / internal MoE; specifically 20% vs DeepSpeed/FSDP1 (record_stream fragmentation), 12% vs FSDP2 (per-parameter eager allocation), 33% vs Megatron on MoE (padding inflation), and 24% vs Megatron on LLaMA-3 (persistent low-precision buffers).
Memory cliff narrowness: FSDP2 OOMs on GPT-OSS-120B at 256 devices but trains at 128 — a clean demonstration that per-parameter even-sharding’s required padding scales adversely with #experts vs #devices.
Headline ablation: 566% throughput / 1630% memory advertised as the upper bound across all baseline configurations (presumably the GPT-OSS OOM regime), in addition to the 1166% MoE number.
Optimizer support: Muon and 8-bit Adam land “natively” (no intrusive code) thanks to RaggedShard’s customizable granularity; reported performance and dev-velocity gains in §6.3 (not on a single figure quoted here).
Planning overhead: heuristic DP completes in polynomial time, vs ILP solvers that “take tens of minutes” or hit system timeouts.
LoC: 7.6K lines of Python; plug-in replacement behind fully_shard; production-tested in ByteDance and portable.

Why it’s interesting

veScale-FSDP makes the operational story behind modern optimizer choices (Muon, Shampoo) and modern quantization choices (block-wise FP8) finally clean — the implicit cost of those techniques in stock FSDP was either intrusive model code, padding-inflated buffers, or both. This is the systems-side dual of the algorithmic stability fork the wiki has been tracking (Training stability at scale) — Controlled LLM Training on Spectral Sphere argues that constraining both weights and updates to a spectral sphere is the right stability locus and ships as a Megatron-LM fork, but its Atomic Module Sharding requirement is exactly what generic FSDP cannot express today; RaggedShard is the missing primitive that lets that style of optimizer drop into a non-bespoke training stack. Complements GLM-5: from Vibe Coding to Agentic Engineering, whose own training-infra notes include “flexible MTP placement under interleaved pipeline parallelism” and “Pipeline ZeRO2 gradient sharding with double-buffered accumulation” — same problem space, in-house solution; veScale-FSDP is the upstreamable form. The 33% Megatron-FSDP MoE padding-inflation number is also the cleanest published quantification of why current open-source FSDP isn’t MoE-native.