Orbax: JAX-Native Checkpointing

• Orbax is a modular, JAX-native checkpointing library that abstracts distributed accelerator and storage complexities to enable efficient asynchronous saving, partial restoration, and resharding.
• Benchmarks indicate that Orbax can achieve up to 3.5× faster saving and 2× faster loading compared to PyTorch competitors, highlighting its performance advantage in distributed settings.
• The library supports multiple API patterns—training-loop checkpointing and file-path-oriented operations—ensuring flexible integration in diverse ML workflows and facilitating model recovery and post-training manipulation.

Searching arXiv for papers directly relevant to Orbax and its documented uses. I’ll look up recent arXiv entries on Orbax and JAX checkpointing to ground the article in cited literature. Orbax is a modular, JAX-native checkpointing library designed to provide a standardized checkpointing solution for distributed machine learning systems built with JAX. It is motivated by a structural gap in the JAX ecosystem: JAX’s modular design philosophy leaves checkpointing fragmented or absent, even though checkpointing is operationally central for recoverability, preemption tolerance, topology changes, and post-training manipulation across the ML model lifecycle. Orbax abstracts distributed accelerator and storage complexity while exposing higher-level and lower-level interfaces for checkpoint orchestration, partial restoration, and format-aware loading. In reported benchmarks, it exceeds comparable PyTorch competitors by up to $3.5\times$ for saving and $2\times$ for loading (Gaffney et al., 21 May 2026).

1. Problem setting and design motivation

Orbax addresses checkpointing as a distributed systems problem rather than as mere serialization. The underlying requirements include recovery from failures and preemptions, minimizing accelerator idle time during saves, restoring onto different accelerator topologies via resharding, and supporting checkpoints beyond training-time restartability, including pre-training, post-training, evaluation, inference, and experimentation such as LoRA-style adaptation (Gaffney et al., 21 May 2026).

The need for such a library arises from JAX’s intentionally minimal core. Prior JAX-related approaches were inconsistent across codebases, sometimes limited to single-host settings, sometimes dependent on manual low-level array and storage manipulation, and sometimes tied to framework-specific stacks such as T5X or Pax. Orbax is positioned as a common abstraction layer that preserves JAX’s modularity while supplying a shared checkpointing substrate for distributed execution, storage coordination, and object-structure-aware restoration (Gaffney et al., 21 May 2026).

A central implication of this framing is that Orbax is not merely a file format or a convenience wrapper. Its scope includes accelerator-to-host transfer, asynchronous storage I/O, metadata coordination, topology-aware restoration, and flexible manipulation of nested model state. This suggests that Orbax is best understood as infrastructure for the entire checkpoint lifecycle, not only for periodic save calls.

2. API model and object abstractions

Orbax organizes its API around two principal usage patterns. The first is sequence-of-steps checkpointing, intended for training loops that save every $n$ steps and require orchestration features such as step tracking, retention, garbage collection, and interval selection. The primary API for this mode is training.Checkpointer. The second is single-checkpoint, file-path-oriented checkpointing, intended for direct control, debugging, and cases where step-based orchestration assumptions are undesirable (Gaffney et al., 21 May 2026).

A defining abstraction is the checkpointable, a discrete, independently serializable subcomponent of training state. Examples explicitly given include model parameters, optimizer state, and dataset iterator state. The abstraction is governed by two principles: checkpointables are separable, so one component can be loaded without requiring others, and distinct types may use different serialization logic and storage representations. This supports workflows such as loading parameters without optimizer state for evaluation or inference, or omitting dataset iterators when they are irrelevant (Gaffney et al., 21 May 2026).

Orbax exposes this design through CheckpointableHandler[T, AbstractT], which provides methods such as save, load, and metadata, and through StatefulCheckpointable, which allows objects to implement save and load directly. At the leaf level, the implementation includes LeafHandler[L, AbstractL] and concrete handlers such as ArrayHandler, NumpyHandler, ScalarHandler, and StringHandler. Format and path control are mediated by abstractions including Context and CheckpointLayout (Gaffney et al., 21 May 2026).

