Sharder

What sets the ceiling

Assume the transformer background. Three quantities decide everything in this project: how we score a step (MFU), whether the step has any single-device headroom left (it is compute-bound), and why we are forced to shard and communicate at all (the memory budget and the cost of collectives).

MFU, the number we optimize

Model FLOPs utilization is the fraction of the accelerator's peak arithmetic rate the training step actually achieves. The step's useful work is counted analytically: a transformer forward pass is about \(2ND\) FLOPs for \(N\) parameters and \(D\) tokens, and the backward pass is roughly twice that, so a step is \(\approx 6ND\). Divide the achieved rate by the device's peak:

$$ \text{MFU} \;=\; \frac{6\,N\,D}{t_{\text{step}}\;\pi_{\text{peak}}} $$

# MFU from the analytical FLOP count, computed straight from measured step time.
flops_per_step = 6 * n_params * n_tokens_per_step   # fwd 2ND + bwd 4ND
achieved       = flops_per_step / step_time_s        # FLOP/s actually delivered
mfu            = achieved / peak_flops               # e.g. ~0.32 on one H100

Compute-bound, so there is no single-device win

A kernel is compute-bound when its arithmetic intensity, FLOPs performed per byte moved, is high enough that it sits under the compute roof rather than the bandwidth roof:

$$ \text{Perf} \;=\; \min\!\bigl(\pi_{\text{peak}},\; I\times\beta\bigr) $$

Large matrix multiplications have high intensity; elementwise work like a softmax does not. The measurement in Part VI shows a full pi0 step running at about 32% MFU with matrix-multiply work at 57.7% of device time and near-zero idle. That is a GEMM-saturated, compute-bound step: there is no speculative kernel-level optimization to chase on one device. The only lever left for throughput is using more devices, which is what the rest of the project is about.

The memory budget forces sharding

Sharding is not optional for this model. Parameters, gradients, and the optimizer state all have to live in device memory, and AdamW alone keeps two extra full-precision copies of every parameter. A rough budget for gemma_2b makes the point:

params (bf16)          2 bytes/param
gradients (bf16)       2 bytes/param
Adam m, v (fp32)       8 bytes/param
--------------------------------------------
                     ~12 bytes per parameter, before activations

That multiplies out well past a single accelerator once activations are included, which is exactly why the design shards the parameters and optimizer state across devices with FSDP rather than replicating them. The question is never "does it fit on one device"; it is "how do we split it so each device holds a slice."

Collectives, and why scaling is a communication problem

Once state is sharded, devices have to exchange it. The two collectives that matter here are all-gather, which reassembles a full parameter matrix from its shards before a layer runs, and reduce-scatter, which sums gradients and splits the result back into shards. Inside a node these run over NVLink through NCCL. Their cost relative to compute is the whole story of the scaling results: if the all-gather for the next layer hides behind the matmul of the current one, scaling is efficient; if it does not, the communication lands on the critical path.

Compiler and kernels

The stack is standard, so only the two places it constrains the project are worth stating: how XLA's fusion distorts profiling, and why cuDNN's flash kernel is unavailable on this model. The rest, Flax NNX for the model, Optax for AdamW, Orbax for checkpoints, is used as-is.

The property we exploit

The training step is a pure function of parameters and batch, so its array layouts are declarative: jax.jit traces it, and if the inputs are already committed to shardings, it infers the entire sharded computation from them. Parallelism is a placement decision, not a code change, which is why the FSDP wiring in Part III is a handful of lines rather than a rewrite.

Fusion, and what it does to a profiler

XLA lowers the step to HLO and fuses adjacent operations into single kernels so intermediates stay in fast memory. The attention softmax is a cluster of elementwise ops that gets fused into the surrounding matrix multiply. That fusion is the reason the naive profiling approach is wrong in two ways at once, both corrected in Part V:

Why cuDNN's flash kernel is out

