RaggedShard: Flexible Tensor Sharding

• RaggedShard is a structure-aware sharding scheme that partitions tensors into atomic, contiguous blocks to preserve alignment and enable efficient computation.
• It employs a global shard planner with heuristic and ILP-inspired methods to balance device buffer sizes while minimizing padding and communication overhead.
• Integration with PyTorch SPMD supports block-wise quantization and matrix-aware optimizers, resulting in improved throughput, memory efficiency, and scalability.

RaggedShard is a flexible, structure-aware sharding scheme for distributed tensor storage and communication, introduced in veScale-FSDP to address the limitations of traditional element- and row-wise sharding formats in large-scale deep learning. Its core innovation is treating each tensor as a sequence of atomic, contiguous memory blocks of arbitrary shape—enabling efficient, alignment-preserving assignment of these blocks to devices. RaggedShard supports complex training paradigms, such as block-wise quantization and matrix-aware optimizers, achieving superior throughput, memory efficiency, and scaling relative to established Fully Sharded Data Parallel (FSDP) approaches (Wang et al., 25 Feb 2026).

1. Motivation and Design Principles

Traditional FSDP and ZeRO-based systems are constrained by element- or row-wise sharding, leading to misalignments and inefficiencies when deploying block-structured computations. For example, 32×32 block-wise quantization and full-matrix optimizers like Shampoo or Muon require tensor subdivisions that do not conform to fixed-grain sharding, forcing manual data realignment or runtime overhead due to padding and bespoke collectives.

RaggedShard is designed to:

• Divide tensors into atomic, fixed-size blocks (shape $\mathbf B$), each indivisible for the purposes of sharding.
• Allow arbitrary mapping of these atomic blocks to devices via a function $f: \{\text{blocks of }T\} \to \{0,\dots,m-1\}$ for $m$ GPUs.
• Preserve contiguous layout of blocks within tensor storage, with any required padding restricted to inter-tensor regions in communication buffers.
• Compose cleanly with existing DTensor placements (Replicate, Shard(dim)), enabling orthogonal parallelism strategies.

These principles enable exact block alignment for quantization, full-matrix operations, and irregular sparsity layouts inherent in Mixture-of-Experts (MoE) models.

2. Data Layout and Mapping

Let a tensor $t$ of shape $\mathbf D_t \in \mathbb N^d$ be partitioned into blocks of shape $\mathbf B_t \in \mathbb N^d$. The number of blocks per dimension is

$U_t(i) = \lceil D_t(i)/B_t(i) \rceil,$

and the total number of blocks is $U_t = \prod_{i=0}^{d-1} U_t(i)$. Each block is indexed as $\mathbf u = (u_0,\dots,u_{d-1})$, $0 \le u_i < U_t(i)$.

The offset of block $f: \{\text{blocks of }T\} \to \{0,\dots,m-1\}$0 in flat storage is

$f: \{\text{blocks of }T\} \to \{0,\dots,m-1\}$1

and its size in elements is $f: \{\text{blocks of }T\} \to \{0,\dots,m-1\}$2.

Device assignment is defined by $f: \{\text{blocks of }T\} \to \{0,\dots,m-1\}$3, with each GPU $f: \{\text{blocks of }T\} \to \{0,\dots,m-1\}$4 storing its contiguous blocks:

$f: \{\text{blocks of }T\} \to \{0,\dots,m-1\}$5

This data layout ensures all computation and quantization within a block are local, with no further communication needed for atomic block operations.

3. Structure-Aware Shard Planning

The global sharding problem—finding an assignment of blocks to GPUs that (a) respects block atomicity, (b) equalizes per-device buffer size $f: \{\text{blocks of }T\} \to \{0,\dots,m-1\}$6, and (c) avoids partial blocks or padding at shard boundaries—is NP-hard by reduction from Partition.

An ILP formulation (not used at runtime) minimizes $f: \{\text{blocks of }T\} \to \{0,\dots,m-1\}$7 and determines offsets $f: \{\text{blocks of }T\} \to \{0,\dots,m-1\}$8:

• $f: \{\text{blocks of }T\} \to \{0,\dots,m-1\}$9, $m$0,
• Non-overlapping intervals $m$1,
• For every device boundary $m$2, either $m$3 is outside $m$4 or aligned on a block boundary: $m$5.

veScale-FSDP implements a polynomial-time heuristic:

• Sort tensors (e.g., by block size).
• Binary search for minimal $m$6 over multiples of least common multiples (l.c.m.) of block sizes.
• Use dynamic programming in CheckValidShard(S) to verify feasibility in $m$7.
• Overall complexity is $m$8.

Simplified pseudocode:

$U_t(i) = \lceil D_t(i)/B_t(i) \rceil,$4 Where $m$9 is the element count, $t$0 block size, and $t$1 the set of tensors.

