Add ahead-of-time (XAOT) compilation for MaxTextTrainingEngine

[NuojCheng]

Description

Stacked on #5104 (merged into #5099's branch). One reviewable commit.

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. MaxTextTrainingEngine had 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 are jax.ShapeDtypeStructs and the mesh is a topology description. It reuses the pre-train path's get_topology_mesh and save_compiled, so compile_topology, compile_xla_flags and compiled_trainstep_file keep the meanings they already have; the three executables are written to three files suffixed by kernel name.

python3 -m maxtext.training_engine.maxtext_engine_compile src/maxtext/configs/base.yml \
  model_name=qwen3-0.6b run_name=engine_aot_qwen3 \
  compile_topology=v6e-4 compile_topology_num_slices=1 \
  per_device_batch_size=4 max_target_length=2048 \
  ici_fsdp_parallelism=4 attention=flash enable_checkpointing=false

qwen3-0.6b for a v6e-4, 38s on a host with no v6e allocation:

kernel	argument	output	alias	temp
fwd_bwd	596 MB	596 MB	0	6.41 GB
fwd_bwd_accum	1.19 GB	596 MB	596 MB	6.59 GB
update	2.39 GB	1.79 GB	1.79 GB	4.7 MB

Two supporting pieces in the engine:

• materialize_weights=False builds the train state abstractly. The moments cannot come from nnx.Optimizer, which allocates them with zeros_like, so the state is built 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 lowered but not run -- fwd_bwd, update and the checkpoint methods raise rather than return something plausible.
• lower() 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.

Two bugs found on the way

_is_jax_dynamic did not count jax.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_kernel was compiling twice per run. 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 (three scalar device_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=auto and 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_kernel never 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.
• qwen3-0.6b compiled for a v6e-4 topology, and main end to end writing one executable per kernel.

pytest tests/post_training/unit/maxtext_engine_xaot_test.py          # 21 passed, 80s
                                                                     # (26 with libtpu, for the v6e-4 topology tests)
pytest tests/post_training/unit/maxtext_engine_test.py \
       tests/post_training/unit/maxtext_engine_deferred_all_reduce_test.py \
       tests/post_training/unit/maxtext_engine_zero1_test.py \
       tests/post_training/unit/maxtext_engine_constructor_test.py   # 90 passed

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):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

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.

[gemini-code-assist]

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)

[NuojCheng]

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.

[codecov]

Codecov Report

❌ Patch coverage is 86.24339% with 26 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
.../maxtext/training_engine/maxtext_engine_compile.py	86.60%	10 Missing and 5 partials ⚠️
src/maxtext/training_engine/maxtext_engine.py	85.71%	6 Missing and 5 partials ⚠️

📢 Thoughts on this report? Let us know!

[gobbleturk]

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

[NuojCheng]

I am hesitant on adding more additional maxtext flags for this feature... Let's do that in followup PRs.

[NuojCheng]

Unless there is a way to avoid adding maxtext flags for this.

[igorts-git]

please add actual pre-compilation here.

[NuojCheng]

add

[igorts-git]

also compile here?

[NuojCheng]

added compile eval