# Never lose your progress: Checkpointing with Orbax

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:

1.  Why checkpointing in JAX requires a different approach than `torch.save()`
    
2.  How Flax NNX exposes model state as something Orbax can save
    
3.  How to save and restore model parameters using `CheckpointManager`
    
4.  How to save multiple pieces of state together, such as both model and optimizer
    
5.  How checkpointing works when your model is sharded across multiple devices
    
6.  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:

```python
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](https://kambale.dev/why-jax-the-numpy-you-know-but-faster), 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](https://kambale.dev/why-jax-the-numpy-you-know-but-faster).

## From Flax NNX object to Pytree: nnx.state

We spent [week 3](https://kambale.dev/flax-nnx) 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.

```python
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](https://kambale.dev/flax-nnx):

```python
# 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`:

```python
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.

```python
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.

```python
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.

```python
# 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.

```python
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

```python
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:

```python
# 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.

```python
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.

```python
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.

```python
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

1.  **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.
    
2.  **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.
    
3.  **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.
    
4.  **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

```python
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*