Orbax: Distributed Checkpointing with JAX

• URL: http://arxiv.org/abs/2605.23066v2
• Date: Tue, 26 May 2026 19:23:57 GMT
• Title: Orbax: Distributed Checkpointing with JAX
• Abstract summary: Orbax is a JAX-native checkpointing library that abstracts the complexities of distributed accelerator systems.<n>We demonstrate performance exceeding comparable PyTorch competitors by up to 3.5$times$ for saving and 2$times$ for loading.
• Score: 2.2139026464835574
• License: http://creativecommons.org/licenses/by/4.0/
• 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$\times$ for saving and 2$\times$ for loading. The library is available at https://github.com/google/orbax.

Related papers

• MASEval: Extending Multi-Agent Evaluation from Models to Systems [77.09381093359598]
  MASEval is a framework-agnostic library that treats the entire system as the unit of analysis.<n>It allows researchers to explore all components of agentic systems, opening new avenues for principled system design.
  arXiv  Detail & Related papers  (2026-03-09T18:46:17Z)
• MPX: Mixed Precision Training for JAX [56.49234478198188]
  Mixed-precision training has emerged as an indispensable tool for enhancing the efficiency of neural network training.<n>We propose MPX, a mixed-precision training toolbox for JAX that simplifies and accelerates the training of large-scale neural networks.<n>MPX seamlessly integrates with popular toolboxes such as Equinox and Flax, allowing users to convert full-precision pipelines to mixed-precision versions.
  arXiv  Detail & Related papers  (2025-07-04T05:47:04Z)
• SwingArena: Competitive Programming Arena for Long-context GitHub Issue Solving [90.32201622392137]
  We present SwingArena, a competitive evaluation framework for Large Language Models (LLMs)<n>Unlike traditional static benchmarks, SwingArena models the collaborative process of software by pairing LLMs as iterations, who generate patches, and reviewers, who create test cases and verify the patches through continuous integration (CI) pipelines.
  arXiv  Detail & Related papers  (2025-05-29T18:28:02Z)
• jinns: a JAX Library for Physics-Informed Neural Networks [0.7366405857677227]
  jinns is an open-source Python library for physics-informed neural networks.<n> Rooted in the JAX ecosystem, it provides a versatile framework for efficiently prototyping real-problems.
  arXiv  Detail & Related papers  (2024-12-18T18:21:41Z)
• JAXbind: Bind any function to JAX [0.0]
  JAXbind provides an easy-to-use Python interface for defining custom, so-called JAX primitives. JAXbind allows a user to interface the JAX function engine with custom derivatives and rules, enabling all JAX transformations for the custom primitive.
  arXiv  Detail & Related papers  (2024-03-13T16:50:04Z)
• DrJAX: Scalable and Differentiable MapReduce Primitives in JAX [9.676195490442367]
  DrJAX is a library designed to support large-scale distributed and parallel machine learning algorithms. DrJAX embeds building blocks for MapReduce computations as primitives in JAX. DrJAX computations can be translated directly to XLA HLO, enabling flexible integration with a wide array of ML training platforms.
  arXiv  Detail & Related papers  (2024-03-11T19:51:01Z)
• BlackJAX: Composable Bayesian inference in JAX [8.834500692867671]
  BlackJAX is a library implementing sampling and variational inference algorithms. It is written in Python, using JAX to compile and run NumpPy-like samplers and variational methods on CPUs, GPUs, and TPUs.
  arXiv  Detail & Related papers  (2024-02-16T16:21:02Z)
• XLB: A differentiable massively parallel lattice Boltzmann library in Python [0.0]
  We introduce XLB library, a Python-based differentiable LBM library based on the JAX platform. XLB's differentiability and data structure is compatible with the extensive JAX-based machine learning ecosystem. XLB has been successfully scaled to handle simulations with billions of cells, achieving giga-scale lattice updates per second.
  arXiv  Detail & Related papers  (2023-11-27T18:50:37Z)
• JaxMARL: Multi-Agent RL Environments and Algorithms in JAX [105.343918678781]
  We present JaxMARL, the first open-source, Python-based library that combines GPU-enabled efficiency with support for a large number of commonly used MARL environments. Our experiments show that, in terms of wall clock time, our JAX-based training pipeline is around 14 times faster than existing approaches. We also introduce and benchmark SMAX, a JAX-based approximate reimplementation of the popular StarCraft Multi-Agent Challenge.
  arXiv  Detail & Related papers  (2023-11-16T18:58:43Z)
• JaxPruner: A concise library for sparsity research [46.153423603424]
  JaxPruner is an open-source library for sparse neural network research. It implements popular pruning and sparse training algorithms with minimal memory and latency overhead.
  arXiv  Detail & Related papers  (2023-04-27T10:45:30Z)
• SequeL: A Continual Learning Library in PyTorch and JAX [50.33956216274694]
  SequeL is a library for Continual Learning that supports both PyTorch and JAX frameworks. It provides a unified interface for a wide range of Continual Learning algorithms, including regularization-based approaches, replay-based approaches, and hybrid approaches. We release SequeL as an open-source library, enabling researchers and developers to easily experiment and extend the library for their own purposes.
  arXiv  Detail & Related papers  (2023-04-21T10:00:22Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.