RaggedShard for Parameterized GSM8K Templates

• The paper introduces RaggedShard, a flexible sharding format that enables variable-sized tensor partitioning for block-wise quantized training and non-element-wise optimizers.
• It details a structure-aware planning algorithm that minimizes communication costs and load imbalance by optimally assigning irregular tensor blocks to thousands of GPUs.
• Integrated in veScale-FSDP, RaggedShard delivers 5–66% throughput improvements and 16–30% lower memory usage compared to fixed sharding approaches.

RaggedShard is a flexible sharding format introduced in the veScale-FSDP system to address the challenges of distributed training for large-scale deep learning models, specifically in the context of block-wise quantized training and non-element-wise optimizers. Unlike fixed element- or row-wise sharding schemes, which impede modern structured computation patterns, RaggedShard enables efficient, structure-aware partitioning of tensors across thousands of GPUs, delivering substantial improvements in throughput and memory usage (Wang et al., 25 Feb 2026).

1. Definition and Motivation

RaggedShard is a variable-sized, structure-aware sharding format designed to represent and distribute tensors whose optimal partitioning follows complex and heterogeneous patterns. Traditional FSDP (Fully Sharded Data Parallel) systems utilize rigid element-wise or row-wise sharding: each tensor is split along a single fixed axis, with each GPU receiving an equally sized shard. While effective for element-wise optimizers and uniform data layouts, this approach fails to accommodate block-wise quantization (which operates on variable-sized blocks rather than rows/elements) and advanced optimizers such as Shampoo and Muon (which consume and produce parameter matrices with nontrivial block or submatrix structure).

Fixed sharding introduces two key inefficiencies in this context:

• Misalignment with computation granularity: Block-wise computations requiring noncontiguous or irregular partitions necessitate costly realignment or redundant data movement.
• Load and memory imbalance: In heterogeneous block designs, fixed sharding produces imbalanced shards, straining memory and network usage.

RaggedShard addresses these limitations by allowing each GPU to receive a "ragged" (i.e., variable-sized and possibly noncontiguous) portion of each tensor, with boundaries adapted to computation and communication patterns.

2. Design and Data Structures

RaggedShard internally represents the mapping from logical tensor blocks to physical GPU memory using metadata structures encoding per-GPU block assignments, offsets, and indices. Each shard is characterized by:

• Block assignment list: For a tensor $T$ partitioned into $K$ blocks $\{B_1, B_2, ..., B_K\}$, RaggedShard maintains a mapping $S: \{1,...,G\} \to 2^{\{1,...,K\}}$ assigning blocks to GPUs.
• Offset tables: For each GPU $g$ and assigned block $B_k$, an offset-indexes tuple $(\text{offset}_k, \text{size}_k)$ specifies the starting index and size in the flattened tensor storage.
• Ragged tensor descriptor: A descriptor defines the block structure, assignment, and strides, enabling contiguous or strided allocation as needed.

The memory footprint for a tensor $T$ distributed across $G$ GPUs is given by:

$$\text{memory\_usage} = \sum_{i=1}^{G} \sum_{k \in S(i)} \text{size}(B_k) + \text{metadata\_overhead},$$

where $K$0 is the block size (in elements or bytes) and $K$1 accounts for per-shard index and offset tables.

Table: Comparison of Sharding Metadata (illustrative layout)

Sharding Format	Block Assignment	Offset Table Complexity
Element-wise	Contiguous, uniform	$K$2 (single stride)
Row-wise	Contiguous rows	$K$3 (row stride)
RaggedShard	Arbitrary blocks	$K$4 (block index/offset)

3. Structure-Aware Planning Algorithm

The structure-aware planning algorithm in RaggedShard assigns tensor blocks to GPU shards to minimize communication costs and maintain load balance, subject to the block structure imposed by quantization or optimizer schedules.

Let $K$5 denote the communication cost for assigning block $K$6 to GPU $K$7, and $K$8 the total size assigned to GPU $K$9. The planning objective is:

$\{B_1, B_2, ..., B_K\}$0

where $\{B_1, B_2, ..., B_K\}$1 weight communication, memory footprint, and load balance.

Step-by-step pseudocode (high-level):

$\{B_1, B_2, ..., B_K\}$8

Constraints include: $\{B_1, B_2, ..., B_K\}$2, and block dependencies for optimizer-specific updates.

4. Integration with veScale-FSDP

In veScale-FSDP, RaggedShard is the native sharding backend, interfacing directly with both runtime and optimizer APIs. Key integration points are:

• Block-wise quantization: RaggedShard supports quantization schedules that partition parameter tensors into nonuniform blocks (e.g., varying block sizes for different matrix bands or model layers) and maintains exact block boundaries on each GPU for fast, local quantization and dequantization.
• Non-element-wise optimizers: For optimizers like Shampoo and Muon, which operate on block-diagonal or structured updates, RaggedShard ensures that each GPU holds exactly the blocks required by its optimizer instances, avoiding redundant storage or communication.
• Data placement: The format enables arbitrary mapping of tensor shards, allowing dynamic assignment and realignment as model topology or optimizer schedules change.
• Communication: Communication operations are orchestrated over ragged shards. For collective operations:
  • All-gather: RaggedShard orchestrates variable-length gathers where each GPU contributes its owned blocks, followed by assembly on target GPUs.
  • Reduce-scatter: Reductions operate over aligned block lists, with only present blocks reduced per GPU.

This design enables veScale-FSDP to efficiently perform end-to-end training (including forward, backward, optimizer step, and quantization) with minimal data movement and maximal parallelism.

5. Performance Analysis

Empirical benchmarks from (Wang et al., 25 Feb 2026) demonstrate that RaggedShard, as integrated in veScale-FSDP, achieves:

• Throughput improvement: 5–66% higher throughput compared to fixed-sharding FSDP systems across a spectrum of model architectures and optimizer schedules.
• Memory footprint: 16–30% lower memory usage, attributed to the absence of alignment padding and block-level memory fragmentation.

Experiments were conducted on large-scale settings, scaling to tens of thousands of GPUs. Performance scaling curves show near-linear throughput improvement with increasing GPU count, subject to communication topology constraints. The flexible data partitioning reduces communication volume by up to 41% for block-wise quantized models. For non-element-wise optimizers, memory waste from misaligned shards is nearly eliminated.

The system demonstrates efficient scaling with number of GPUs $\{B_1, B_2, ..., B_K\}$3, with runtime per step approximating:

$\{B_1, B_2, ..., B_K\}$4

where $\{B_1, B_2, ..., B_K\}$5 is total computation, $\{B_1, B_2, ..., B_K\}$6 the per-GPU communication, and $\{B_1, B_2, ..., B_K\}$7 system-dependent scaling factors.

6. Comparative Discussion

Compared to fixed element-wise or row-wise sharding, RaggedShard exhibits the following properties:

Aspect	Fixed Element-wise/Row-wise	RaggedShard
Block-structured compat.	Limited (requires copying)	Native support
Load balance	Only for uniform block sizes	Maintains balance for arbitrary blocks
Shard size variance	None (uniform)	Arbitrary
Metadata overhead	Minimal	Higher (per block)
Communication pattern	Simple (fixed schedule)	Complex (per-block, ragged)

Limitations: RaggedShard incurs higher metadata management overhead and increased scheduling complexity. The block-assignment problem is computationally richer, leading to possibly higher planning latency for very large-scale deployments.

Future directions include optimizing planning algorithms for extreme scale, further reducing communication overhead with topology-aware assignment, and tighter integration with emerging non-uniform memory access architectures.

(Wang et al., 25 Feb 2026)