# Orbax: Distributed Checkpointing with JAX

arXiv: https://arxiv.org/abs/2605.23066 · HTML: https://mlsanity.com/paper/2605.23066/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 7 of its 17 references that are in the map.
- First posted on arXiv 21 May 2026.

## Builds on (most cited first, 7 of 7)

- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [NeMo: a toolkit for building AI applications using Neural Modules](https://mlsanity.com/paper/1909.09577/) (2019)
- [Scaling Up Models and Data with textttt5x and textttseqio](https://mlsanity.com/paper/2203.17189/) (2022)
- [The rising costs of training frontier AI models](https://mlsanity.com/paper/2405.21015/) (2024)
- [Revisiting Reliability in Large-Scale Machine Learning Research Clusters](https://mlsanity.com/paper/2410.21680/) (2024)
- [ByteCheckpoint: A Unified Checkpointing System for Large Foundation Model Development](https://mlsanity.com/paper/2407.20143/) (2024)
- [Adaptive Fault Tolerance Mechanisms of Large Language Models in Cloud Computing Environments](https://mlsanity.com/paper/2503.12228/) (2025)

## Abstract (arXiv)

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.