Orbax Checkpoint

Orbax is a modular and customizable JAX checkpointing library designed for high-performance, large-scale model persistence and recovery. It provides a JAX-native approach to managing distributed array storage and checkpoint lifecycles.

Key capabilities include:

• Distributed Checkpointing: Unified API for single- and multi-process checkpointing, abstracting the complexities of persisting distributed arrays.
• High Performance: Memory-efficient and fast checkpointing designed to minimize impact on training loops.
• Lifecycle Management: Facilitates training loop integration through metadata management, garbage collection, and saving policies.
• Advanced Workflows: Supports topology-agnostic loading (resharding), partial loading, and incremental saving.
• Extensibility: Customizable handler interfaces for user-defined types and logic.
• Model Exporting: Through the companion library orbax-export, JAX models can be exported to the TensorFlow SavedModel format.

Orbax v1 introduces a fundamental shift from flat configuration flags to structured policy composition.

In v0, you passed independent arguments (like save_interval_steps or max_to_keep) to the CheckpointManager, which used internal priority logic to resolve conflicts. In v1, you must explicitly compose your logic using Policy Objects and pass them to the Checkpointer constructor.

Key differences:

• No Implicit Priority: In v1, the Checkpointer executes exactly what you compose. If you want multiple rules (e.g., 'save every 10 steps' AND 'save on preemption'), you must explicitly combine them using an AnyPolicy wrapper.
• Explicit Composition: Instead of setting multiple flags, you construct a logic tree of policies.
• Context vs. Checkpointer: Environmental/IO settings (like async or multiprocessing) move to ocp.Context, while training lifecycle logic (save/preserve rules) moves to the Checkpointer constructor.

The ocp.Context manages environment-specific and detailed I/O settings. It is intended to be set once for a training task and used as a context manager to apply settings globally to all operations within that scope.

Key specialized options classes within Context include:

• AsyncOptions: Configures asynchronous checkpoint saving, including timeout_secs and post-finalization callbacks.
• FileOptions: Configures directory and file management (permissions, atomicity, custom paths).
• MultiprocessingOptions: Configures multiprocessing behavior (primary host designation, active process subsets).
• DeletionOptions: Configures checkpoint deletion behavior.
• Component-Specific Options: Granular control for PyTreeOptions and ArrayOptions (storage formats, dtypes, concurrent I/O limits).
• PathwaysOptions: Configures Pathways-specific saving and loading.

Orbax uses SaveDecisionPolicy to determine when a checkpoint should be saved during a training loop. Instead of manually checking conditions, you provide a policy to the checkpoint manager. The policy evaluates a DecisionContext (which contains information about the current step, etc.) to return a boolean indicating whether a save operation should proceed.

Common implementations include:

• FixedIntervalPolicy: Saves every $N$ steps.
• SpecificStepsPolicy: Saves at a predefined list of step numbers.
• ContinuousCheckpointingPolicy: Designed for continuous saving patterns.
• PreemptionCheckpointingPolicy: Handles saving when a preemption signal (like a preemptible VM shutdown) is detected.
• InitialSavePolicy: Ensures a save occurs at the very beginning.
• AnySavePolicy: A catch-all policy.

from orbax.checkpoint.checkpoint_managers import save_decision_policy

# Example: Save every 100 steps
policy = save_decision_policy.FixedIntervalPolicy(interval=100)
# This policy is then passed to your CheckpointManager configuration