Orbax Checkpointing in Keras

Author: Samaneh Saadat, Amit Srivastava Date created: 2025/11/20 Last modified: 2026/03/04 Description: Save and load Orbax checkpoints with distributed resharding.
Orbax is the recommended checkpointing library for the JAX ecosystem. It provides high-level functionality for checkpoint management, composable serialization, and multi-host coordination.
Starting with Keras 3.14, the built-in keras.callbacks.OrbaxCheckpoint callback makes it easy to:
Install the Orbax checkpointing library:
Set the Keras backend to JAX, configure virtual devices for the distributed demo, and import the required libraries.
OrbaxCheckpoint works like ModelCheckpoint — pass it as a callback to model.fit(). No boilerplate classes or wrappers are needed.
The checkpoint directory now contains a step-directory for each saved epoch.
Use keras.saving.load_model() to reload a full model (config + weights + optimizer state) from an Orbax checkpoint directory.
If you already have a model instance and just want to load the weights, use load_weights():
When you pass OrbaxCheckpoint to a new fit() call, training resumes from the correct step number. Pass initial_epoch to continue the epoch count from where you left off.
Monitor a metric and keep only the best checkpoint:
Save every N batches instead of every epoch by setting save_freq to an integer:
OrbaxCheckpoint works seamlessly with the Keras Distribution API. When you set a distribution, model variables are automatically sharded according to your LayoutMap and checkpoints capture the distributed state.
Note: The example below uses virtual devices to simulate a multi-device environment. In production, use your actual accelerators (GPUs/TPUs).
After training, you can inspect the sharding of saved variables:
One of the most powerful features of Orbax checkpointing in Keras is the ability to load a checkpoint saved with one sharding layout and restore it under a different layout. Keras automatically provides the target shardings to Orbax so arrays are resharded on load.
For example, a checkpoint saved with dense_1/kernel sharded on the "model" axis can be loaded with that same kernel sharded on the "data" axis instead:
The dense_1/kernel variable was originally sharded as (None, "model") but is now sharded as ("data", None). Orbax handles the data movement automatically during loading.