Picture this. You have spent six hours training a model. The loss curve looks beautiful, accuracy is climbing, and you are one epoch away from a result worth writing home about. Then the power goes out. Or the Colab runtime disconnects. Or your cloud instance gets preempted because you are on a spot pricing tier to save money. Six hours of compute, gone, with nothing to show for it but a chart that no longer exists anywhere except your memory of watching it.
This is not a hypothetical for anyone building AI systems outside of a well funded lab with guaranteed, uninterrupted access to a GPU cluster. Power instability is a genuine constraint in many parts of the world. Free tier compute gets reclaimed without warning. Long training runs get interrupted by things entirely outside your control.
The solution is checkpointing: periodically saving the complete state of your training run to disk, so that if something goes wrong, you can restore exactly where you left off instead of starting over from scratch.
Today we cover Orbax, the checkpointing library built for the JAX ecosystem. Orbax does for training state what Git does for code: it gives you reliable, versioned snapshots you can always return to. By the end of this article you will understand:
Why checkpointing in JAX requires a different approach than torch.save()
How Flax NNX exposes model state as something Orbax can save
How to save and restore model parameters using CheckpointManager
How to save multiple pieces of state together, such as both model and optimizer
How checkpointing works when your model is sharded across multiple devices
How to build a complete, crash resilient training loop
Let's make sure we never lose six hours of compute again.
Why checkpointing in JAX is different
In PyTorch, saving a model is famously simple:
torch.save(model.state_dict(), "checkpoint.pt")
This works because state_dict() returns a Python dictionary of tensors, and PyTorch's serialization just pickles that dictionary to disk. Loading is equally simple: build the model architecture, then call model.load_state_dict(torch.load("checkpoint.pt")).
JAX, as we have discussed since week 1, has no hidden state. There is no single object holding "the model" that you can wave a wand over and serialize. What exists instead are pytrees, explicit, nested structures of arrays that your training loop passes around by hand.
This has real advantages once you understand it. Because JAX state is always an explicit, structured pytree, checkpointing can be done with the same rigor and predictability as any other JAX operation. There is no hidden magic to go wrong. But it does mean the mental model is different, and that difference is worth understanding before writing any code.
Orbax is designed around this reality. It does not try to serialize an opaque model object. It serializes pytrees, the same fundamental data structure we have been working with since week 1.
From Flax NNX object to Pytree: nnx.state
We spent week 3 learning that Flax NNX gives us PyTorch style, object oriented models while secretly managing pure JAX state underneath. Checkpointing is where that "secretly" becomes explicit and useful.
The bridge function is nnx.state(). Given an NNX module, it extracts every nnx.Variable inside it, including all nnx.Param weights, into a single pytree called an nnx.State. This is exactly the kind of structure Orbax knows how to save.
from flax import nnx
model = CNN(rngs=nnx.Rngs(0))
# Extract the model's state as a pytree
state = nnx.state(model)
There is a closely related pair of functions worth knowing, nnx.split and nnx.merge, which we previewed back in week 3:
# Split separates structure (GraphDef) from data (State)
graphdef, state = nnx.split(model)
# Merge reconstructs a live module from structure and data
reconstructed_model = nnx.merge(graphdef, state)
The GraphDef captures the static architecture, the layout of layers, and it never changes during training. The State captures the actual numbers, the weights that update every step. Orbax only ever needs to save the State. The GraphDef can always be recreated simply by instantiating your model class again in code.
This distinction, static architecture versus dynamic state, is the conceptual key to everything that follows.
The core Orbax components
Orbax is typically imported under the alias ocp:
import orbax.checkpoint as ocp
There are a few pieces worth knowing by name.
ocp.CheckpointManager is the component you will use in almost every real project. It manages an entire directory of checkpoints across many training steps: deciding when to save, how many old checkpoints to retain, and how to find the most recent one when you want to resume.
mngr = ocp.CheckpointManager(
directory="/path/to/checkpoints",
options=ocp.CheckpointManagerOptions(
max_to_keep=3,
save_interval_steps=100,
),
)
ocp.Checkpointer, and specifically StandardCheckpointer, is the lower level component that actually handles serialization of a single pytree. In most workflows you will not touch this directly. CheckpointManager uses it internally.
ocp.args is a namespace of argument objects that tell Orbax exactly what kind of save or restore operation you are performing.
ocp.args.StandardSave(pytree) saves a standard pytree.
ocp.args.StandardRestore(abstract_pytree) restores a pytree, and requires you to supply an abstract structure describing what shape and dtype you expect back.
ocp.args.Composite(**kwargs) lets you save or restore several distinctly named items together, for example model parameters and optimizer state as two separate named entries in a single checkpoint.
We will use all three shortly.
Saving model state: A first example
Let's start with the simplest possible case: saving just the model's parameters at a given training step.
import orbax.checkpoint as ocp
from flax import nnx
# Assume model is an initialized nnx.Module
mngr = ocp.CheckpointManager("/tmp/my_checkpoints")
# Extract the state
state_to_save = nnx.state(model)
# Save at step 100
mngr.save(100, args=ocp.args.StandardSave(state_to_save))
# IMPORTANT: wait for the (possibly asynchronous) save to actually finish
mngr.wait_until_finished()
mngr.close()
Two details here deserve emphasis.
First, mngr.wait_until_finished(). Orbax can save asynchronously in the background, which means your training loop does not have to stall while gigabytes of weights are written to disk. This is a meaningful performance win for large models. But it also means that if your program exits immediately after calling save(), the write might not have actually completed yet. Always call wait_until_finished() before your program ends, or before you rely on that checkpoint being present.
Second, mngr.close(). Like any resource that manages file handles and background threads, close it cleanly when you are done, particularly at the very end of a training script.
Restoring model state
Restoring is slightly more involved than saving, because Orbax needs to know the exact shape and dtype structure it is restoring into before it reads a single byte from disk. This is where nnx.eval_shape becomes essential.
nnx.eval_shape builds an "abstract" version of your model: an object with the exact same structure as a real model, but where every array is replaced by a lightweight ShapeDtypeStruct describing its shape and dtype, with no actual memory allocated. This gives Orbax a blueprint to restore into, without you needing to fully instantiate a real model first, which matters enormously for very large models where a wasted allocation could mean an out of memory error before you have even loaded anything.
# Build an abstract model, no real memory allocated
abstract_model = nnx.eval_shape(lambda: CNN(rngs=nnx.Rngs(0)))
# Split into structure and abstract state
graphdef, abstract_state = nnx.split(abstract_model)
# Ask the manager for the latest available step
step_to_restore = mngr.latest_step()
if step_to_restore is not None:
restored_state = mngr.restore(
step_to_restore,
args=ocp.args.StandardRestore(abstract_state),
)
# 4. Reconstruct a live, usable model
restored_model = nnx.merge(graphdef, restored_state)
mngr.close()
mngr.latest_step() is a small but genuinely useful convenience. Rather than hardcoding which step to restore, it automatically finds the most recent checkpoint written to the directory, which is exactly what you want in the common case of "resume from wherever I last left off."
There is also a second restoration pattern worth knowing: updating an existing model object in place, rather than creating a brand new one.
existing_model = CNN(rngs=nnx.Rngs(0))
nnx.update(existing_model, restored_state)
Use nnx.merge when you are building a fresh model from nothing. Use nnx.update when you already have a live model object and simply want to overwrite its weights with restored values, for example loading a pretrained checkpoint into a model you have already partially configured.
Saving more than just the model
A trained model alone is not enough to resume training seamlessly. If you restore only the weights and restart with a fresh optimizer, you lose Adam's momentum and variance estimates, effectively resetting your optimizer's memory of the training trajectory so far. You also lose your exact step count, which matters if you are using a learning rate schedule tied to step number, as we built in Week 5.
For a real resumable training run, you want to checkpoint model parameters and optimizer state together. This is what ocp.args.Composite is for.
Saving composite state
import optax
from flax import nnx
tx = optax.adamw(learning_rate=0.001)
optimizer = nnx.Optimizer(model, tx=tx, wrt=nnx.Param)
# Extract each piece of state separately
params_state = nnx.split(optimizer, nnx.Param)[1]
optimizer_state = nnx.state(optimizer)
save_items = {
"params": ocp.args.StandardSave(params_state),
"optimizer": ocp.args.StandardSave(optimizer_state),
}
mngr.save(
optimizer.step.value,
args=ocp.args.Composite(**save_items),
)
mngr.wait_until_finished()
Notice that we key the step number off optimizer.step.value. NNX optimizers track their own step count internally, which is a convenient single source of truth for "how far along is this training run."
Restoring composite state
Restoration follows the same abstract shape pattern as before, but now for two named items instead of one:
# Build abstract model and optimizer
abs_model = nnx.eval_shape(lambda: CNN(rngs=nnx.Rngs(0)))
abs_optimizer = nnx.eval_shape(
lambda: nnx.Optimizer(abs_model, optax.adamw(0.001), wrt=nnx.Param)
)
graphdef, abs_params_state = nnx.split(abs_model, nnx.Param)
abs_optimizer_state = nnx.state(abs_optimizer)
# Define what we're restoring into
restore_targets = {
"params": ocp.args.StandardRestore(abs_params_state),
"optimizer": ocp.args.StandardRestore(abs_optimizer_state),
}
# Restore
step = mngr.latest_step()
restored = mngr.restore(step, args=ocp.args.Composite(**restore_targets))
# Build real, live instances and update them with restored data
model_instance = CNN(rngs=nnx.Rngs(1))
optimizer_instance = nnx.Optimizer(model_instance, optax.adamw(0.001), wrt=nnx.Param)
nnx.update(model_instance, restored["params"])
nnx.update(optimizer_instance, restored["optimizer"])
mngr.close()
Notice something subtle here: we initialize model_instance with a fresh nnx.Rngs(1), a completely different seed than we originally used. This is intentional and it demonstrates an important point. The random seed only matters for the initial values before restoration. Once nnx.update overwrites those values with the checkpointed weights, the original seed used to create the model is irrelevant. What matters is that the abstract structure, the shapes and dtypes, matches exactly.
A complete, resumable training loop
Let's bring this together into something you would actually deploy: a training loop that checks for an existing checkpoint on startup, resumes if one exists, and saves periodically as it trains.
import jax
import jax.numpy as jnp
from flax import nnx
import optax
import orbax.checkpoint as ocp
CHECKPOINT_DIR = "/tmp/resilient_training"
SAVE_EVERY = 100
NUM_STEPS = 2000
def build_model_and_optimizer(seed=0):
model = CNN(rngs=nnx.Rngs(seed))
tx = optax.adamw(learning_rate=0.001)
optimizer = nnx.Optimizer(model, tx=tx, wrt=nnx.Param)
return model, optimizer
@nnx.jit
def train_step(model, optimizer, batch):
def loss_fn(model):
logits = model(batch["image"])
loss = optax.softmax_cross_entropy_with_integer_labels(
logits=logits, labels=batch["label"]
).mean()
return loss, logits
(loss, logits), grads = nnx.value_and_grad(loss_fn, has_aux=True)(model)
optimizer.update(model, grads)
return loss
def main():
mngr = ocp.CheckpointManager(
CHECKPOINT_DIR,
options=ocp.CheckpointManagerOptions(max_to_keep=3, save_interval_steps=SAVE_EVERY),
)
model, optimizer = build_model_and_optimizer(seed=0)
start_step = 0
# Attempt to resume
latest_step = mngr.latest_step()
if latest_step is not None:
print(f"Found existing checkpoint at step {latest_step}. Resuming.")
abs_model = nnx.eval_shape(lambda: build_model_and_optimizer(seed=0)[0])
abs_optimizer = nnx.eval_shape(
lambda: build_model_and_optimizer(seed=0)[1]
)
graphdef, abs_params_state = nnx.split(abs_model, nnx.Param)
abs_optimizer_state = nnx.state(abs_optimizer)
restored = mngr.restore(
latest_step,
args=ocp.args.Composite(
params=ocp.args.StandardRestore(abs_params_state),
optimizer=ocp.args.StandardRestore(abs_optimizer_state),
),
)
nnx.update(model, restored["params"])
nnx.update(optimizer, restored["optimizer"])
start_step = latest_step + 1
else:
print("No checkpoint found. Starting fresh.")
# Training loop
for step in range(start_step, NUM_STEPS):
# In a real project this batch would come from the Grain
# pipeline built in Week 8
dummy_batch = {
"image": jnp.ones((32, 28, 28, 1)),
"label": jnp.zeros((32,), dtype=jnp.int32),
}
loss = train_step(model, optimizer, dummy_batch)
if step % SAVE_EVERY == 0:
params_state = nnx.split(optimizer, nnx.Param)[1]
optimizer_state = nnx.state(optimizer)
mngr.save(
step,
args=ocp.args.Composite(
params=ocp.args.StandardSave(params_state),
optimizer=ocp.args.StandardSave(optimizer_state),
),
)
print(f"Step {step:4d} | Loss: {loss:.4f} | Checkpoint saved")
mngr.wait_until_finished()
mngr.close()
print("Training complete.")
if __name__ == "__main__":
main()
Run this once, kill it midway through with Ctrl+C, and run it again. You will see it detect the existing checkpoint and resume from exactly where it left off, rather than starting over. This single pattern is what separates a fragile prototype from a production ready training pipeline.
Sharded checkpointing for distributed training
Everything above works cleanly for a model that fits comfortably on a single device. But we are building toward week 10, where we will shard models across multiple GPUs or TPUs. Checkpointing sharded state introduces one additional wrinkle worth understanding now, so it is not a surprise later.
When saving, there is genuinely nothing special to do. If your state is already sharded, for example because it was produced inside a jax.jit call executed within a Mesh context, Orbax transparently handles writing each device's shard to the correct location on disk.
mngr.save(step, args=ocp.args.StandardSave(sharded_state_pytree))
Restoring sharded state is where care is required. The abstract target you provide to StandardRestore must include the sharding information describing how the restored arrays should be distributed across devices, not just their shape and dtype.
def create_abstract_sharded_target():
abstract_model = nnx.eval_shape(lambda: CNN(rngs=nnx.Rngs(0)))
_, abstract_state = nnx.split(abstract_model)
sharding_specs = nnx.get_partition_spec(abstract_state)
return jax.lax.with_sharding_constraint(abstract_state, sharding_specs)
with mesh: # a jax.sharding.Mesh, covered fully in Week 10
abstract_target = jax.jit(create_abstract_sharded_target)()
restored_sharded_state = mngr.restore(
step,
args=ocp.args.StandardRestore(abstract_target),
)
This matters because when you restore a large sharded model, you do not want Orbax to load the entire thing onto a single device and then redistribute it, which could easily exceed that single device's memory. By providing sharding information up front, Orbax can restore each shard directly to the correct device, keeping memory usage proportional to what a single device actually needs to hold.
It is also worth flagging a real world migration detail. Newer versions of Flax NNX changed how checkpoints structure RNG state, particularly for layers like Dropout or BatchNorm that carry their own random state. If you are restoring a checkpoint saved with an older NNX version into a newer one, this can require a specific migration step. Keep this in mind if you are ever debugging a restore that fails only for models containing stateful, randomness dependent layers.
Other Orbax features worth knowing
Asynchronous checkpointing is enabled by default in CheckpointManager and is what allows mngr.save() to return almost immediately while the actual write happens on a background thread. This is why wait_until_finished() exists and why it matters.
Atomicity is guaranteed by CheckpointManager. A checkpoint write either completes fully or is treated as if it never happened. You will never end up restoring a half written, corrupted checkpoint, which is exactly the kind of failure mode you most want to avoid after a crash.
The TensorStore backend is what Orbax often uses under the hood for efficient array serialization, particularly valuable for very large arrays and for cloud storage backends where efficient, chunked I/O matters.
Common pitfalls
A few mistakes are common enough to name directly.
Forgetting wait_until_finished() before the program exits. Because saves can be asynchronous, a script that calls save() and then immediately exits can leave you with an incomplete or missing checkpoint.
Restoring with the wrong abstract structure. If your abstract_state does not exactly match the structure of what you saved, for example because you changed your model's architecture between the save and the restore, Orbax will raise a clear error rather than silently loading garbage. Treat that error as useful information, not an obstacle to work around.
Not setting max_to_keep. Without a limit, CheckpointManager will happily keep every single checkpoint you ever save, which for a large model can consume enormous amounts of disk space over a long training run. Set max_to_keep deliberately based on how much history you actually need.
Saving too frequently. Checkpointing has real cost, both in I/O and in the memory required to stage the write. Saving every single step is rarely necessary. A cadence tied to a meaningful interval, such as every hundred or every thousand steps, is usually the right balance.
Exercises
Simulate a crash. Take the resumable training loop above, let it run for a few hundred steps, then interrupt it. Restart it and confirm the loss picks up from a similar value rather than the high loss of a freshly initialized model.
Inspect a checkpoint on disk. After saving a checkpoint, use view or your terminal to look at the directory structure Orbax creates. Note how it organizes different steps and named items.
Add Grain iterator state. Using what you learned in Week 8, extend the training loop to also checkpoint the Grain data iterator's position, so that resuming a training run does not repeat or skip any batches.
Test a shape mismatch. Deliberately change your model's hidden_dim after saving a checkpoint, then try to restore into the new architecture. Read the resulting Orbax error carefully and confirm you understand exactly what it is telling you.
Quick reference
import orbax.checkpoint as ocp
from flax import nnx
# Manager Setup
mngr = ocp.CheckpointManager(
directory,
options=ocp.CheckpointManagerOptions(max_to_keep=3, save_interval_steps=100),
)
# Saving a Single Pytree
state = nnx.state(model)
mngr.save(step, args=ocp.args.StandardSave(state))
mngr.wait_until_finished()
# Restoring a Single Pytree
abstract_model = nnx.eval_shape(lambda: MyModel(rngs=nnx.Rngs(0)))
graphdef, abstract_state = nnx.split(abstract_model)
restored_state = mngr.restore(
mngr.latest_step(), args=ocp.args.StandardRestore(abstract_state)
)
restored_model = nnx.merge(graphdef, restored_state)
# Composite Save/Restore (model + optimizer)
mngr.save(step, args=ocp.args.Composite(
params=ocp.args.StandardSave(params_state),
optimizer=ocp.args.StandardSave(optimizer_state),
))
restored = mngr.restore(step, args=ocp.args.Composite(
params=ocp.args.StandardRestore(abs_params_state),
optimizer=ocp.args.StandardRestore(abs_optimizer_state),
))
# Updating Live Objects
nnx.update(model, restored["params"])
nnx.update(optimizer, restored["optimizer"])
# Cleanup
mngr.close()
What's next
We can now train models that survive crashes, power outages, and preempted cloud instances. But everything we have built so far assumes a single device. Real large scale models, the kind used for foundational speech to text systems or large language models, do not fit on one GPU at all.
Next week we tackle sharding and parallelism. We will learn how to split a model across multiple devices using jax.sharding.Mesh and PartitionSpec, how to combine this with the checkpointing techniques from today, and how to think about data parallelism, tensor parallelism, and fully sharded data parallelism as complementary strategies rather than competing ones.
Next week: Splitting the Load: Sharding Across Devices
