Orbax: Distributed Checkpointing with JAX
Distributed, Parallel, and Cluster Computing
2026-05-28 v2 Machine Learning
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.5 for saving and 2 for loading. The library is available at https://github.com/google/orbax.
Keywords
Cite
@article{arxiv.2605.23066,
title = {Orbax: Distributed Checkpointing with JAX},
author = {Colin Gaffney and Shutong Li and Daniel Ng and Anastasia Petrushkina and Niket Kumar and Adam Cogdell and Mridul Sahu and Yaning Liang and Nikhil Bansal and Justin Pan and Angel Mau and Abhishek Agrawal and Marco Berlot and Ruoxin Sang and Kiranbir Sodhia and Rakesh Iyer},
journal= {arXiv preprint arXiv:2605.23066},
year = {2026}
}
Comments
18 pages, 5 tables, 6 figures