Orbax also formalizes the distinction between concrete and abstract checkpointables. The abstract form is a lightweight representation of properties such as shapes, dtypes, and shardings. This abstract state defines the contract between expected object structure and on-disk bytes, and it is used for metadata description, validation, reshaping, casting, and resharding at load time. A plausible implication is that abstract checkpointables are the mechanism that lets Orbax separate logical state compatibility from specific device layouts.

3. Save, load, and distributed execution semantics

Orbax save operations are split into synchronous and asynchronous phases. Before returning control to training, the synchronous phase performs validations, existence checks, and device-to-host copying of model weights. The copy from accelerator to host is identified as the most critical blocking task, both because weights must not be mutated concurrently with checkpoint creation and because some hardware cannot write directly from accelerator memory to storage. Once the host copy is complete, the asynchronous phase can proceed in a background thread, allowing training to continue while metadata coordination and distributed writes run concurrently (Gaffney et al., 21 May 2026).

Checkpoint creation is atomic at the directory level. Orbax writes into a temporary directory and finalizes only after all operations complete. Where the filesystem supports atomic directory renames, the temporary directory is renamed into the final path; otherwise an indicator mechanism is used. This prevents partially written checkpoints from being interpreted as complete checkpoints (Gaffney et al., 21 May 2026).

Metadata is layered. A global metadata file records tree structure, key names, key types such as dict, list, and tuple, and empty nodes. Per-process metadata records array shapes, shard shapes, and storage chunk shapes. A leader process writes global metadata and validates consistency across workers during finalization. For the actual data path, each host writes to its own subdirectory, after which a primary host merges per-host shard metadata into a global index. That merged index references original data rather than copying it, so finalization is metadata-heavy rather than data-heavy (Gaffney et al., 21 May 2026).

Orbax is designed for both multi-controller JAX and single-controller Pathways. In multi-controller JAX, checkpointing requires leader election, barriers, and careful coordination of directory creation, metadata writing, merging, and finalization. In single-controller Pathways, JAX’s Colocated Python API allows a central controller to handle metadata and orchestration while worker hosts perform direct device, host, and storage I/O. For host-resident objects such as NumPy arrays, strings, and scalars, serialization logic is similar across modes; for on-device arrays, Orbax uses different leaf handling to schedule transfers directly on remote worker hosts (Gaffney et al., 21 May 2026).

Loading is driven by abstract state when available. With an abstract state, Orbax validates structure against checkpoint metadata and applies specified properties such as casting, sharding, and reshaping or other transformations. Without an abstract state, it falls back to checkpoint metadata and validates against current topology, but does not infer a sharding when the device set has changed. Partial loading is supported: if the abstract state is a subset of the checkpoint, only that subset is loaded; if it is a superset, missing values are filled with ... placeholders. In distributed loading, each host reads the shards it needs, moves data from storage to host and then host to device, and constructs global arrays without device-to-device movement (Gaffney et al., 21 May 2026).

4. Storage formats, resharding, and measured performance

Orbax uses TensorStore-backed formats. With TensorStore’s zarr driver, each weight is stored as a separate subdirectory and each shard as a distinct file. Orbax recommends ocdbt layered on zarr as the default for large-scale models. OCDBT, described as an optionally-cooperative distributed B-tree, coalesces multiple shards into larger files, stores keys in a B+ tree, and places array metadata in a .zarray key. The paper states that OCDBT can group data into files near a target size of about $2\ \mathrm{GiB}$, which is advantageous at scale, although it can add overhead at small scales (Gaffney et al., 21 May 2026).

A major Orbax feature is load-time resharding. TensorStore stores data in chunked form, and Orbax refers to smaller read units as subchunks. If restore sharding differs from save sharding, read requests may not align with write chunks; suitably chosen subchunk boundaries at save time can greatly reduce over-reading. The paper illustrates this with a tensor of shape (x, y) saved across 64 devices with (16, 4) partitioning and restored on 16 devices using (16, 1) or (64, 1) partitioning. When chunk boundaries align, reads are efficient; when they do not, extra data is loaded and discarded (Gaffney et al., 21 May 2026).

