Add ahead-of-time (XAOT) compilation for MaxTextTrainingEngine - #5112
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces ahead-of-time (AOT) compilation and lowering capabilities for the MaxTextTrainingEngine without requiring physical hardware or weight materialization. It adds a materialize_weights flag to the engine, implements abstract model and state creation using jax.ShapeDtypeStruct, and introduces a new train_compile.py script alongside comprehensive parity tests. The review feedback identifies potential AttributeError crashes in the newly added _to_aval and _rehome_aval helper functions when processing non-array leaves (such as Python scalars or step counts) in the train or optimizer states, and suggests adding guards to return these leaves as-is.
| def _rehome_aval(aval: Any, mesh: jax.sharding.Mesh) -> Any: | ||
| """Returns `aval` with its sharding spec re-expressed on `mesh`. | ||
|
|
||
| Propagation hands back `NamedSharding`s on whichever mesh was active during the trace -- | ||
| an `AbstractMesh`, or the explicit view above. `_mesh_sharding` compares meshes by | ||
| equality to decide whether a leaf belongs to this engine's, so a spec that is right but | ||
| homed elsewhere would be silently discarded and replaced by a replicated one. | ||
| """ | ||
| spec = getattr(getattr(aval, "sharding", None), "spec", None) | ||
| target = jax.sharding.NamedSharding(mesh, spec) if spec is not None else None | ||
| return jax.ShapeDtypeStruct(aval.shape, aval.dtype, sharding=target) |
There was a problem hiding this comment.
Similar to _to_aval, the _rehome_aval helper assumes that the input aval is always an array-like object with shape and dtype attributes. Since _rehome_aval is mapped over the entire model_pure and state_pure trees in _create_abstract_state, any non-array leaves (such as Python scalars) will cause an AttributeError when trying to construct the new ShapeDtypeStruct.\n\nAdding a guard to return the leaf as-is if it lacks shape or dtype attributes will prevent potential crashes during abstract state creation.
def _rehome_aval(aval: Any, mesh: jax.sharding.Mesh) -> Any:\n """Returns `aval` with its sharding spec re-expressed on `mesh`.\n\n Propagation hands back `NamedSharding`s on whichever mesh was active during the trace --\n an `AbstractMesh`, or the explicit view above. `_mesh_sharding` compares meshes by\n equality to decide whether a leaf belongs to this engine's, so a spec that is right but\n homed elsewhere would be silently discarded and replaced by a replicated one.\n """\n if aval is None:\n return None\n if not hasattr(aval, "shape") or not hasattr(aval, "dtype"):\n return aval\n spec = getattr(getattr(aval, "sharding", None), "spec", None)\n target = jax.sharding.NamedSharding(mesh, spec) if spec is not None else None\n return jax.ShapeDtypeStruct(aval.shape, aval.dtype, sharding=target)There was a problem hiding this comment.
Taken, in ecdcf56 — but as consistency rather than as a crash fix, and it is worth saying which, since the two helpers are not equally exposed.
I instrumented _rehome_aval and recorded the type of every leaf it is handed, across adamw, adam_pax, sgd and Zero-1:
adamw/dp: {'ShapeDtypeStruct': 159}
zero1: {'ShapeDtypeStruct': 159}
sgd: {'ShapeDtypeStruct': 116}
adam_pax: {'ShapeDtypeStruct': 158}
No non-array leaf, and one of the two call sites cannot produce one by construction: state_pure is the output of jax.eval_shape, which returns ShapeDtypeStructs and nothing else. The other, model_pure, comes from create_nnx_abstract_model, which is itself built out of avals. So the AttributeError is not reachable on any path I can construct today — optax's count and the optimizer's step, the scalars this would most plausibly be about, are jnp.zeros([], int32) and so are arrays.
Still worth adding. _to_aval and _place_state_on_mesh both guard, and _rehome_aval being the one that does not is the kind of asymmetry that reads as an oversight later. Returning the leaf untouched is also the right answer semantically: something with no shape has no sharding to re-home.
9989cd3 to
ecb35e6
Compare
5c937f1 to
fcec4fc
Compare
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
4bc4e9c to
ecdcf56
Compare
fcec4fc to
923c670
Compare
a003635 to
5324209
Compare
1c69d6e to
50d512a
Compare
| from maxtext.utils import maxtext_utils | ||
| from maxtext.utils import model_creation_utils | ||
|
|
||
| KERNEL_NAMES = ("fwd_bwd", "fwd_bwd_accum", "update") |
There was a problem hiding this comment.
can we give users the choice of which to compile, maybe just one at a time? compilation can be slow, a user might only be interested in one of these, not all of them
There was a problem hiding this comment.
I am hesitant on adding more additional maxtext flags for this feature... Let's do that in followup PRs.
There was a problem hiding this comment.
Unless there is a way to avoid adding maxtext flags for this.
| @@ -1359,6 +1421,50 @@ | |||
| dynamic_batch, static_batch = _split_static_and_dynamic(self._prepare_batch(dummy_data)) | |||
| self._compile_for_batch(dynamic_batch, static_batch) | |||
There was a problem hiding this comment.
please add actual pre-compilation here.
| in_shardings=eval_in_shardings, | ||
| out_shardings=eval_out_shardings, | ||
| ) | ||
| self._compiled_eval_signature = _batch_signature(dynamic_batch, static_batch) |
There was a problem hiding this comment.
added compile eval
997687f to
71b53d5
Compare
`trainers/pre_train/train_compile.py` lets a pre-training configuration be
costed and memory-checked against a target topology from a host that does not
own it. The engine had no equivalent, so every iteration on a sharding or a
batch size meant booking the hardware.
Adds `training_engine/maxtext_engine_compile.py`, which does the same for the
engine's three kernels (first forward/backward, accumulating forward/backward,
update) and reports each one's cost and memory. Nothing is materialized:
weights and optimizer moments are `jax.ShapeDtypeStruct`s and the mesh is a
topology description.
Almost all of that lives in the new module, as `AbstractMaxTextEngine`, a
`MaxTextTrainingEngine` subclass. Its moments cannot come from `nnx.Optimizer`,
which allocates them with `zeros_like`, so it builds the train state under a
trace instead -- `nnx.eval_shape` for the module graph and `jax.eval_shape`
over an all-Explicit view of the mesh for the layouts, since that is where JAX
carries a parameter's sharding into the moment allocated from it. Such an
engine can be compiled but not run; `fwd_bwd`, `update` and the checkpoint
methods raise rather than return something plausible.
The engine itself gains `compile_kernels()`, which routes through the same
`_compile_for_batch` the live engine calls on its first `fwd_bwd`, so the
ahead-of-time path cannot drift from the live one by construction, plus four
one-line hooks -- `_build_model`, `_build_optimizer`, `_checkpoint_dir` and
`_place_leaf` -- for the subclass to override.
`compile()` now compiles rather than only staging `jax.jit` closures for XLA to
run on the first `fwd_bwd`. It goes through `compile_kernels()`, which lowers,
applies `compile_xla_flags` (or a caller's `compiler_options`) and hands back
`{kernel name: Compiled}`; `compile()` installs those three executables and the
ahead-of-time entry point returns them for its report, so the live path and the
topology path compile through one piece of code. The kernel names both are keyed
by are `maxtext_engine.KERNEL_NAMES`. Lowering stays private in
`_lower_kernels`, since `jax.jit` offers no way to compile without lowering
first and no caller wants the halfway artifact; `_compile_for_batch` keeps the
jitted wrappers by name in `_jitted_kernels`, which is what it traces, because
an executable cannot be lowered again.
`compile()` also compiles the forward-only eval kernel, so a first `eval_step`
of the same shape dispatches rather than stalling on XLA. It is not part of an
update, so it is not one of `KERNEL_NAMES` and the ahead-of-time report does
not cover it.
`_is_jax_dynamic` now counts a `jax.ShapeDtypeStruct` as dynamic. Without that
an ahead-of-time batch is classified static and closed over as a constant,
which `jax.jit` rejects outright: an aval is not a valid JAX type.
Also fixes a pre-existing double compile. `nnx.Optimizer` builds optax's `count`
and its own `step` with `jnp.zeros` under no mesh, so they reach the first
update uncommitted on device 0 and come back from it committed across the mesh.
That is a second argument signature and a second full compile of the largest
kernel in the engine, for a program that runs once. `_place_state_on_mesh`
settles them up front, which also lets an ahead-of-time report describe the
steady state rather than the first step.
`tests/post_training/unit/maxtext_engine_xaot_test.py` asserts the optimized
HLO of all three kernels is identical, after the usual source-location
normalization, between a live engine that has stepped and an abstract one --
across data parallelism, Zero-1, FSDP, `shard_mode=auto` and bfloat16
gradients, plus qwen3-0.6b compiled for a v6e-4 topology. The parity classes
run in a re-execed subprocess: they need four CPU devices, and by the time
pytest imports the file a sibling module has already initialized the backend,
so setting `XLA_FLAGS` at import time is a no-op that would skip them green.
Both HLO rigs -- that file's and the data-parallel suite's -- spy on the
dispatch handles for the avals a kernel was called with, and recompile through
`_jitted_kernels` now that the handles are executables.
71b53d5 to
2f5cd1f
Compare
3 participants
Description
Stacked on #5104 (merged into #5099's branch). One reviewable commit.
trainers/pre_train/train_compile.pylets a pre-training configuration be costed and memory-checked against a target topology from a host that does not own it.MaxTextTrainingEnginehad no equivalent, so every iteration on a sharding or a batch size meant booking the hardware.This adds
training_engine/maxtext_engine_compile.py, which does the same for the engine's three kernels -- first forward/backward, accumulating forward/backward, update -- and reports each one's cost and memory. Nothing is materialized: weights and optimizer moments arejax.ShapeDtypeStructs and the mesh is a topology description. It reuses the pre-train path'sget_topology_meshandsave_compiled, socompile_topology,compile_xla_flagsandcompiled_trainstep_filekeep the meanings they already have; the three executables are written to three files suffixed by kernel name.qwen3-0.6b for a v6e-4, 38s on a host with no v6e allocation:
fwd_bwdfwd_bwd_accumupdateTwo supporting pieces in the engine:
materialize_weights=Falsebuilds the train state abstractly. The moments cannot come fromnnx.Optimizer, which allocates them withzeros_like, so the state is built under a trace instead:nnx.eval_shapefor the module graph, andjax.eval_shapeover an all-Explicitview of the mesh for the layouts, since that is where JAX carries a parameter's sharding into the moment allocated from it. Such an engine can be lowered but not run --fwd_bwd,updateand the checkpoint methods raise rather than return something plausible.lower()routes through the same_compile_for_batchthe live engine calls on its firstfwd_bwd, so the ahead-of-time path cannot drift from the live one by construction.Two bugs found on the way
_is_jax_dynamicdid not countjax.ShapeDtypeStruct. Left alone, an AOT batch is classified static, closed over as a constant, and the kernel lowered with no batch argument at all -- a smaller HLO and a memory report missing every activation, with nothing anywhere saying so._update_kernelwas compiling twice per run.nnx.Optimizerbuilds optax'scountand its ownstepwithjnp.zerosunder no mesh, so they reach the first update uncommitted on device 0 and come back from it committed across the mesh. That is a second argument signature and a second full compile of the largest kernel in the engine, for a program that runs once._place_state_on_meshsettles them up front (three scalardevice_puts). A pre-existing startup cost independent of AOT, but it also decided whether an AOT report describes the steady state or only step one.Tests
New:
tests/post_training/unit/maxtext_engine_xaot_test.py, 21 tests, on four simulated CPU devices plus a v6e-4 topology.The HLO of all three kernels is asserted byte-identical between a live engine that has actually stepped and an abstract one -- StableHLO compared exactly, optimized HLO after the repo's usual source-location normalization -- across data parallelism, Zero-1, FSDP,
shard_mode=autoand bfloat16 gradients. Plus:test_the_comparison_can_fail-- a byte-equality assertion is only as strong as its ability to tell two things apart, so this perturbs the model width and requires every comparison to notice. Width, not sequence length:_update_kernelnever sees a sequence, so the obvious knob would leave one of the three kernels silently untested.test_lowering_reproduces_the_kernels_training_ran-- licenses the test rig itself, which re-lowers from recorded argument avals.test_every_step_after_the_first_runs_the_same_kernels-- pins the double-compile fix.mainend to end writing one executable per kernel.No end-to-end workload: this change compiles but never executes a training step, and the engine's execution paths are covered by the four suites above.
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.