---
title: 'RaggedShard: Flexible Tensor Sharding'
url: https://www.emergentmind.com/topics/raggedshard
type: topic
---

# RaggedShard: Flexible Tensor Sharding

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 [2602.22437].

## 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 $\mathbf u$ in flat storage is
$$
\mathrm{off}_t(\mathbf u) = \sum_{i=0}^{d-1} \left( u_i \cdot B_t(i) \right) \cdot \prod_{j<i} D_t(j),
$$
and its size in elements is $\prod_i \min\{ B_t(i), D_t(i) - u_i B_t(i) \}$.

Device assignment is defined by $f_t(\mathbf u)$, with each GPU $p$ storing its contiguous blocks:
$$
L_{t,p} = \bigcup_{\mathbf u: f_t(\mathbf u) = p} [ \mathrm{off}_t(\mathbf u),\, \mathrm{off}_t(\mathbf u) + |\mathbf u| ).
$$
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 $S$, 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 $S$ and determines offsets $\{\ell_t, r_t\}$:
- $r_t - \ell_t = e_t$, $r_t \le mS$,
- Non-overlapping intervals $[\,\ell_t, r_t)$,
- For every device boundary $kS$, either $kS$ is outside $[\,\ell_t, r_t)$ or aligned on a block boundary: $(kS - \ell_t) \equiv 0 \pmod{g_t}$.

veScale-FSDP implements a polynomial-time heuristic:
- Sort tensors (e.g., by block size).
- Binary search for minimal $S$ over multiples of least common multiples (l.c.m.) of block sizes.
- Use dynamic programming in `CheckValidShard(S)` to verify feasibility in $O(m\log E)$.
- Overall complexity is $O(|\mathcal T|^2 m \log E \log(|\mathcal T|m))$.

**Simplified pseudocode:**
```python
def PlanShards({e_t, g_t}, m, g_coll):
    G = sort({g_t})
    for gamma in G:
        base = lcm(g_coll, gamma)
        S = binary_search_smallest(base, CheckValidShard)
        keep_minimal_S
    return S, {ell_t, r_t}
```
Where $e_t$ is the element count, $g_t$ block size, and $\mathcal T$ 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 $m \times S$ 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., $\mathbf B_t = (32,32)$), 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 $r$. 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**:
  $$
  2\left( E - \tfrac{E}{m} \right) + 2P
  $$
  elements communicated for all-gather and reduce-scatter, where $E$ is total parameter elements, $m$ number of devices, and $P$ aggregate padding ($<3\%$ of $E$ in practice).

- **Memory Overhead per GPU**:
  $$
  \tfrac{E}{m} + P + \text{peak staging}
  $$
  much improved compared to up to $30\%$ padding in fixed-sharding.

- **Throughput Model**:
  $$
  T \propto \frac{ \text{compute} }{ \text{compute} + \alpha\,\mathrm{Comm} + \beta\,\mathrm{latency} }
  $$
  Communication reduction proportional to $O(P)$ increases $T$ by $5–66\%$.

- **Empirical observations**:
  - Throughput for $70$B–$120$B parameter LLMs at $1$K GPUs is $5\%$ higher than DeepSpeed/ZeRO-3 and PyTorch FSDP, $11–66\%$ higher on MoE models.
  - Memory consumption is $16–30\%$ lower per-GPU than all baselines.
  - Padding is $<3\%$ for typical $32 \times 32$ or $16 \times 16$ blocks.
  - Planner overhead is $<0.3$ seconds for hundreds of tensors on thousands of GPUs.
  - Scaling demonstrated up to $10,000+$ GPUs with near-linear weak and strong scaling.
  - Model scaling: MFU for $2.4$T parameters on $1$K GPUs remains around $50$–$55\%$.

## 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 $90$–$100\%$ 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 [2602.22437].

Source: https://www.emergentmind.com/topics/raggedshard