For resharding benchmarks, the reported mean load times show large gains from subchunking. For the 70B model, the 4-8-4 → 1-128-1 case is 13.41 ± 1.44 s with subchunking and 45.23 ± 2.69 s without, a speedup of 3.37×; the 4-8-4 → 1-32-1 case is 14.40 ± 0.51 s versus 34.40 ± 2.84 s, a speedup of 2.39×. For the 8B model, the 8-2-1 → 1-16-1 case is 3.59 ± 0.17 s versus 20.40 ± 2.56 s, a speedup of 5.68×, while 8-2-1 → 1-4-1 is 10.88 ± 0.49 s versus 21.07 ± 3.12 s, a speedup of 1.94× (Gaffney et al., 21 May 2026).

The headline comparison against PyTorch Distributed Checkpoint was conducted on NVIDIA A100 (40 GB) GPUs with single-region Google Cloud Storage, using Llama 3.1 8B, 70B, and 405B models on 4, 16, and 64 devices respectively, with 20 isolated save and 20 isolated load operations per model (Gaffney et al., 21 May 2026).

Model	Save P50: PyTorch vs Orbax	Load P50: PyTorch vs Orbax
8B	14.4s vs 33.9s	45.1s vs 22.6s
70B	211.0s vs 104.6s	109.7s vs 75.3s
405B	595.0s vs 173.0s	169.5s vs 117.2s

These numbers correspond to median save and load times, with background-save improvements reported as 0.6× for 8B, 2.0× for 70B, and 3.4× for 405B, and load improvements reported as 2.0×, 1.5×, and 1.4× respectively. The paper also notes that blocking save numbers are mostly attributable to device-to-host copying via JAX/XLA rather than uniquely to Orbax itself (Gaffney et al., 21 May 2026).

In multi-slice CPU-only scalability experiments on Llama 3.1 70B, Orbax introduces Replica-Parallel saving, in which all slices participate and each array is divided into $n$ segments so each slice writes its $1/n$ fraction. Reported mean save times improve from 49.20 ± 12.28 to 29.40 ± 3.00 at 2 slices (1.67×), from 57.46 ± 20.19 to 24.47 ± 6.97 at 8 slices (2.35×), and from 77.53 ± 18.88 to 32.51 ± 13.88 at 16 slices (2.39×). For replicated loading, the paper evaluates a broadcast-based method in which one slice reads from storage and then broadcasts via JAX collectives. This is slower at 2 slices, 35.31 ± 1.69 versus 27.88 ± 2.10 (0.79×), but improves substantially at larger counts, reaching 4.06× at 16 slices and 4.93× at 32 slices (Gaffney et al., 21 May 2026).

5. Documented roles in published JAX workflows

In the DAC-JAX reimplementation of the Descript Audio Codec, Orbax appears as one of the core libraries in “Google’s JAX ecosystem of Flax, Optax, Orbax, AUX, and CLU.” The paper establishes Orbax as part of the infrastructure layer of the project, but does not provide Orbax-specific implementation details: there is no dedicated Orbax section, no discussion of checkpointing APIs, no mention of Orbax-managed model serialization, no example of saving or restoring training state, and no explanation of Orbax integration in the training loop. The surrounding JAX workflow does include a training and fine-tuning script, device parallelism via JAX’s pmap, on-device data augmentation, reimplementation of the original DAC loss functions and evaluation metrics, CLU-based metrics, and TensorBoard-based inspection of metrics and reconstructed audio. The codebase also enables reuse of model weights from the original PyTorch DAC, and the two implementations produce equivalent token sequences and decoded audio for the same input. This suggests a JAX-native training stack into which Orbax would plausibly fit, but the paper does not confirm its operational role beyond inclusion in the software stack (Braun, 2024).