4. Integration with FSDP and Runtime Mechanics

RaggedShard operates as a new placement for DTensor, natively supported in PyTorch SPMD environments:

• Initialization: Modules are wrapped as FullyShardedDataParallel(..., sharding=RaggedShard(...)). The system plans global buffer sizes and offsets, then allocates a $t$2 DBuffer and sets up zero-copy slicing for local shards.
• Forward pass: Prior to execution, FSDP invokes DBuffer.all_gather(parameters), converting placements from RaggedShard(Shard) to RaggedShard(Replicate). This enables a single, large NCCL AllGather operation.
• Backward pass: Following gradient computation, DBuffer.reduce_scatter(gradients) performs a unified NCCL ReduceScatter to revert to RaggedShard(Shard) format, followed by local optimizer steps.

This unified buffer strategy eliminates per-tensor padding, redundant copying, and inefficient collective sizes.

5. Enabling Block-Wise Quantization and Matrix-Aware Optimizers

RaggedShard enables precise block-level control required by structure-aware training components:

• Block-wise 8-bit Adam: Given a quantization tile size (e.g., $t$3), each block is assigned locally, with quantization and scaling factors computed entirely on-device, removing the need for cross-device scale factor gathering or manual padding.
• Matrix-aware optimizers (e.g., Muon): RaggedShard allows all blocks of a 2D tensor to reside on a root GPU $t$4. Operations such as Newton–Schulz iterations for full-matrix preconditioning proceed entirely locally. Redistribute operations handle block collection and broadcast, abstracted away by DTensor and RaggedShard logic, eliminating manual intervention and custom communication code.

6. Performance and Scaling Characteristics

Performance and efficiency metrics for RaggedShard in veScale-FSDP include:

• Communication Volume:

$t$5

elements communicated for all-gather and reduce-scatter, where $t$6 is total parameter elements, $t$7 number of devices, and $t$8 aggregate padding ($t$9 of $\mathbf D_t \in \mathbb N^d$0 in practice).

• Memory Overhead per GPU:

$\mathbf D_t \in \mathbb N^d$1

much improved compared to up to $\mathbf D_t \in \mathbb N^d$2 padding in fixed-sharding.

• Throughput Model:

$\mathbf D_t \in \mathbb N^d$3

Communication reduction proportional to $\mathbf D_t \in \mathbb N^d$4 increases $\mathbf D_t \in \mathbb N^d$5 by $\mathbf D_t \in \mathbb N^d$6.

• Empirical observations:
  • Throughput for $\mathbf D_t \in \mathbb N^d$7B–$\mathbf D_t \in \mathbb N^d$8B parameter LLMs at $\mathbf D_t \in \mathbb N^d$9K GPUs is $\mathbf B_t \in \mathbb N^d$0 higher than DeepSpeed/ZeRO-3 and PyTorch FSDP, $\mathbf B_t \in \mathbb N^d$1 higher on MoE models.
  • Memory consumption is $\mathbf B_t \in \mathbb N^d$2 lower per-GPU than all baselines.
  • Padding is $\mathbf B_t \in \mathbb N^d$3 for typical $\mathbf B_t \in \mathbb N^d$4 or $\mathbf B_t \in \mathbb N^d$5 blocks.
  • Planner overhead is $\mathbf B_t \in \mathbb N^d$6 seconds for hundreds of tensors on thousands of GPUs.
  • Scaling demonstrated up to $\mathbf B_t \in \mathbb N^d$7 GPUs with near-linear weak and strong scaling.
  • Model scaling: MFU for $\mathbf B_t \in \mathbb N^d$8T parameters on $\mathbf B_t \in \mathbb N^d$9K GPUs remains around $U_t(i) = \lceil D_t(i)/B_t(i) \rceil,$0–$U_t(i) = \lceil D_t(i)/B_t(i) \rceil,$1.

7. Significance and Influence

RaggedShard provides a uniform, block-aware sharding abstraction, underpinned by a provably near-optimal planner and high-performance DBuffer primitive. It furnishes structure-aware training methods—block quantization, full-matrix preconditioners, and sparse MoE layouts—with consistent, boundary-aligned semantics and high runtime efficiency. Block-wise quantized Adam and matrix-aware Muon optimizers converge identically to data-parallel (DDP) baselines, but execute at $U_t(i) = \lceil D_t(i)/B_t(i) \rceil,$2–$U_t(i) = \lceil D_t(i)/B_t(i) \rceil,$3 of full-precision throughput without code changes.

By decoupling atomic block sharding from rigid element/row constraints, RaggedShard eliminates manual boundary checks and padding overhead, substantially outperforming fixed-grain formats on both communication and memory metrics, particularly as model and system scales increase (Wang et al., 25 Feb 2026).