How does a production JAX trainer build its device mesh?

From a handful of integers, not from hand-placed devices. MaxText's create_device_mesh takes one parallelism degree per named axis — ici_fsdp_parallelism, ici_tensor_parallelism, dcn_data_parallelism, and so on, most left at -1 to be "recommended to be auto-sharded" — fills in whatever is left unspecified, and reshapes the flat device list accordingly. Its own docstring states the default plainly: it "creates a device mesh with each slice in its own data parallel group," and single-slice jobs still get two replicas. The result is handed straight to jax.sharding.Mesh along with the axis names, and everything downstream refers to those names, never to a device id.

What is the difference between a logical axis and a physical mesh axis in a JAX trainer?

A logical axis names what a tensor dimension *is* — embed, heads, layers, a batch — independent of any hardware. A physical axis names an actual dimension of the device mesh — data, fsdp, tensor, stage, expert. A rule table (config.logical_axis_rules) maps each logical name onto a tuple of physical axes, and a small real function, remove_size_one_mesh_axis, "removes mesh axes from a PartitionSpec where the axis size is 1" — so a rule can list five physical axes and still shard along only the ones your mesh actually gives more than one device to. Shrink the mesh and rules do not change; they simply resolve to fewer active axes.

Why does a large training job checkpoint asynchronously instead of blocking on the write?

Because the write is enormous and the accelerators are the most expensive thing in the building. A 70B-parameter train state (weights, gradients, and Adam moments) is on the order of 1.1 TB; at a default checkpoint_period of 10,000 steps over a 150,001-step run that is roughly 15 checkpoints and on the order of 16 TB written over the life of one job. MaxText's async_checkpointing defaults to true precisely so training resumes the instant the device-to-host copy lands, while a background thread does the slow part. The save function even measures itself honestly — it calls jax.block_until_ready(state) first "so that our checkpointing metrics measure only checkpointing time, not training time."

What happens to a JAX array after a jitted function donates its buffer?

It stops being usable. A trainer's compiled train step is jitted with donate_argnums = 0 — in the source, the comment reads plainly, "this is the index of the state - we allow the compiler to make use of this memory" — which lets XLA overwrite the incoming state's HBM buffers in place to build the new state, instead of allocating a second full copy of a multi-hundred-gigabyte train state every single step. The price is that the old Python handle to that state is no longer valid: touch it again and JAX raises an error, because the bytes it pointed to may already hold the next step's numbers.

What happens when a host crashes in the middle of a large JAX training run?

The job stops cleanly rather than corrupting silently. A preemption or a runtime error from the accelerator surfaces as a jax.errors.JaxRuntimeError, which the training loop catches and re-raises as a specific exceptions.StopTraining, first waiting on any in-flight checkpoint write so nothing is lost mid-save. From there either the orchestrator relaunches the job — which resumes from checkpoint_manager.latest_step(), possibly onto a completely different number of hosts — or, if elastic training is enabled, elastic_retry restarts it automatically in place. Either path depends on the checkpoint being durable and restorable onto a new mesh shape, which is its own mechanism.

Anatomy of a Production JAX LLM Trainer

One training loop, five subsystems that all have to agree: a device mesh, a table of sharding rules, a jitted train step, an asynchronous checkpointer, and a data pipeline that never hands two devices the same example. We build the mesh from four integers, watch a logical axis resolve onto physical hardware, dial a live mesh and shard a toy tensor, then follow one step through jit, one checkpoint through an async write, and one host through a crash — the whole trainer, end to end.

System design · AI / ML. The source ↗

A free, interactive, animated visual explainer of Anatomy of a Production JAX LLM Trainer — built to be understood, not skimmed.

Questions

How does a production JAX trainer build its device mesh?: From a handful of integers, not from hand-placed devices. MaxText's create_device_mesh takes one parallelism degree per named axis — ici_fsdp_parallelism, ici_tensor_parallelism, dcn_data_parallelism, and so on, most left at -1 to be "recommended to be auto-sharded" — fills in whatever is left unspecified, and reshapes the flat device list accordingly. Its own docstring states the default plainly: it "creates a device mesh with each slice in its own data parallel group," and single-slice jobs still get two replicas. The result is handed straight to jax.sharding.Mesh along with the axis names, and everything downstream refers to those names, never to a device id.
What is the difference between a logical axis and a physical mesh axis in a JAX trainer?: A logical axis names what a tensor dimension *is* — embed, heads, layers, a batch — independent of any hardware. A physical axis names an actual dimension of the device mesh — data, fsdp, tensor, stage, expert. A rule table (config.logical_axis_rules) maps each logical name onto a tuple of physical axes, and a small real function, remove_size_one_mesh_axis, "removes mesh axes from a PartitionSpec where the axis size is 1" — so a rule can list five physical axes and still shard along only the ones your mesh actually gives more than one device to. Shrink the mesh and rules do not change; they simply resolve to fewer active axes.
Why does a large training job checkpoint asynchronously instead of blocking on the write?: Because the write is enormous and the accelerators are the most expensive thing in the building. A 70B-parameter train state (weights, gradients, and Adam moments) is on the order of 1.1 TB; at a default checkpoint_period of 10,000 steps over a 150,001-step run that is roughly 15 checkpoints and on the order of 16 TB written over the life of one job. MaxText's async_checkpointing defaults to true precisely so training resumes the instant the device-to-host copy lands, while a background thread does the slow part. The save function even measures itself honestly — it calls jax.block_until_ready(state) first "so that our checkpointing metrics measure only checkpointing time, not training time."
What happens to a JAX array after a jitted function donates its buffer?: It stops being usable. A trainer's compiled train step is jitted with donate_argnums = 0 — in the source, the comment reads plainly, "this is the index of the state - we allow the compiler to make use of this memory" — which lets XLA overwrite the incoming state's HBM buffers in place to build the new state, instead of allocating a second full copy of a multi-hundred-gigabyte train state every single step. The price is that the old Python handle to that state is no longer valid: touch it again and JAX raises an error, because the bytes it pointed to may already hold the next step's numbers.
What happens when a host crashes in the middle of a large JAX training run?: The job stops cleanly rather than corrupting silently. A preemption or a runtime error from the accelerator surfaces as a jax.errors.JaxRuntimeError, which the training loop catches and re-raises as a specific exceptions.StopTraining, first waiting on any in-flight checkpoint write so nothing is lost mid-save. From there either the orchestrator relaunches the job — which resumes from checkpoint_manager.latest_step(), possibly onto a completely different number of hosts — or, if elastic training is enabled, elastic_retry restarts it automatically in place. Either path depends on the checkpoint being durable and restorable onto a new mesh shape, which is its own mechanism.

Anatomy of a Production JAX LLM Trainer

Questions

Related explainers