In contrast, the Gemma 4 31B TPU study assigns Orbax an explicit operational function. There, Orbax is the checkpointing backend for JAX/Tunix training and is used to persist training state during TPU fine-tuning, survive preemption by writing checkpoints directly to GCS, save the LoRA-adapted JAX/Tunix model state in a JAX-native tree format, provide the source material for a custom post-training merge script, and enable conversion from Orbax checkpoints into a merged HuggingFace-style safetensors model for inference. The paper gives a concrete checkpoint regime: Orbax via orbax-checkpoint, storage on Google Cloud Storage, checkpoints every 100 optimizer steps, and retention of the 3 most recent checkpoints (Kishnani et al., 25 May 2026).

The Gemma workflow also illustrates Orbax’s interoperability limits and the corresponding conversion logic. Orbax checkpoints are not directly usable for inference because they store JAX trees rather than safetensors, Tunix JAX weight names and layouts differ from HuggingFace, and one JAX module may map to multiple HuggingFace tensors, particularly because of the fused kv_einsum parameterization. The authors therefore implement a custom converter, orbax_to_peft.py, that restores Orbax checkpoints, reconstructs the LoRA-augmented JAX model, traverses nnx.LoRAParam tensors, computes the LoRA update

$$W_{\text{merged}} = W_{\text{base}} + \frac{\alpha}{r}\cdot AB,$$

and writes a single merged model.safetensors file (Kishnani et al., 25 May 2026).

The same paper gives module-specific layout conversions. For q_einsum, JAX uses $A:(\mathrm{in}, r)$ and $B:(r, \mathrm{heads}, \mathrm{head\_dim})$, with conversion via $(AB)^T$ into HuggingFace (out, in) layout. For kv_einsum, $2\times$0 is fused for K and V and must be split along dimension 1 before separate $2\times$1 deltas are computed. For attn_vec_einsum, $2\times$2 is flattened to (heads * hd, r) before computing $2\times$3. For gate_proj, up_proj, and down_proj, the conversion again uses $2\times$4 to match HuggingFace conventions. This case study presents Orbax as a robust TPU-native persistence layer, but not as the final deployment format (Kishnani et al., 25 May 2026).

6. Significance, boundaries, and recurrent misconceptions

Orbax’s broader significance lies in standardizing checkpointing for JAX without collapsing JAX’s modular design philosophy into a monolithic framework. Its documented capabilities include asynchronous saving, topology-aware restoration, resharding, partial loading, support for both multi-controller and single-controller execution, multiple storage layouts, and compatibility with external formats through CheckpointLayout, including a documented safetensors loading path via: $2\times$5 This positions Orbax as a checkpointing layer that spans training recovery, experimentation, model editing, and deployment-adjacent transformations (Gaffney et al., 21 May 2026).

At the same time, the literature makes several boundaries clear. First, inclusion of Orbax in a JAX software stack does not by itself document Orbax-specific behavior; DAC-JAX is an example where Orbax is listed as part of the ecosystem but not otherwise analyzed (Braun, 2024). Second, Orbax is not synonymous with a serving format. In the Gemma 4 TPU pipeline, Orbax is an intermediate persistence layer between TPU fine-tuning and safetensors export, and a custom postprocessing stage is required to bridge JAX pytrees and HuggingFace-compatible artifacts (Kishnani et al., 25 May 2026). Third, Orbax performance is strongly conditioned by layout choices and scale: OCDBT can add overhead at small scales, broadcast-based replicated loading can be counterproductive at 2 slices, and blocking save time is dominated in part by JAX/XLA device-to-host copying rather than by checkpoint library logic alone (Gaffney et al., 21 May 2026).

A plausible implication is that Orbax’s principal contribution is infrastructural regularization: it supplies a common, JAX-native checkpoint contract for distributed state while leaving format translation and deployment interoperability to higher-level workflows when model ecosystems diverge.