Orbax: Distributed Checkpointing with JAX

arXiv:2605.23066v1 Announce Type: cross Abstract: In a landscape of high-performance distributed ML systems, JAX has emerged as a framework of choice. However, JAX's modular design philosophy leaves it without a standardized checkpointing solution. In this paper, we introduce Orbax, a modular, JAX-native checkpointing library that abstracts the complexities of distributed accelerator systems while also providing flexibility for user-friendly checkpoint manipulations throughout the ML model lifecycle. We demonstrate performance exceeding comparable PyTorch competitors by up to 3.5imes for saving and 2imes for loading. The library is available at https://github.com/google/orbax.

Source: arXiv cs.LG | 2026-05-25