How does Orbax decide which host saves each shard of a checkpoint?

By computation, not communication. Every host runs the same deterministic planner (get_replica_slices in replica_slices.py) over the same sharding metadata it was already handed. In the default single-replica policy, a host keeps an addressable shard only if the shard's replica_id matches the target replica id — so of all the identical copies of a shard across the cluster, exactly one replica's copy survives the filter and gets written. No election, no lock, no RPC: identical inputs to the same function produce identical answers on every host.

What is replica-parallel checkpointing in Orbax?

A cooperative mode where, instead of one replica writing a whole shard while the others idle, ownership of each shard is split evenly across the R replicas that hold it — each writes 1/R. The file finds the first shard dimension evenly divisible by the replica count, divides it into R equal bands, and each replica takes the band starting at replica_id × band_size. The offsets tile the shard exactly, with no overlap and no gap, and each host computes its own band locally from its replica id.

How can hosts agree on checkpoint ownership without any communication?

Because everything the decision needs is already replicated identically on every process. A jax.sharding.Sharding's devices_indices_map says which slice of the global array every device holds; counting devices that map to the same slice gives the replica count; shard replica ids are assigned deterministically by JAX. Ownership is then a pure function of that shared metadata, so every host derives the same global plan by running the same code — agreement by construction rather than by consensus protocol.

What happens if no dimension divides evenly by the replica count?

The planner degrades honestly instead of erroring. calculate_replica_parallel_axis_and_local_shape scans the shard's dimensions for one where axis_size % replica_count == 0; if none exists (say a (4096, 4096) shard across 3 replicas), it returns nothing, maybe_pick_replica_parallel returns None, and the caller falls back to single-replica — one replica writes the whole shard and the rest write nothing. The shard is still written exactly once; you just lose the parallelism. A code TODO notes the future fix: let the last replica carry a smaller slice.

Why not just let every replica write its copy of the shard?

Cost and correctness. A tensor replicated across R hosts written R times is R× the storage egress and R× the request pressure on the object store — and R concurrent writers racing on the same object. At training scale (thousands of arrays, saved every few minutes from hundreds of hosts) that multiplies real money and real failure modes. The whole point of replica_slices.py is that deduplicating writers costs nothing: the filter is one comparison per shard against metadata each host already holds.

Zero-RPC Sharding: How 1,000 Hosts Agree Who Writes Which Bytes

At checkpoint time every data-parallel replica holds the same shard — so who saves it? We walk the real Orbax planner that assigns ownership with zero coordination: each host runs one deterministic function over the sharding it already has, and every unique shard gets exactly one writer. Single-replica, replica-parallel, and the divisibility check that decides between them.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of Zero-RPC Sharding: How 1,000 Hosts Agree Who Writes Which Bytes — built to be understood, not skimmed.

Questions

How does Orbax decide which host saves each shard of a checkpoint?: By computation, not communication. Every host runs the same deterministic planner (get_replica_slices in replica_slices.py) over the same sharding metadata it was already handed. In the default single-replica policy, a host keeps an addressable shard only if the shard's replica_id matches the target replica id — so of all the identical copies of a shard across the cluster, exactly one replica's copy survives the filter and gets written. No election, no lock, no RPC: identical inputs to the same function produce identical answers on every host.
What is replica-parallel checkpointing in Orbax?: A cooperative mode where, instead of one replica writing a whole shard while the others idle, ownership of each shard is split evenly across the R replicas that hold it — each writes 1/R. The file finds the first shard dimension evenly divisible by the replica count, divides it into R equal bands, and each replica takes the band starting at replica_id × band_size. The offsets tile the shard exactly, with no overlap and no gap, and each host computes its own band locally from its replica id.
How can hosts agree on checkpoint ownership without any communication?: Because everything the decision needs is already replicated identically on every process. A jax.sharding.Sharding's devices_indices_map says which slice of the global array every device holds; counting devices that map to the same slice gives the replica count; shard replica ids are assigned deterministically by JAX. Ownership is then a pure function of that shared metadata, so every host derives the same global plan by running the same code — agreement by construction rather than by consensus protocol.
What happens if no dimension divides evenly by the replica count?: The planner degrades honestly instead of erroring. calculate_replica_parallel_axis_and_local_shape scans the shard's dimensions for one where axis_size % replica_count == 0; if none exists (say a (4096, 4096) shard across 3 replicas), it returns nothing, maybe_pick_replica_parallel returns None, and the caller falls back to single-replica — one replica writes the whole shard and the rest write nothing. The shard is still written exactly once; you just lose the parallelism. A code TODO notes the future fix: let the last replica carry a smaller slice.
Why not just let every replica write its copy of the shard?: Cost and correctness. A tensor replicated across R hosts written R times is R× the storage egress and R× the request pressure on the object store — and R concurrent writers racing on the same object. At training scale (thousands of arrays, saved every few minutes from hundreds of hosts) that multiplies real money and real failure modes. The whole point of replica_slices.py is that deduplicating writers costs nothing: the filter is one comparison per shard against metadata each host already holds.

Zero-RPC Sharding: How 1,000 Hosts Agree Who Writes Which Bytes

Questions

Related explainers