Orbax: Distributed Checkpointing with JAX
Not on the map. Why?
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.
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
7 of 17 references.
References · 7
- 2021LoRA: Low-Rank Adaptation of Large Language Modelscited by 14,380
- 2019NeMo: a toolkit for building AI applications using Neural Modulescited by 235
- 2022Scaling Up Models and Data with textttt5x and textttseqiocited by 187
- 2024The rising costs of training frontier AI modelscited by 95
- 2024Revisiting Reliability in Large-Scale Machine Learning Research Clusterscited by 28
- 2024ByteCheckpoint: A Unified Checkpointing System for Large Foundation Model Developmentcited by 24
- 2025Adaptive Fault Tolerance Mechanisms of Large Language Models in Cloud Computing Environmentscited by 12
Built on by
Not cited yet.