cuDNN's fused flash-attention kernel is the fastest path when it applies, but it constrains the head dimension. gemma_2b's is 256; the kernel requires at most 128, so it rejects the shape. Confirmed by dropping to head dim 128, where cuDNN runs. Since 256 is fixed by the model, the usable fused path is XLA's implementation="xla", not cuDNN, which is exactly what the FlashAttention wiring in Part V uses.

The mechanisms

With the hardware and libraries in place, four mechanisms do the actual work of scaling. Each has a simple picture behind it.

The device mesh

JAX arranges the available accelerators into a logical grid called a mesh, and gives its axes names. Sharder uses a two-dimensional mesh with a batch axis and an fsdp axis. A tensor is then sharded by declaring which of its dimensions map to which mesh axis: the data batch is split along batch, and the model parameters are split along fsdp.

In code, the mesh is built once, then a tensor's layout is declared with a PartitionSpec that names which mesh axis each of its dimensions maps to. Parameters are sharded along fsdp; the data batch is split along the data axis.

from jax.sharding import NamedSharding, PartitionSpec as P

mesh = sharding.make_mesh(num_fsdp_devices=n_gpus)      # 2D: (batch, fsdp)

state_sharding = sharding.fsdp_sharding(params, mesh)   # one PartitionSpec per param leaf
data_sharding  = NamedSharding(mesh, P(("batch", "fsdp")))  # shard batch dim 0, any rank

FSDP: shard, gather, compute, free, scatter

Fully-sharded data parallelism keeps each parameter matrix split across the fsdp devices, so no device ever holds the whole model. To run a layer, the shards are all-gathered into the full matrix just in time; the matrix multiply runs; the full matrix is freed; and in the backward pass the gradients are reduce-scattered back into shards. Memory stays sharded, and only the one matrix currently in use is ever whole.

The elegant part is that none of this is written by hand. The parameters are placed on their shards, the inputs on the data sharding, and jax.jit infers the whole sharded computation from the committed input layouts. XLA is what emits the all-gather before each matmul and the reduce-scatter after the backward pass.

# place state on its shards; jit reads the layouts off the placed inputs
params   = jax.device_put(params, state_sharding)
obs, act = jax.device_put((obs, act), data_sharding)

step = jax.jit(train_step)     # no manual collectives: XLA emits gather / reduce-scatter
loss, grads = step(params, obs, act)

Data sharding, and why it must be exact

When several processes train together, each reads a different slice of the data. The naive way, hand each process every N-th example, changes which examples land in each step compared to a single-process run. For a portfolio-grade contribution that is not good enough: the multi-process run should produce bit-identical loss and gradients to the single-process baseline, so nothing about correctness depends on the device count.

Sharder's sampler shards within each global batch. Every process takes a contiguous slice of the same deterministic global order, and concatenating those slices in process order reconstructs exactly the batch a single process would have formed.

The sampler is the whole trick. Each process takes a contiguous slice within every global batch of one shared, deterministic order, so the slices reassemble the exact single-process batch.

