JAX 501: systems topics

Everything so far assumed one Python process driving some devices. These pages are about running JAX as part of a larger system: many processes and hosts, long-lived and restartable jobs, artifacts that outlive the process that created them, and the machinery that keeps it all fast and safe.

1. Introduction to multi-controller JAX (aka multi-process/multi-host JAX) — multi-controller JAX: running one process per host, process-spanning meshes and arrays, runtime-level pipeline parallelism with jax.device_put, and building global arrays from per-process data.

2. Distributed data loading — distributed data loading: getting each batch’s shards onto the right hosts, for data-parallel and model-parallel workloads.

3. Fault Tolerant Distributed JAX — fault-tolerant distributed JAX: surviving machine failures with jax.live_devices, barrier semantics, and recovery, with worked training examples.

4. Security considerations — securing JAX’s network services: mTLS for the coordination service and the TPU runtime’s internal services, and which connections remain unprotected.

5. Exporting and serializing staged-out computations — exporting and serializing staged-out computations for later or cross-platform execution, and Shape polymorphism — exporting with symbolic shapes so one artifact serves many input sizes.

6. Persistent compilation cache — the persistent compilation cache: skipping recompilation across process restarts and sharing compiled artifacts across nodes.

7. Transfer guard — logging or disallowing unintended host-device transfers.

Smaller notes

A few systems topics are small enough to cover right here.

Concurrency and threads. JAX has limited support for Python concurrency: it’s fine to call JAX APIs like jax.jit() or jax.grad() from multiple Python threads, but you must not use threads to manipulate JAX trace values inside a traced function — the likely outcome is a mysterious error. In multi-controller JAX (Introduction to multi-controller JAX (aka multi-process/multi-host JAX)) threading has a further hazard: every process must enqueue the same operations in the same order on a given device, and threads can schedule work in different orders in different processes, causing non-deterministic crashes. The jax.thread_guard() context manager helps detect this: once set, an error is raised if a JAX operation is issued from a thread other than the one where the guard was set.

Multi-process coordination helpers. The jax.experimental.multihost_utils module collects small utilities for multi-controller programs: sync_global_devices (a named cross-process barrier), broadcast_one_to_all (make process 0’s value everyone’s value), process_allgather (gather a value from every process), assert_equal (check that a value agrees across processes), and host_local_array_to_global_array / global_array_to_host_local_array (convert between per-host and global views of arrays).

Compatibility policies for long-lived artifacts. Serialized artifacts outlive the process that made them, so here is what’s promised: exported modules have explicit compatibility windows (see the compatibility guarantees in Exporting and serializing staged-out computations); persistent compilation cache entries make no cross-version promises, but their keys include the jaxlib version, so version changes cause cache misses rather than misbehavior; and JAX’s general API stability rules are described in the API compatibility policy.