The JAX Stack, From the Inside
A bottom-up path through the modern JAX training/serving stack, for reading real open-source code rather than interview prep — jax, flax, orbax, xla, qwix, maxtext — from tracing to sharding to checkpointing to quantization to production.
Layer 1 — The core mechanic
One trick — trace a Python function into a typed jaxpr — becomes jit, grad, and vmap alike, and the reason a model is just a pytree of arrays.
- Tracing → jaxpr: the One Trick Behind Every JAX Transform — Run a Python function once with shape-and-dtype stand-ins and you get a typed jaxpr — the record every other JAX transform reads. Start here; everything below assumes it.
- jit(grad(vmap(f))): Why Transform Order Changes the Answer — jit, grad, and vmap all look like decorators you can stack in any order. They are not — vmap(grad(f)) computes per-example gradients, and grad(vmap(f)) does not even type-check until you bolt on a reduction.
- Params Are Data: the Functional Model Behind Flax — A Flax model never holds its own weights — init() hands back a plain pytree of arrays, apply() takes it back in. That discipline is what lets jax.grad, sharding, and checkpointing treat a whole model like any other array.
Layer 2 — Splitting across devices
From a PartitionSpec on one array, to the compiler passes that insert the collectives, to the kernel that makes attention fast once the bytes land on one chip.
- Sharding in JAX — Name your devices into a mesh, annotate an array with a PartitionSpec, and the compiler places the bytes and inserts the communication — the whole story from one device to a full mesh.
- The SPMD Partitioner: One Program Across a Device Mesh — One program, annotated once, becomes N identical per-device programs — the partitioner that inserts every all-gather, reduce-scatter, and halo exchange your sharding implies.
- Inside XLA: ~200 Passes and the Fusion Decision — Between your traced program and the GPU kernels sit roughly 200 compiler passes. Fusion, the one that decides which ops share a kernel, runs a stopwatch on every merge before it commits.
- Anatomy of a FlashAttention Kernel — Attention is slow not because of the math but because it writes an N×N score matrix to memory and reads it back. FlashAttention tiles the keys and never writes that matrix at all.
Layer 3 — State at scale
The full checkpoint story, save to load: who writes which bytes, how thousands of per-host files become one atomic manifest, and how a multihost cluster reads back exactly its own shards.
- The Checkpoint Lifecycle: What Happens Between save() and Durable — Follow one save() from the training step that triggers it to the moment it's durable and restorable on a different mesh — async, host-coordinated, atomic.
- Zero-RPC Sharding: How 1,000 Hosts Agree Who Writes Which Bytes — Every data-parallel replica holds the same shard at checkpoint time, so who actually writes it? A deterministic function every host runs alone, with zero coordination RPCs.
- The B-Tree Merge: Thousands of Checkpoint Files into One Atomic Manifest — Thousands of per-host checkpoint files become one atomic manifest through a single TensorStore transaction — all of it commits, or none of it does.
- What "Atomic" Means on a Filesystem vs. an Object Store — Atomic means something different on a POSIX filesystem than it does on an object store with no rename at all — the real code that guarantees it either way.
- Loading a Safetensors Checkpoint on a Multihost Cluster, Line by Line — A safetensors file is one flat byte blob per tensor, but a big model has to land as sharded arrays across many hosts. The loader that maps each host to exactly its own byte range.
Layer 4 — Making it cheap
Quantization without touching model code — the interception trick, then the algorithm it hides, read line by line.
- How Qwix Quantizes Any Flax Model Without Touching Its Code — Quantize a model without touching its code: patch jax.lax.dot_general and jnp.einsum for the duration of one forward pass, and every matmul routes through int8 underneath.
- GPTQ, Line by Line: Hessian-Based Weight Quantization — Round one weight and the weights you haven't rounded yet absorb the error — the real 220-line JAX implementation, Hessian, dampened Cholesky, and all.
Layer 5 — Where it all meets
The production capstones — a real training loop and a real RL cluster, built from every layer above.
- Anatomy of a Production JAX LLM Trainer — One training loop, five subsystems that all have to agree: a device mesh, a sharding-rules table, a jitted train step, an async checkpointer, a data pipeline. Everything above, assembled.
- The RL Cluster: Five Roles, One Mesh Dial — Five roles sharing one accelerator fleet, and one line of config that decides whether they share memory or ship weights across the wire.