Multi-controller JAX

Open a JAX training script built for a thousand-host pod and you might look for a coordinator: some process handing out work, some scheduler routing batches. Stock JAX has none. It starts one ordinary Python process per host and hands every one of them the exact same script; jax.distributed.initialize() is the single call that lets those processes discover each other, and after it returns nothing else about how you write the program changes.

The architecture has a name in the community: McJAX, for multi-controller JAX, because every host is its own controller running an identical copy of the program rather than one controller dispatching work to workers. Chapter 11's IFRT array is what lets a single logical array span every one of those processes in your Python code even though the processes never share memory; this chapter sits one layer above that, at the level where the processes themselves come into existence and find each other.

Each host runs its copy of the program forward on its own, at its own pace, with no message passing except at one specific kind of instruction: a collective. An all-reduce or an all-gather is the sole channel between hosts, and outside of those calls, a host genuinely cannot tell whether the others are ahead of it, behind it, or have crashed.

Every host runs the same program and believes it is alone.

That isolation has a cost on the other side of the same coin: because every host runs the same program in lockstep, one host dying kills the whole gang. There is no spare controller to reroute around a dead worker, because there is no controller at all, only peers running the same script. Elasticity in the face of a failed host is not something this architecture offers; it is the problem chapter 13's design exists to solve.

Before any of this can dispatch anything, JAX has to find a backend to talk to, and xla_bridge is the module that does the finding. CPU support is built in; everything else, TPU, GPU, and any third-party accelerator, arrives as a PJRT plugin, registered through Python entry points that ship inside packages like jax[tpu] or jax[cuda]. Install one of those extras and xla_bridge discovers it automatically at import time.

When more than one backend is available and you want to choose explicitly rather than let discovery decide, the JAX_PLATFORMS environment variable does that job. Set it before your script runs and the backend selection is fixed before a single line of your program executes.

The path a single jax.jit call takes, once a backend is chosen, follows the chapters before this one almost exactly: trace to a jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr →, lower that jaxpr to StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo →, hand the StableHLO to the backend's PJRT or IFRT client to compile, and keep the resulting executable cached in memory so the next call with the same signature skips straight to execution. Nothing here is new machinery; it is the same pipeline chapters 1 through 10 built, now running once per host, every host compiling the identical program in parallel.

What is new is what happens when the process restarts. An in-memory cache dies with the process that held it, and recompiling a large program through chapter 4's two hundred passes is not free. A persistent compilation cache, which you opt into, serializes compiled executables to disk, keyed on a fingerprint of the program and the compiler version, so a second run of the same script, even in a fresh process, can skip compilation entirely and load the executable straight from disk. It is the production answer to a cost this track measured directly two stages back.