What is the Orbax V1 API?

What is the Orbax V1 API?#

Orbax has developed a new API, called V1, which can be accessed using:

from orbax.checkpoint import v1 as ocp

This contrasts with the legacy API, called V0, which is comprised of all other APIs located at the top level orbax.checkpoint package.

Motivation#

Users have often complained that the V0 API has excessive complexity, verbosity, and unintuitive interfaces. We have introduced the V1 API centered around strong support for core use cases, a more opinionated interface, consolidation of advanced configuration options, and removal of unused customization layers.

The new API will be made available at the top-level orbax.checkpoint package in a future version 1.0.0. We do not expect to make this change before late 2026; due to widespread adoption of the legacy API, large-scale migration is difficult.

Core Tasks#

Here is a comparison of core checkpointing functionalities and their entry points in both APIs:

Task

V0 Interface

V1 Interface

Checkpointing in a training loop (sequence of steps)

CheckpointManager

v1.training.Checkpointer

Checkpointing with an arbitrary path

Checkpointer / AsyncCheckpointer

v1.save/v1.load

Customization

CheckpointHandler

v1.CheckpointableHandler / v1.StatefulCheckpointable

Configuration

(no single solution)

v1.Context

 
import orbax.checkpoint as ocp
import jax
from jax import numpy as jnp
from etils import epath
import grain

root_v0 = epath.Path('/tmp/what-is-v1/v0')
root_v1 = epath.Path('/tmp/what-is-v1/v1/')
state = {'a': jnp.arange(8), 'b': jnp.ones(8)}
abstract_state = jax.tree.map(ocp.arrays.to_shape_dtype_struct, state)
dataset = iter(grain.MapDataset.range(30).batch(3).map(lambda x: x.tolist()))

Comparison#

Eliminates unnecessary objects and moves closer toward a purely functional interface.

### V0 ###

ckptr = ocp.PyTreeCheckpointer()
ckptr.save(root_v0 / 'ckpt1', state)

### V1 ###

ocp.v1.save(root_v1 / 'ckpt1', state)

Allows for easily handling distinct checkpointables with minimal verbosity, including objects (like Grain iterators) not defined by Orbax.

### V0 ###

ckptr = ocp.Checkpointer(ocp.CompositeCheckpointHandler())
ckptr.save(
    root_v0 / 'ckpt2',
    args=ocp.args.Composite(
        state=ocp.args.PyTreeSave(state),
        dataset=grain.checkpoint.CheckpointSave(dataset),
    )
)

### V1 ###

ocp.v1.save_checkpointables(
    root_v1 / 'ckpt2', dict(state=state, dataset=dataset))

Provides easy mechanisms for configuring specialized saving and loading behavior via Context. Eliminates duplicative and

### V0 ###

ckptr = ocp.PyTreeCheckpointer()
partial_item = {**abstract_state}
del partial_item['b']
restore_args = {'a': ocp.ArrayRestoreArgs()}
ckptr.restore(
    root_v0 / 'ckpt1',
    args=ocp.args.PyTreeRestore(
        item=partial_item,
        restore_args=restore_args,
        partial_restore=True
    )
)

/home/docs/checkouts/readthedocs.org/user_builds/orbax/envs/stable/lib/python3.12/site-packages/orbax/checkpoint/_src/serialization/jax_array_handlers.py:749: UserWarning: Sharding info not provided when restoring. Populating sharding info from sharding file. Please note restoration time will be slightly increased due to reading from file. Note also that this option is unsafe when restoring on a different topology than the checkpoint was saved with.
  warnings.warn(

{'a': Array([0, 1, 2, 3, 4, 5, 6, 7], dtype=int32)}

### V1 ###

ctx = ocp.v1.Context()
ctx.pytree_options.loading.partial_load = True
partial_abstract_state = {**abstract_state}
del partial_abstract_state['b']
ocp.v1.load(root_v1 / 'ckpt1', )

WARNING:absl:TensorStore data files not found in checkpoint path /tmp/what-is-v1/v1/ckpt1. This may be a sign of a malformed checkpoint, unless your checkpoint consists entirely of strings or other non-standard PyTree leaves.

{'a': Array([0, 1, 2, 3, 4, 5, 6, 7], dtype=int32),
 'b': Array([1., 1., 1., 1., 1., 1., 1., 1.], dtype=float32)}

The above is not an exhaustive list, but we hope it gives an intutitive sense of the improvements offered by the new API. For more context, we previously shared a design proposal, which can provide a helpful historical perspective, though do bear in mind that the actual implementation has evolved somewhat.

Contents