- The paper introduces Orbax, a modular and performant solution that standardizes distributed checkpointing in the JAX ecosystem.
- It employs a layered API design with synchronous and asynchronous saving to boost throughput and enable efficient model recovery.
- Benchmark results demonstrate up to 3.5× speedup in background saving and improved loading speeds, ensuring scalability across multi-device environments.
Orbax: Distributed Checkpointing with JAX
Motivation and Context
The JAX ecosystem, despite its widespread adoption for high-performance ML research, historically lacked a standardized, robust solution for distributed checkpointing. This gap exacerbated fragmentation among JAX codebases and constrained the utilization of advanced distributed accelerator hardware. The absence of a unified checkpointing infrastructure led to idiosyncratic, often non-scalable implementations, limiting recovery performance, usability for rapid iterations, and experimentation with model manipulation techniques. To address this, Orbax is introduced as a modular, performant, and JAX-native checkpointing library, designed to abstract distributed storage complexities while enabling flexible checkpoint orchestration.
Existing checkpointing solutions in JAX—such as Flax, T5X, and Pax—were generally constrained by lack of scalability, object size limits, and fragmented APIs. TensorStore-based integrations provided low-level access but required considerable manual distributed coordination. PyTorch, by contrast, offers the DCP library for distributed checkpointing and NeMo's FPS strategy; however, these are not natively accessible within JAX workflows and impose user-side orchestration for step management. Pathways' closed-source checkpointing and other bespoke frameworks further underscored the need for a standardized, open approach within JAX.
Architecture and API
Orbax's API is designed around progressive disclosure and modularity. The central abstractions are:
- Sequence-of-Steps Checkpointing: The
training.Checkpointer interface supports step-level orchestration, allowing object-oriented management of step histories, ongoing saves, and curation tasks (e.g., garbage collection).
- Singular Checkpoints: Purely functional interfaces enable direct manipulation and debugging via explicit checkpoint file paths.
- Checkpointables: Fine-grained, independently-serializable components of the training state, supporting separability and distinct serialization logic (model parameters, optimizer state, dataset iterators).
- CheckpointableHandler & StatefulCheckpointable: Interfaces allow custom serialization behaviors, enhancing modularity and easing integration with external components (e.g., Grain dataset iterators).
- Abstract vs. Concrete Checkpointables: Abstract types record metadata (shape, dtype, sharding) to formalize the contract between model structure and checkpoint disk representation, supporting flexible restoration and partial loading.
This layered approach allows developers to both exploit high-level orchestration and exercise precise control over checkpoint lifecycle, customization, and compatibility across codebases and storage backends.
Saving and Loading Logic
Orbax's operational pipeline is optimized for minimizing device idle time and maximizing storage I/O throughput:
- Saving: Divided into synchronous (blocking) and asynchronous phases. Validation and device-to-host copies occur synchronously; all expensive storage I/O is delegated to background threads. Atomicity is enforced via temporary directory usage and finalization by a leader host, ensuring robustness against interruptions and race conditions.
- Metadata Management: Global metadata (key names, structures) is independently managed per process, then merged and validated on finalization. Serialization leverages TensorStore drivers (Zarr, OCDBT) for scalable chunking and metadata indexing, accommodating diverse model scales.
- Loading: Restoration is guided by user-provided abstract state or checkpoint metadata, supporting dynamic casting, reshaping, and topology-aware resharding. Partial loading and placeholder fill-in (for missing keys) are supported, enabling selective recovery and manipulation.
Multi-controller and single-controller (Pathways) coordination is transparently handled, leveraging JAX's Colocated Python API for high-bandwidth data transfers and reducing the central controller's role as an I/O bottleneck.
Benchmarking against PyTorch's DCP library using large-scale Llama models (8B, 70B, 405B) demonstrates several key results:
- Orbax achieves up to 3.5× speedup in background saving and 2× speedup in loading for LLM-scale checkpoints on A100 hardware with GCS backend.
- OCDBT-backed checkpointing yields increased throughput for large models and topologies (64 devices, 405B), with up to 1.34× speedup.
- Subchunked checkpoint representation enables up to 5.68× faster resharded loading for model migration across devices and sharding configurations.
- Replica-parallel saving distributes I/O load across slices, up to 2.39× speedup in multi-slice environments.
- Broadcast-based loading reduces duplicated storage reads, up to 4.93× speedup in large-scale multi-slice settings.
These strong numerical results validate Orbax's scalability and efficiency relative to established PyTorch solutions and prior JAX-specific implementations.
Ecosystem Compatibility and Usability
Orbax formalizes JAX checkpointing, minimizing fragmentation by standardizing serialization formats and APIs, while facilitating compatibility with popular external representations (e.g., Safetensors). The CheckpointLayout abstraction enables loading of diverse formats without bespoke logic, streamlining cross-framework model sharing.
Support for advanced model surgery—runtime checkpoint transformation, partial loading, and manipulation—empowers parameter-efficient fine-tuning (e.g., LoRA), rapid experimentation, and integration with evolving memory optimization strategies. Modular handlers extend to custom data types without explicit dependencies, promoting ecosystem-wide interoperability.
Practical and Theoretical Implications
Practically, distributed checkpointing at extreme scales (10-100K nodes) and flexible storage backend tuning are critical for reducing training recovery times, developer iteration latency, and maximizing expensive accelerator resource utilization. The abstraction of checkpointable objects and the separation of abstract vs. concrete types formalize model representation contracts, facilitating robust migration across topologies and model evolution.
Theoretically, Orbax's layered design—combining scalable serialization, metadata management, and topology-aware loading—lays foundational groundwork for future exploration of model persistence workflows, fine-grained recovery algorithms, and distributed I/O optimization. By modularizing checkpointing primitives, Orbax supports experimentation with alternative storage layout strategies and online checkpoint transformation, essential for adapting to increasingly heterogeneous model architectures, distributed environments, and emergent ML paradigms.
Conclusion
Orbax provides a scalable, performant, and modular solution for distributed checkpointing in JAX, formalizing core use cases and introducing abstractions for flexible, reliable model state management. Benchmarking demonstrates superior throughput and efficiency relative to PyTorch DCP and prior JAX solutions, especially at large scale. The implications extend to ecosystem-wide standardization, usability, and future-proofing for emerging ML research directions. Further development targets extreme scaling, advanced storage tuning, and support for model surgery and checkpoint transformation, positioning Orbax as a key component in the JAX ML infrastructure (2605.23066).