Checkpointing

Saving and loading model state at scale: sharded arrays, multihost clusters, and reading exactly the bytes each machine needs.

Explainers

• What "Atomic" Means on a Filesystem vs. an Object Store — A checkpoint must appear all-or-nothing, even if the job dies with one byte left to write. We walk the real Orbax code that guarantees it two ways — an atomic rename on a POSIX filesystem, and a commit_success.txt marker on an object store that has no rename at all.
• 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.
• The B-Tree Merge: Thousands of Checkpoint Files into One Atomic Manifest — Every host in a training run writes its own OCDBT checkpoint files into a per-process subdirectory — no coordination, no contention. Turning those scattered per-process B-trees into one global manifest a reader can trust is a single atomic TensorStore transaction. We walk the real Orbax file that does it: glob the subdirs, stage every copy_range under ts.Transaction(atomic=True), validate the would-be-merged store, then commit once — all of it, or none of it.
• The Checkpoint Lifecycle: What Happens Between save() and Durable — A training run that saves synchronously leaves a rack of accelerators idle for the seconds it takes to push a terabyte to storage. So the save goes async — a fast device-to-host copy, then a background thread that writes, coordinates every host, commits the directory as one atomic unit, and garbage-collects the old ones. Follow one save from the training step that triggers it to the moment it is durable and restorable onto a different mesh — drawn, computed, and animated.
• Loading a Safetensors Checkpoint on a Multihost Cluster, Line by Line — A safetensors file is one flat byte blob per tensor, but a big model must land as sharded jax.Arrays across many hosts. Walk the real loader that maps each host’s shards to byte ranges — and reads exactly those.
• Model Surgery: Rewriting a Checkpoint’s Parameters, Line by Line — A model is saved under one structure and needed under another — HuggingFace names, a split gate/up projection, 64 separately-saved experts. We walk Orbax’s real model-surgery module: six transforms that rename, fuse, repeat, and stack a pytree of weights, and the host-memory care that lets them run on tens of gigabytes.