# process pi of pc: a contiguous slice inside each global batch of the shared order.
# concatenating slice(P0) .. slice(Pn) rebuilds the single-process global batch exactly.
def process_index_stream(order, global_batch, pi, pc):
    local = global_batch // pc
    return np.concatenate([
        order[s*global_batch + pi*local : s*global_batch + (pi + 1)*local]
        for s in range(len(order) // global_batch)
    ])

FlashAttention

Attention compares every token to every other token, forming a score matrix of size sequence × sequence. The naive implementation writes that whole matrix to memory, so its memory grows with the square of the sequence length. FlashAttention computes the same result by streaming the matrix in tiles and never storing it whole, which makes it faster and, more importantly, keeps its memory flat as the sequence grows. Part VI shows how decisive that becomes at longer context.

Wiring it into Gemma is a drop-in replacement for the einsum-softmax-einsum, with three details that have to be right: the query is already scaled by head_dim**-0.5 so the kernel scale is 1.0, the block-causal mask is passed through, and the grouped-query layout (eight query heads, one key/value head) is handled by the kernel.

# fused path in gemma attention (behind use_flash_attention, default on)
encoded = jax.nn.dot_product_attention(
    q, k, v,                       # q pre-scaled; k, v are the single GQA kv head
    mask=attn_mask,                # block-causal prefix + causal suffix
    scale=1.0,                     # scaling already applied to q
    implementation="xla",          # cuDNN rejects head_dim 256; xla is the usable fused path
)   # bit-identical loss to the naive path, verified in pi0_flash_test.py

openpi and the pi0 model

openpi trains vision-language-action models, or VLAs: models that take camera images and a language instruction and output robot actions. The model this project profiles is pi0, a flow-based VLA.

The pi0 model

pi0 has three parts. A SigLIP vision encoder turns each camera image into a sequence of visual tokens. A Gemma language backbone, gemma_2b in this configuration, processes the visual tokens together with the text instruction. A flow-matching action expert then produces the continuous action trajectory. The Gemma backbone is where most of the compute lives, and its configuration, eighteen layers of width 2048 with an MLP dimension of 16384, eight query heads sharing one key/value head, and a head dimension of 256, is exactly what determines the numbers later on.

The training step

A training step is the pure function from Part II applied to this model: run the forward pass to a flow-matching loss, differentiate to get gradients, and apply the AdamW update. That single step is the unit everything is measured on. Its data comes from a pipeline that loads and prepares the next batch on the host CPU while the accelerator is still computing the current one, so the load is hidden behind compute rather than stalling it.

The project, gap by gap

Sharder was built as a set of numbered goals, G1 through G6, each closing one gap between openpi's single-host script and a multi-host, fault-tolerant one. For each, the interesting content is the same three things: what was missing, how it was actually done in the code, and what it bought.

G1 · Bringing up multi-host

Motivation. To train across machines, JAX has to fuse the separate processes into one logical program, so that a collective like all-gather spans every device. That handshake is jax.distributed.initialize.

Mechanism. A small module reads the coordinator address, the process count, and this process's index from the environment and performs the handshake, and does nothing at all when there is a single process. The single-host path is therefore untouched.

def maybe_initialize():
    addr = os.environ.get("JAX_COORDINATOR_ADDRESS")
    if not addr:
        return                              # single process: no-op, upstream path untouched
    jax.distributed.initialize(
        coordinator_address=addr,
        num_processes=int(os.environ["JAX_NUM_PROCESSES"]),
        process_id=int(os.environ["JAX_PROCESS_ID"]),
    )                                       # now jax.device_count() spans every host

G2 / G3 · Sharding with bit-parity

Motivation. openpi's loader raised NotImplementedError for more than one process. Removing that guard is easy; doing it so the result matches single-process training exactly is the real work.

Mechanism. The within-batch sampler from Part III gives each process a contiguous slice of one deterministic global order, so the union of slices is the single-process batch. A test asserts that the multi-process loss and gradients equal the single-process reference.

G6 · The instrument, and two corrections

Motivation. The project's rule is to never optimize a bottleneck it has not measured. That requires an instrument that attributes a step's device time to categories, attention, matmul, optimizer, and so on, and reports utilization honestly.

Mechanism. The instrument maps each compiled operation back to its source scope and assigns its time to a category. Building it surfaced two measurement errors that, uncorrected, told the wrong story.

Result. A gemma_2b step runs at about 32% MFU on one H100, with matrix-multiply work at 57.7% of device time and near-zero idle. The step is compute-bound and GEMM-saturated: there is no speculative single-GPU optimization to chase, so the lever is parallelism.

G4 · Exact resume, without storing the iterator

Motivation. A crashed run must resume on exactly the examples it had reached, with none repeated or skipped. The usual approach checkpoints the data iterator's internal state, which is fragile.

Mechanism. Because the example order is a pure function of the seed and the epoch, and the training state already counts steps, the resume position can be derived rather than stored: the step counter maps to an (epoch, offset), and the sampler seeks there. No iterator blob is ever checkpointed. A test confirms the resumed stream continues exactly, including across an epoch boundary.

# the entire "iterator state" is one integer: the step counter already in TrainState.
def resume_position(consumed_batches, dataset_len, global_batch):
    n = dataset_len // global_batch          # batches per epoch
    return consumed_batches // n, consumed_batches % n   # -> (epoch, offset)

# on restore: seek the deterministic sampler; no separate iterator checkpoint exists
epoch, offset = resume_position(int(state.step), dataset_len, global_batch)
sampler.seek(epoch, offset)

G5 · Elastic restart

Motivation. On a large cluster, processes and nodes die. Training should survive that by restarting from the last checkpoint rather than starting over.

Mechanism. A supervisor runs the training command and, if it exits with an error, relaunches it with a resume flag, up to a retry limit with backoff. It composes the earlier pieces: G1 brings the cluster back up, Orbax restores the checkpoint, and G4 seeks the data to the right step. A fault-injection test kills a run partway and asserts it resumes with no lost or duplicated work.

for attempt in range(max_retries + 1):
    cmd = command + (["--resume"] if attempt else [])   # fresh start, then resume
    rc = subprocess.run(cmd).returncode
    if rc == 0:
        break
    time.sleep(backoff * 2**attempt)                    # a dead process becomes a resume
# G1 re-inits the cluster, Orbax restores the checkpoint, G4 seeks the data to state.step

FlashAttention, wired and verified

Motivation. Attention's cost was hidden by fusion; once measured it was worth fusing for memory and for longer context.

Mechanism. The fused kernel is wired into the Gemma attention behind a flag, correctly handling the block-causal mask, the pre-scaling of the query, and the grouped-query head layout. It is verified to produce the same loss as the naive path, to bit precision, for both training and inference.

Result. A free 9.3% end-to-end speedup, enabled by default because it changes no outputs. The measured detail is in the results.

Results

Everything measured on a four-GPU H100 node, on gemma_2b in bf16. The figures below are generated from the recorded numbers by scripts/plot_results.py.

Scaling

Two ways to add GPUs. Strong scaling keeps the global batch fixed and splits it more finely; weak scaling keeps the work per GPU fixed and grows the batch. They behave very differently.

Under strong scaling, four GPUs means one sample per GPU, the parameter all-gather dominates, and per-GPU MFU falls to 12.5%. Weak scaling keeps every device fed and reaches 2.85× on four GPUs, 980 TFLOP/s at 24.8% MFU, against strong scaling's 1.44×.

The gap from ideal is FSDP communication that does not overlap compute. Tuning XLA's overlap flags recovers a little of it, +2.3% at two GPUs and +3.9% at four, growing with device count. Recent compilers already overlap most collectives, so the remainder is largely communication that cannot be hidden rather than a tuning miss.

Attention scales as sequence squared

At pi0's current sequence length, attention is a small slice of the step. But its cost grows with the square of the sequence, so it does not stay small: doubling the tokens roughly quadruples the attention work.

Honest limitations

Three things are deliberately not claimed. Cross-process NCCL collectives hung on the test node, which is consistent with an environment issue rather than a code defect, so genuine multi-node training is not yet demonstrated. RLDS-based data resume is coarse rather than exact. And no full convergence run has been done, since that needs base checkpoints and a real dataset. The systems groundwork for each is in place; the demonstrations are future work.

What it demonstrates

Sharder is a small contribution with a specific shape: it makes openpi trainable across devices and hosts without disturbing the single-host path, and it does so on the back of measurement rather than intuition. The intellectual core is three ideas: bit-precise parity between the multi-process and single-process runs, derive-don't-store resume that needs no iterator checkpoint, and two profiling corrections that turned a misleading picture into an actionable one.

The full source, the measurement scripts, and every figure on this page are in the repository.

→ github.com/virkvarjun/OpenPI-GPU

@software{virk2026sharder,
  author = {Virk, Arjun},
  title  = {Sharder: Fault-Tolerant Multi-Host JAX Training for openpi},
  year   = {2026},
  url    = {https://github.com/virkvarjun/OpenPI-GPU}
}