From 2f5cd1fa6be8b15b2607243da2c0ec8edad7a563 Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Fri, 4 Sep 2026 22:00:01 +0000 Subject: [PATCH] Add ahead-of-time compilation for MaxTextTrainingEngine `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. --- src/maxtext/training_engine/maxtext_engine.py | 232 +++++- .../training_engine/maxtext_engine_compile.py | 312 +++++++++ .../unit/maxtext_engine_data_parallel_test.py | 25 +- .../unit/maxtext_engine_xaot_test.py | 660 ++++++++++++++++++ 4 files changed, 1194 insertions(+), 35 deletions(-) create mode 100644 src/maxtext/training_engine/maxtext_engine_compile.py create mode 100644 tests/post_training/unit/maxtext_engine_xaot_test.py diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py index 9153180a85..d572fcacff 100644 --- a/src/maxtext/training_engine/maxtext_engine.py +++ b/src/maxtext/training_engine/maxtext_engine.py @@ -61,6 +61,10 @@ # Where the same split puts the `nnx.Optimizer`, including the optax state Zero-1 shards. _OPTIMIZER_STATE_KEY = "optimizer" +# The kernels an update is made of, in the order a step runs them. `compile_kernels()` is +# keyed by these, and so is the ahead-of-time entry point in `maxtext_engine_compile`. +KERNEL_NAMES = ("fwd_bwd", "fwd_bwd_accum", "update") + _PURE_STATE_FALLBACK_WARNING = ( "Cannot keep the train state as a pure pytree across steps (%s), so every fwd_bwd and " "update will re-walk the NNX module graph. That is correct but slow -- the two " @@ -75,14 +79,16 @@ def _is_jax_dynamic(value: Any) -> bool: A `gen_model_input_fn` returns the loss function's keyword arguments, and only some of them are arrays. Tunix's GRPO adapter, for instance, returns a `TrainExample` alongside an `algo_config` object and integer `pad_id`/`eos_id`. The arrays must be traced; the - rest must be closed over, or `jax.jit` rejects the call outright. + rest must be closed over, or `jax.jit` rejects the call outright. `jax.ShapeDtypeStruct` + counts too: `compile_kernels()` drives the whole path on shapes alone, and an aval closed over as a + constant cannot be traced at all. """ leaves = jax.tree.leaves(value) if not leaves: # An all-`None` subtree (e.g. an unset `ref_per_token_logps`) flattens to nothing. It # carries no data either way, so tracing it is harmless and keeps the treedef intact. return True - return any(isinstance(leaf, (jax.Array, np.ndarray, np.generic)) for leaf in leaves) + return any(isinstance(leaf, (jax.Array, jax.ShapeDtypeStruct, np.ndarray, np.generic)) for leaf in leaves) def _split_static_and_dynamic(batch: Any) -> tuple[Any, dict[str, Any]]: @@ -348,6 +354,13 @@ def _normalize_loss_output(out: Any, has_aux: bool) -> abstract_engine.LossOutpu ) +def _to_aval(value: Any) -> Any: + """Returns `value` as a `jax.ShapeDtypeStruct` on the sharding it carries, or unchanged if it has no shape.""" + if not hasattr(value, "shape") or not hasattr(value, "dtype"): + return value + return jax.ShapeDtypeStruct(value.shape, value.dtype, sharding=getattr(value, "sharding", None)) + + @struct.dataclass(frozen=True, kw_only=True) class RouterReplayTrainerPayload(abstract_engine.TrainerPayload): """A TrainerPayload extension carrying forced router-replay expert decisions. @@ -538,26 +551,15 @@ def __init__( # every eager caller compile or make a deferred compile never happen. self._compile_requested = False self._compiled_signature: Any = None + # `{kernel name: the jitted wrapper}` staged by the last `_compile_for_batch`; empty until + # then, and only ever read straight after one, since a recompile replaces every entry. + self._jitted_kernels: dict[str, Any] = {} self._compiled_eval: Any = None self._compiled_eval_signature: Any = None self._signature_compare_warned: bool = False if not training_config.model_name: raise ValueError("training_config.model_name must be specified") - model_or_model_mesh_pair = model_creation_utils.from_pretrained( - config=self._config, - mesh=self._mesh, - model_mode=common_types.MODEL_MODE_TRAIN, - rng_key=self._init_rng, - wrap_with_tunix_adapter=wrap_with_tunix_adapter, - tokenizer_pad_id=tokenizer_pad_id, - ) - # `from_pretrained` returns `(model, mesh)` when it had to derive the mesh itself, and just the model - # when one was supplied. Adopt the derived mesh so `self._model` is always a module and `compile()` can - # still build shardings. - if self._mesh is None: - self._model, self._mesh = model_or_model_mesh_pair - else: - self._model = model_or_model_mesh_pair + self._model = self._build_model(wrap_with_tunix_adapter, tokenizer_pad_id) self._state: Any = None # Pure-pytree mirror of the model and train state, carried across steps so the step path # never re-walks the module graph. `None` means "not cached". @@ -592,11 +594,11 @@ def __init__( # `checkpointing.CheckpointState` expects an nnx.Optimizer too, so wrap it here. `wrt=nnx.Param` # covers every parameter, which is correct only because LoRA is rejected above. self._learning_rate_schedule, tx = train_utils.create_training_optimizer(self._config, self._model) - self._optimizer = nnx.Optimizer(self._model, tx, wrt=nnx.Param) + self._optimizer = self._build_optimizer(tx) self._train_step: int = 0 self._checkpoint_manager = checkpointing.CheckpointManager( - checkpoint_dir=self._config.checkpoint_dir, + checkpoint_dir=self._checkpoint_dir(), config=self._config, ) self._metrics_recorder = metrics_module.MetricsRecorder() @@ -605,6 +607,35 @@ def __init__( self._throttler = inflight_throttler.InflightThrottler(config=self._config, metrics_logger=self._metrics_logger) self._raiden_sync: Any = None + def _build_model(self, wrap_with_tunix_adapter: bool, tokenizer_pad_id: int | None) -> Any: + """Returns the model to train, adopting a mesh when this engine was given none.""" + model_or_model_mesh_pair = model_creation_utils.from_pretrained( + config=self._config, + mesh=self._mesh, + model_mode=common_types.MODEL_MODE_TRAIN, + rng_key=self._init_rng, + wrap_with_tunix_adapter=wrap_with_tunix_adapter, + tokenizer_pad_id=tokenizer_pad_id, + ) + # `from_pretrained` returns `(model, mesh)` only when it had to derive the mesh itself. Adopt the + # derived one so `self._model` is always a module and `compile()` can still build shardings. + if self._mesh is not None: + return model_or_model_mesh_pair + model, self._mesh = model_or_model_mesh_pair + return model + + def _build_optimizer(self, tx: Any) -> Any: + """Returns the `nnx.Optimizer` for `self._model`. + + A subclass that cannot allocate moments overrides this, and may install `self._state` and + rebind `self._model` on the way: `__init__` does not touch either again. + """ + return nnx.Optimizer(self._model, tx, wrt=nnx.Param) + + def _checkpoint_dir(self) -> str: + """Returns the directory this engine checkpoints through; an empty string disables Orbax entirely.""" + return self._config.checkpoint_dir + @property def model(self) -> Any: """Returns the NNX model instance.""" @@ -883,6 +914,42 @@ def target(leaf, base): return jax.tree.map(target, params_pure, params_shardings) + def _place_leaf(self, leaf: Any, target: jax.sharding.Sharding) -> Any: + """Returns one train-state leaf committed to `target`.""" + return jax.device_put(leaf, target) + + def _place_state_on_mesh(self) -> None: + """Commits every train-state leaf to this mesh, in place, before anything is compiled. + + `nnx.Optimizer` builds optax's `count` and its own `step` with `jnp.zeros` under no mesh, so + they reach the first update uncommitted and come back from it committed -- a second argument + signature, and a second compile of the largest kernel in the engine. Settling them up front + makes steps one and two the same program. Also covers state that arrives later, from a + restore or a public setter. + """ + if self._mesh is None or self._state is None: + return + moved = False + + def place(leaf): + nonlocal moved + # `device_put` would turn a Python scalar in the state into a device array. + if not hasattr(leaf, "shape") or not hasattr(leaf, "dtype"): + return leaf + leaf_sharding = getattr(leaf, "sharding", None) + if isinstance(leaf_sharding, jax.sharding.NamedSharding) and leaf_sharding.mesh == self._mesh: + return leaf + moved = True + return self._place_leaf(leaf, self._mesh_sharding(leaf)) + + placed = jax.tree.map(place, self._read_state_pure()) + if not moved: + return + with self._sharding_ctx(): + nnx.update(self._state, placed) + self._invalidate_pure_state() + self._refresh_pure_state() + def _shard_optimizer_state_over_data(self) -> None: """Moves the optimizer's parameter-shaped state onto the Zero-1 layout, in place. @@ -913,7 +980,7 @@ def place(leaf): if target is None: return leaf moved = True - return jax.device_put(leaf, target) + return self._place_leaf(leaf, target) optimizer_pure = jax.tree.map(place, state_pure[_OPTIMIZER_STATE_KEY]) if not moved: @@ -1205,6 +1272,8 @@ def _compile_for_batch(self, dynamic_batch: Any, static_batch: dict[str, Any]) - # The only place the graphs are walked: a recompile is when they may legitimately have # changed shape, and everything after is maintained as plain pytrees. self._refresh_pure_state() + # Before the Zero-1 pass and the shardings read off the state below. + self._place_state_on_mesh() # Before the shardings below are read off the state: this is what puts the optimizer # moments on the Zero-1 layout, and `state_mesh_shardings` has to see them there. self._shard_optimizer_state_over_data() @@ -1302,6 +1371,14 @@ def accum_kernel(params, rest, dynamic, acc_grads, acc_denom): out_shardings=update_out_shardings, donate_argnums=(0,), ) + # The same three wrappers by name, which is what `_lower_kernels()` traces: `compile()` overwrites + # the attributes above with the executables they lower to, and an executable cannot be + # lowered again. + self._jitted_kernels = { + "fwd_bwd": self._compiled_fwd_bwd, + "fwd_bwd_accum": self._compiled_fwd_bwd_accum, + "update": self._compiled_update, + } self._compiled_signature = _batch_signature(dynamic_batch, static_batch) self._compiled = True @@ -1329,14 +1406,19 @@ def kernel(params, rest, dynamic): ) self._compiled_eval_signature = _batch_signature(dynamic_batch, static_batch) - def compile(self, dummy_data: abstract_engine.TrainerPayload) -> None: - """Triggers SPMD JIT compilation of fwd_bwd and update steps. + def compile( + self, + dummy_data: abstract_engine.TrainerPayload, + compiler_options: dict[str, Any] | None = None, + ) -> None: + """Triggers SPMD compilation of the fwd_bwd, update and eval steps. Args: dummy_data: Sample TrainerPayload providing representative tensor shapes. Its shapes must match the real batches, or the first `fwd_bwd` simply recompiles. When it is `None` the engine cannot know the input shapes, so it stays on the eager path and compiles lazily on the first `fwd_bwd` instead. + compiler_options: Optional dictionary of compilation options passed to XLA compiler. """ # Recorded even when compilation is deferred: it is what tells `fwd_bwd` the caller # wants the compiled path at all. Engines that never call `compile` stay eager. @@ -1346,19 +1428,117 @@ def compile(self, dummy_data: abstract_engine.TrainerPayload) -> None: if dummy_data is None: # Callers driving a generic worker lifecycle (Tunix's `TrainerWorker.compile`) pass - # nothing to compile against. Deferring costs nothing measurable: `jax.jit` is lazy, - # so even with a payload this method only stages the wrappers and XLA still runs on - # the first `fwd_bwd`. Logged at info, not warning -- the first `fwd_bwd` compiles - # against the real batch, whose shapes are right by construction. + # nothing to compile against. When dummy_data is None the engine cannot know the + # input shapes, so it compiles against the first fwd_bwd payload instead. logging.info( "MaxTextTrainingEngine.compile() was called without dummy_data; compiling " "against the first fwd_bwd payload instead." ) return + compiled = self.compile_kernels(dummy_data, compiler_options) + self._compiled_fwd_bwd = compiled["fwd_bwd"] + self._compiled_fwd_bwd_accum = compiled["fwd_bwd_accum"] + self._compiled_update = compiled["update"] + self._compile_eval(dummy_data, compiler_options) + + def compile_kernels( + self, + dummy_data: abstract_engine.TrainerPayload, + compiler_options: dict[str, Any] | None = None, + ) -> dict[str, jax.stages.Compiled]: + """Lowers and compiles every kernel, and hands them back rather than installing them. + + The body of `compile()` with the engine's own bookkeeping left out, so a caller that + only wants the executables -- the ahead-of-time path, which compiles for a topology it + cannot run on -- gets them from the same code the live path uses. + + Args: + dummy_data: As `compile()`'s, but required: there is no first `fwd_bwd` to defer to. + compiler_options: XLA options, defaulting to `config.compile_xla_flags`. + + Returns: + `{kernel name: jax.stages.Compiled}`, keyed by `KERNEL_NAMES`. + """ + options = self._xla_options(compiler_options) + lowered = self._lower_kernels(dummy_data) + with self._sharding_ctx(): + return {name: lowered[name].compile(compiler_options=options) for name in KERNEL_NAMES} + + def _xla_options(self, compiler_options: dict[str, Any] | None) -> dict[str, Any] | None: + """Returns the XLA options to compile with: the caller's, or `config.compile_xla_flags`.""" + if compiler_options is None and getattr(self._config, "compile_xla_flags", ""): + return max_utils.parse_libtpu_flags_to_dict(self._config.compile_xla_flags) + return compiler_options + + def _compile_eval(self, dummy_data: abstract_engine.TrainerPayload, compiler_options: dict[str, Any] | None) -> None: + """Compiles the forward-only eval kernel, so the first `eval_step` does not stall on XLA. + + The eval kernel is not part of an update, so it is not one of `KERNEL_NAMES` and the + ahead-of-time report does not cover it; this is only for the live engine. An eval batch + shaped unlike `dummy_data` still recompiles inside `eval_step`, as an unforeseen training + batch does inside `fwd_bwd`. + """ + dynamic_batch, static_batch = _split_static_and_dynamic(self._prepare_batch(dummy_data)) + self._compile_eval_for_batch(dynamic_batch, static_batch) + params_pure, rest_pure = self._read_model_pure(getattr(self._state, _MODEL_STATE_KEY, self._model)) + with self._sharding_ctx(): + # `_compile_eval_for_batch` leaves the jitted wrapper here; lowering it replaces it with + # what it compiles to, which is what `eval_step` then dispatches through. + self._compiled_eval = self._compiled_eval.lower( + jax.tree.map(_to_aval, params_pure), + jax.tree.map(_to_aval, rest_pure), + jax.tree.map(_to_aval, dynamic_batch), + ).compile(compiler_options=self._xla_options(compiler_options)) + + def _lower_kernels(self, dummy_data: abstract_engine.TrainerPayload) -> dict[str, jax.stages.Lowered]: + """Lowers every kernel this engine runs, the half of `compile_kernels` before XLA runs. + + Not public: `jax.jit` offers no way to compile without lowering first, so this exists because + `compile_kernels` needs it, not because a caller does. Routes through `_compile_for_batch`, + the same method the live path calls on its first `fwd_bwd`, so the shapes and the shardings + are the live ones by construction. The accumulating kernel is lowered even for a + single-micro-batch run, where the live engine never traces it: omitting the kernel that holds + the extra parameter-sized accumulator would understate the peak that decides whether a + configuration fits. + + Args: + dummy_data: One micro-batch, real or abstract, whose structure must match the batches the + engine will be given, exactly as `compile()`'s does. + + Returns: + `{kernel name: jax.stages.Lowered}`, keyed by `KERNEL_NAMES`. + + Raises: + ValueError: If `dummy_data` is None. + """ + if dummy_data is None: + raise ValueError( + "compile_kernels() needs a dummy payload -- unlike compile(), it cannot defer to the first real batch." + ) dynamic_batch, static_batch = _split_static_and_dynamic(self._prepare_batch(dummy_data)) self._compile_for_batch(dynamic_batch, static_batch) + state_aval = jax.tree.map(_to_aval, self._read_state_pure()) + params_pure, rest_pure = self._read_model_pure(getattr(self._state, _MODEL_STATE_KEY, self._model)) + params_aval = jax.tree.map(_to_aval, params_pure) + rest_aval = jax.tree.map(_to_aval, rest_pure) + batch_aval = jax.tree.map(_to_aval, dynamic_batch) + mean_loss_aval = jax.ShapeDtypeStruct((), jnp.float32) if self._config.skip_step_on_spikes else None + + with self._sharding_ctx(): + fwd_bwd = self._jitted_kernels["fwd_bwd"].lower(params_aval, rest_aval, batch_aval) + # Off the kernel's own outputs, not predicted from the parameters: the gradients differ by + # `grad_dtype` and, under deferral, an `unreduced` tag. + _, _, _, grads_aval, denominator_aval = fwd_bwd.out_info + return { + "fwd_bwd": fwd_bwd, + "fwd_bwd_accum": self._jitted_kernels["fwd_bwd_accum"].lower( + params_aval, rest_aval, batch_aval, grads_aval, denominator_aval + ), + "update": self._jitted_kernels["update"].lower(state_aval, grads_aval, denominator_aval, mean_loss_aval), + } + def fwd_bwd(self, payload: abstract_engine.TrainerPayload, **kwargs: Any) -> None: """Executes a micro-batch forward-backward pass and accumulates gradients. diff --git a/src/maxtext/training_engine/maxtext_engine_compile.py b/src/maxtext/training_engine/maxtext_engine_compile.py new file mode 100644 index 0000000000..caa961dcb5 --- /dev/null +++ b/src/maxtext/training_engine/maxtext_engine_compile.py @@ -0,0 +1,312 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Ahead-of-time (XAOT) compilation of `MaxTextTrainingEngine`'s training step. + +`trainers/pre_train/train_compile.py` does this for `train.py`'s single fused `train_step`. +The engine splits the same work across three kernels -- one forward/backward for the first +micro-batch of an update, an accumulating one for every later micro-batch, and the optimizer +update -- so this compiles all three and reports the cost and memory of each. + +Nothing is materialized: the weights, the optimizer moments and the batch are all +`jax.ShapeDtypeStruct`s, and the device mesh is a topology description rather than hardware. +So a v5e-256 configuration can be compiled from a workstation, and an out-of-memory one +reports the same `RESOURCE_EXHAUSTED` it would report on the target -- before the target is +booked. + +Example, qwen3-0.6b on four v6e chips: + + 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 + +Add `compiled_trainstep_file=/tmp/engine_qwen3.pickle` to serialize the executables; each +kernel is written to its own file, suffixed with the kernel name. +""" + +import os +from typing import Any, Sequence + +from absl import app +from flax import nnx +import jax +from maxtext.common import common_types +from maxtext.common import train_state_nnx +from maxtext.configs import pyconfig +from maxtext.trainers.pre_train import train_compile as pre_train_compile +from maxtext.training_engine import maxtext_engine +from maxtext.utils import gcs_utils +from maxtext.utils import max_utils +from maxtext.utils import maxtext_utils +from maxtext.utils import model_creation_utils + +# Re-exported: which kernels there are is the engine's to say, and this module reports on +# whatever it lowers. +KERNEL_NAMES = maxtext_engine.KERNEL_NAMES + +# Both `dump_hlo` filters default to `jit_train_step`, `train.py`'s fused step; the engine's +# kernels lower as `jit_first_kernel`, `jit_accum_kernel` and `jit__update_kernel`, so on those +# defaults the dump comes back empty. +HLO_DUMP_DEFAULTS = { + "dump_hlo_local_module_name": "jit_.*kernel", + "dump_hlo_module_name": "kernel", +} + + +def _propagation_mesh(mesh: jax.sharding.Mesh) -> jax.sharding.Mesh: + """Returns a stand-in for `mesh` that `jax.eval_shape` will propagate shardings across. + + Nothing runs on it. `jax.eval_shape` carries a value's layout through an operation only on + `Explicit` axes, so under `shard_mode=auto` -- where every axis is `Auto` -- the moments would + all come back replicated. Marking the axes `Explicit` is how the layouts are *observed*; they + are re-homed onto the real mesh afterwards, and its axis types decide what actually runs. + """ + axis_types = getattr(mesh, "axis_types", None) + if axis_types is not None and all(axis_type == jax.sharding.AxisType.Explicit for axis_type in axis_types): + return mesh + return jax.sharding.Mesh( + mesh.devices, + mesh.axis_names, + axis_types=(jax.sharding.AxisType.Explicit,) * len(mesh.axis_names), + ) + + +def _rehome_aval(aval: Any, mesh: jax.sharding.Mesh) -> Any: + """Returns `aval` with its sharding spec re-expressed on `mesh`. + + The engine's `_mesh_sharding` compares meshes by equality, so a spec that is right but homed + on the trace's mesh would be silently replaced by a replicated one. + """ + if not hasattr(aval, "shape") or not hasattr(aval, "dtype"): + return aval + 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) + + +class AbstractMaxTextEngine(maxtext_engine.MaxTextTrainingEngine): + """A `MaxTextTrainingEngine` whose weights and moments are shapes rather than arrays. + + Enough to trace and compile every kernel, which is all `compile_kernels()` needs, while nothing + is allocated and no checkpoint, tokenizer or network is touched. Nothing can be executed. + """ + + def __init__(self, training_config: pyconfig.HyperParameters, mesh: jax.sharding.Mesh) -> None: + """Initializes an engine that can be lowered but not run. + + Args: + training_config: MaxText HyperParameters configuration instance. + mesh: The mesh to compile against, typically a topology this host does not own. + + Raises: + ValueError: If `mesh` is None. With no weights there is no device set to read one off. + """ + if mesh is None: + raise ValueError( + "AbstractMaxTextEngine requires a mesh: with no weights there is nothing to read a device set " + "off, and the point of the abstract path is to compile against a mesh this host does not own -- " + "build one with `trainers.pre_train.train_compile.get_topology_mesh`." + ) + super().__init__(training_config, mesh=mesh) + + def _build_model(self, wrap_with_tunix_adapter: bool, tokenizer_pad_id: int | None) -> Any: + """Returns the model with `jax.ShapeDtypeStruct` weights on their real shardings. + + The same `create_nnx_abstract_model` call `from_pretrained` makes before it materializes + anything, minus the checkpoint load -- so no weights, no HF token and no network. + """ + del wrap_with_tunix_adapter, tokenizer_pad_id # `__init__` accepts neither. + _, abstract_model = model_creation_utils.create_nnx_abstract_model( + model_creation_utils.verify_and_sync_scan_layers(self._config), + self._mesh, + model_mode=common_types.MODEL_MODE_TRAIN, + rng_key=self._init_rng, + ) + return abstract_model + + def _build_optimizer(self, tx: Any) -> Any: + """Installs the traced train state and returns the optimizer inside it. + + `self._model` is rebound to the model inside that state so the two stay one graph. + """ + self._state = self._trace_train_state(tx) + self._model = self._state.model + return self._state.optimizer + + def _trace_train_state(self, tx: Any) -> Any: + """Returns the `TrainStateNNX` for this model, moments included, as avals. + + `nnx.Optimizer` allocates the moments eagerly with `zeros_like`, so they are traced instead. + Two traces, because neither alone answers both questions: `nnx.eval_shape` gives the module + graph but drops shardings, and `jax.eval_shape` under `_propagation_mesh` gives the + layouts, because that is where JAX carries a parameter's sharding through the `zeros_like` + inside `tx.init` into the moment allocated from it. The result is merged back onto the real + mesh, whose axis types -- not the stand-in's -- decide what the compiled kernels do. + + Both run under the engine's own `_sharding_ctx`, so the rules the MaxText layers are written + against are the live ones rather than a second copy that can drift from them. + """ + model_graphdef, model_pure = nnx.split(self._model) + + def build(model_state): + model = nnx.merge(model_graphdef, model_state) + return train_state_nnx.TrainStateNNX(model, nnx.Optimizer(model, tx, wrt=nnx.Param)) + + propagation_mesh = _propagation_mesh(self._mesh) + with self._sharding_ctx(): + state_graphdef, _ = nnx.split(nnx.eval_shape(build, model_pure)) + # Displaces the real mesh for this trace only: the one above needs no propagation, and a + # stand-in set around it collides with the config's own AbstractMesh under `shard_mode=auto`. + with jax.set_mesh(propagation_mesh): + state_pure = jax.eval_shape( + lambda model_state: nnx.split(build(model_state))[1], + jax.tree.map(lambda aval: _rehome_aval(aval, propagation_mesh), model_pure), + ) + return nnx.merge(state_graphdef, jax.tree.map(lambda aval: _rehome_aval(aval, self._mesh), state_pure)) + + def _checkpoint_dir(self) -> str: + """Returns no directory: Orbax creates whatever it is given, and this engine can never save.""" + return "" + + def _place_leaf(self, leaf: Any, target: jax.sharding.Sharding) -> Any: + """Restates the aval on `target`: there is nothing to move, and `device_put` takes no aval.""" + return jax.ShapeDtypeStruct(leaf.shape, leaf.dtype, sharding=target) + + def _cannot_run(self, operation: str) -> RuntimeError: + """Returns the error every execution entry point raises instead of running.""" + return RuntimeError( + f"AbstractMaxTextEngine.{operation}() needs real weights, and this engine has only shapes. It " + "exists to be traced and compiled: call `compile_kernels()`, or build a MaxTextTrainingEngine instead." + ) + + def fwd_bwd(self, payload: Any, **kwargs: Any) -> None: + raise self._cannot_run("fwd_bwd") + + def update(self, **kwargs: Any) -> int: + raise self._cannot_run("update") + + def save_checkpoint(self, metadata: Any, **kwargs: Any) -> None: + raise self._cannot_run("save_checkpoint") + + def restore_checkpoint(self, **kwargs: Any) -> Any: + raise self._cannot_run("restore_checkpoint") + + +def with_engine_hlo_dump_defaults(argv: Sequence[str]) -> list[str]: + """Returns `argv` with the HLO dump filters pointed at the kernels, if a dump was asked for. + + On `argv` rather than the config, and only under `dump_hlo`: `pyconfig.initialize` bakes the + regex into `XLA_FLAGS` whether or not a dump was asked for, so widening it unconditionally + would leave every compile writing a dump nobody asked for. + """ + given = dict(arg.split("=", 1) for arg in argv if "=" in arg) + if given.get("dump_hlo", "").strip().lower() not in ("true", "1"): + return list(argv) + return list(argv) + [f"{key}={value}" for key, value in HLO_DUMP_DEFAULTS.items() if key not in given] + + +def get_shaped_micro_batch(config: pyconfig.HyperParameters) -> dict[str, jax.ShapeDtypeStruct]: + """Returns the abstract batch one `fwd_bwd` call is given. + + `maxtext_utils.get_shaped_batch` shapes the *global* batch, because `train.py`'s fused step + folds gradient accumulation inside itself. The engine's caller drives one `fwd_bwd` per + micro-batch, so compiling against the global batch would size every activation by the + accumulation factor and report a peak memory no step ever reaches. + """ + shaped_batch = maxtext_utils.get_shaped_batch(config) + micro_batch_size = int(config.micro_batch_size_to_train_on) + + def to_micro_batch(aval: jax.ShapeDtypeStruct) -> jax.ShapeDtypeStruct: + if not aval.shape or aval.shape[0] == micro_batch_size: + return aval + return jax.ShapeDtypeStruct((micro_batch_size,) + aval.shape[1:], aval.dtype) + + return {key: to_micro_batch(aval) for key, aval in shaped_batch.items()} + + +def compile_engine_kernels(config: pyconfig.HyperParameters, topology_mesh: jax.sharding.Mesh) -> dict[str, Any]: + """Lowers and compiles every kernel the engine runs, on `topology_mesh`. + + Returns: + `{kernel name: jax.stages.Compiled}`, keyed by `KERNEL_NAMES`. + """ + return AbstractMaxTextEngine(config, topology_mesh).compile_kernels(get_shaped_micro_batch(config)) + + +def kernel_save_path(compiled_trainstep_file: str, kernel_name: str) -> str: + """Returns where one kernel's executable goes: `/tmp/engine.pickle` -> `/tmp/engine_fwd_bwd.pickle`.""" + stem, extension = os.path.splitext(compiled_trainstep_file) + return f"{stem}_{kernel_name}{extension}" + + +def main(argv: Sequence[str]) -> None: + """Compiles the engine's kernels for `compile_topology` and reports what they cost.""" + jax.config.update("jax_default_prng_impl", "unsafe_rbg") + os.environ["LIBTPU_INIT_ARGS"] = ( + os.environ.get("LIBTPU_INIT_ARGS", "") + " --xla_tpu_spmd_rng_bit_generator_unsafe=true" + ) + print("Starting training_engine/maxtext_engine_compile.py...", flush=True) + + config = pyconfig.initialize(with_engine_hlo_dump_defaults(argv)) + pre_train_compile.validate_config(config) + if config.enable_diloco: + raise NotImplementedError( + "enable_diloco is not supported here: MaxTextTrainingEngine has no DiLoCo outer step, so the " + "numbers reported would describe a different computation." + ) + + topology_mesh = pre_train_compile.get_topology_mesh(config) + + # After the topology is built, so this does not initialize the local backend first. + max_utils.print_system_information() + + print("Jitting and compiling the engine's kernels...", flush=True) + compiled = compile_engine_kernels(config, topology_mesh) + print("Jitting and compilation complete!", flush=True) + + for name in KERNEL_NAMES: + print(f"--- {name} ---") + print(f"Cost analysis: {compiled[name].cost_analysis()}") + print(f"Memory analysis: {compiled[name].memory_analysis()}") + + if config.compiled_trainstep_file != "": + for name in KERNEL_NAMES: + save_path = kernel_save_path(config.compiled_trainstep_file, name) + pre_train_compile.save_compiled(compiled[name], save_path) + print(f"Successfully saved compiled {name} kernel as {save_path}") + + print("Finished training_engine/maxtext_engine_compile.py successfully!", flush=True) + + if config.dump_hlo: + # `upload_dump` deletes what it uploaded; say which filter was too narrow rather than raise + # from the rmtree of a directory XLA never wrote. + if not os.path.isdir(config.dump_hlo_local_dir): + raise FileNotFoundError( + f"dump_hlo is set but XLA wrote nothing to {config.dump_hlo_local_dir}: " + f"dump_hlo_local_module_name={config.dump_hlo_local_module_name!r} matched none of the engine's " + f"kernels (jit_first_kernel, jit_accum_kernel, jit__update_kernel)." + ) + gcs_utils.upload_dump( + config.dump_hlo_local_dir, + config.dump_hlo_gcs_dir, + module_name=config.dump_hlo_module_name, + delete_local_after=config.dump_hlo_delete_local_after, + all_host_upload=config.dump_hlo_upload_all, + ) + + +if __name__ == "__main__": + app.run(main) diff --git a/tests/post_training/unit/maxtext_engine_data_parallel_test.py b/tests/post_training/unit/maxtext_engine_data_parallel_test.py index 6e3b553da8..bf831681f5 100644 --- a/tests/post_training/unit/maxtext_engine_data_parallel_test.py +++ b/tests/post_training/unit/maxtext_engine_data_parallel_test.py @@ -257,18 +257,20 @@ def _gathered_elements(hlo: str) -> int: class _KernelHlo: """Captures the arguments the engine passes one jitted kernel, to re-lower it later. - `jax.jit` keeps no handle on the executable it cached, so the only way to read a kernel's - optimized HLO is to lower it again. Recording `ShapeDtypeStruct`s rather than the arrays - themselves keeps that independent of donation -- `_compiled_update` donates its state, so - by the time a test asks for the HLO those buffers are gone. + An executable keeps no handle on the lowering it came from, so the only way to read a + kernel's optimized HLO is to lower it again, through the jitted wrapper `compile()` built it + from. Recording `ShapeDtypeStruct`s rather than the arrays themselves keeps that independent + of donation -- `_compiled_update` donates its state, so by the time a test asks for the HLO + those buffers are gone. Lowering is lazy and memoized: it is a full XLA compile, and only a third of the tests below want HLO at all. It needs the mesh set, because the fwd/bwd kernels apply the `reduced` tag while tracing, and by then the run that recorded the avals is long over. """ - def __init__(self, engine, attr: str, mesh): - self._jitted = getattr(engine, attr) + def __init__(self, engine, kernel: str, attr: str, mesh): + self._dispatched = getattr(engine, attr) + self._jitted = engine._jitted_kernels[kernel] # pylint: disable=protected-access self._mesh = mesh self._avals = None self._text = None @@ -282,7 +284,7 @@ def _spy(self, *args): lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype, sharding=getattr(x, "sharding", None)), args, ) - return self._jitted(*args) + return self._dispatched(*args) def text(self) -> str: if self._avals is None: @@ -308,7 +310,12 @@ def overrides(self) -> dict[str, Any]: return {"shard_optimizer_over_data": self.zero1, "opt_type": self.opt_type} -_KERNELS = {"first": "_compiled_fwd_bwd", "accum": "_compiled_fwd_bwd_accum", "update": "_compiled_update"} +# Probe name -> the engine's name for that kernel, and the attribute it dispatches through. +_KERNELS = { + "first": ("fwd_bwd", "_compiled_fwd_bwd"), + "accum": ("fwd_bwd_accum", "_compiled_fwd_bwd_accum"), + "update": ("update", "_compiled_update"), +} class _Run: @@ -330,7 +337,7 @@ def __init__(self, recipe: _Recipe): stack.enter_context(jax.set_mesh(self.mesh)) self.engine = maxtext_engine.MaxTextTrainingEngine(self.config, mesh=self.mesh) self.engine.compile(_batch(self.config, 0)) - self._probes = {name: _KernelHlo(self.engine, attr, self.mesh) for name, attr in _KERNELS.items()} + self._probes = {name: _KernelHlo(self.engine, kernel, attr, self.mesh) for name, (kernel, attr) in _KERNELS.items()} norms = [] for step in range(recipe.steps): for micro in range(recipe.micro_batches): diff --git a/tests/post_training/unit/maxtext_engine_xaot_test.py b/tests/post_training/unit/maxtext_engine_xaot_test.py new file mode 100644 index 0000000000..7edeb7e05f --- /dev/null +++ b/tests/post_training/unit/maxtext_engine_xaot_test.py @@ -0,0 +1,660 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Ahead-of-time compilation of the engine compiles the same thing training runs. + +`training_engine/maxtext_engine_compile.py` reports a configuration's cost and peak memory from +shapes alone, with no weights and no hardware. That report is worth only the claim behind it -- +that the kernels it compiled are the kernels training will run -- and `AbstractMaxTextEngine` +reaches them differently: it traces its optimizer moments rather than allocating them, and +derives their shardings by propagation rather than off real arrays. + +So these tests compare the optimized HLO of all three kernels against a live engine that has +actually stepped. `test_the_comparison_can_fail` guards the vacuous case. + +The parity tests run in a subprocess. Their shardings need more than one device, and the CPU +backend reads `--xla_force_host_platform_device_count` only at backend initialization -- which is +already past by the time pytest imports this file, because sibling modules under +`tests/post_training` touch JAX while being collected. Setting the flag at import time here is +therefore a no-op that leaves `jax.device_count() == 1` and skips the parity classes green. +`test_engine_aot_parity_on_a_four_device_cpu_mesh` re-execs the module with the flag appended +instead, and refuses to pass unless the child reports tests actually run. +""" + +# Every `_engine._private` below reads a member of the class under test from its own test. +# pylint: disable=protected-access + +import os +import re +import subprocess +import sys +import unittest + +from absl.testing import absltest +from absl.testing import parameterized +from flax import nnx +import jax +from maxtext.configs import pyconfig +from maxtext.trainers.pre_train import train_compile as pre_train_compile +from maxtext.training_engine import maxtext_engine +from maxtext.training_engine import maxtext_engine_compile +from maxtext.utils import maxtext_utils +import numpy as np +import pytest + +from tests.utils.test_helpers import get_test_config_path + +# training_engine imports tunix, so these tests need the post-training dependency bundle. +pytestmark = [pytest.mark.post_training] + +_REQUIRED_DEVICES = 4 +_SENTINEL = "MAXTEXT_ENGINE_XAOT_TESTS_PASSED" +_RAN = re.compile(rf"{_SENTINEL} ran=(\d+)") + + +@pytest.mark.cpu_only +def test_engine_aot_parity_on_a_four_device_cpu_mesh(): + """Runs the two parity classes below in a child process with four CPU devices. + + See the module docstring for why re-exec is the only option. The flag is *appended* rather + than defaulted so it wins over whatever a sibling module set -- XLA takes the last occurrence + of a repeated flag. + """ + env = os.environ.copy() + env["XLA_FLAGS"] = f"{env.get('XLA_FLAGS', '')} --xla_force_host_platform_device_count={_REQUIRED_DEVICES}".strip() + env["JAX_PLATFORMS"] = "cpu" + # The child imports `tests.utils.test_helpers`, which pytest puts on the path for us and a + # bare interpreter does not. + repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) + env["PYTHONPATH"] = os.pathsep.join([repo_root, env["PYTHONPATH"]]) if env.get("PYTHONPATH") else repo_root + + result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False) + + report = f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert result.returncode == 0, report + ran = _RAN.search(result.stdout) + # An exit status of 0 is also what a run that skipped everything produces. + assert ran, f"the child did not report a completed run\n{report}" + assert int(ran.group(1)) > 0, f"every test in the child skipped\n{report}" + + +def _config(**overrides) -> pyconfig.HyperParameters: + """A tiny real decoder, big enough that every sharding decision is visible in the HLO. + + `adamw` rather than `sgd`: its moments are the part of the train state the abstract path has + to invent, and so the part most likely to be wrong. + """ + argv = [ + "maxtext_engine_xaot_test.py", + get_test_config_path("base.yml"), + "model_name=default", + "run_name=engine_xaot_test", + "enable_checkpointing=False", + "convert_checkpoint_if_possible=False", + "skip_jax_distributed_system=True", + "enable_tensorboard=False", + "record_internal_nn_metrics=False", + "enable_dropout=False", + "init_weights_seed=0", + "dtype=float32", + "weight_dtype=float32", + "grad_dtype=float32", + "remat_policy=none", + "scan_layers=False", + "attention=dot_product", + "shard_mode=explicit", + f"ici_data_parallelism={_REQUIRED_DEVICES}", + "ici_fsdp_parallelism=1", + "ici_tensor_parallelism=1", + "per_device_batch_size=1", + "vocab_size=128", + "base_emb_dim=64", + "base_mlp_dim=128", + "base_num_decoder_layers=2", + "base_num_query_heads=4", + "base_num_kv_heads=4", + "head_dim=16", + "max_target_length=32", + "opt_type=adamw", + "learning_rate=1e-2", + "gradient_clipping_threshold=0.0", + "warmup_steps_fraction=0.0", + "learning_rate_final_fraction=1.0", + "gradient_accumulation_steps=1", + ] + argv.extend(f"{key}={value}" for key, value in overrides.items()) + return pyconfig.initialize(argv) + + +def _config_and_mesh(**overrides) -> tuple[pyconfig.HyperParameters, jax.sharding.Mesh]: + cfg = _config(**overrides) + return cfg, maxtext_utils.get_mesh_from_config(cfg) + + +def _batch(cfg: pyconfig.HyperParameters, seed: int) -> dict[str, np.ndarray]: + """One micro-batch shaped for `maxtext.trainers.pre_train.train.loss_fn`. + + NumPy rather than `jnp`, as a real driver's batch is: a committed array built here would + arrive replicated and be rejected by `jax.jit`'s exact `in_shardings` match. + """ + batch, seq = int(cfg.micro_batch_size_to_train_on), cfg.max_target_length + rng = np.random.default_rng(seed) + tokens = rng.integers(1, cfg.vocab_size, size=(batch, seq)).astype(np.int32) + positions = np.tile(np.arange(seq, dtype=np.int32), (batch, 1)) + segmentation = np.ones((batch, seq), dtype=np.int32) + return { + "inputs": tokens, + "targets": np.roll(tokens, -1, axis=-1), + "inputs_position": positions, + "inputs_segmentation": segmentation, + "targets_segmentation": segmentation, + } + + +def _abstract(batch: dict[str, np.ndarray]) -> dict[str, jax.ShapeDtypeStruct]: + """The same batch with the data taken out -- all an AOT compile is given.""" + return jax.tree.map(lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype), batch) + + +def _script_argv(*overrides: str) -> tuple[str, ...]: + """`main`'s argv for a model small enough to compile inside a test.""" + return ( + "", + get_test_config_path("base.yml"), + "compile_topology=v6e-4", + "compile_topology_num_slices=1", + "per_device_batch_size=1", + "max_target_length=128", + "base_emb_dim=128", + "base_mlp_dim=128", + "base_num_decoder_layers=2", + "enable_checkpointing=false", + ) + overrides + + +class _LiveKernels: + """Recompiles the kernels a running engine actually dispatched. + + An executable keeps no handle on the arguments it was compiled from, so reading a kernel's HLO + means compiling it a second time -- through the engine's own jitted wrappers, from the avals of + the arguments it was dispatched with. Avals rather than the arrays, because `_compiled_update` + donates its state. That recompilation is only sound if it reproduces the original, which + `test_recompiling_reproduces_the_kernels_training_ran` checks and the rest of the file rests on. + """ + + _ATTRIBUTES = { + "fwd_bwd": "_compiled_fwd_bwd", + "fwd_bwd_accum": "_compiled_fwd_bwd_accum", + "update": "_compiled_update", + } + + def __init__(self, engine: maxtext_engine.MaxTextTrainingEngine): + self._engine = engine + self._dispatched = {} + self.calls: dict[str, list] = {} + for name, attribute in self._ATTRIBUTES.items(): + self._dispatched[name] = getattr(engine, attribute) + setattr(engine, attribute, self._spy(name)) + + def _spy(self, name: str): + """Returns a stand-in for one kernel's dispatch handle that records its arguments and forwards.""" + + def call(*args): + self.calls.setdefault(name, []).append( + # The sharding as well as the shape: `jax.jit` keys its cache on both. + jax.tree.map( + lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype, sharding=getattr(x, "sharding", None)), + args, + ) + ) + return self._dispatched[name](*args) + + return call + + def compiled(self) -> dict[str, jax.stages.Compiled]: + """The executable behind each kernel's first dispatch. + + `jax.jit` exposes no way to compile without lowering first, so this goes through + `Lowered.compile()` -- the same two steps the engine's `compile_kernels` takes. + """ + missing = sorted(set(self._ATTRIBUTES) - set(self.calls)) + if missing: + raise AssertionError(f"these kernels were never dispatched, so they have no HLO to read: {missing}") + # The mesh and the logical axis rules have to be live while tracing, exactly as they were at + # dispatch: without the rules every `maybe_shard_with_logical` is a no-op. + with self._engine._sharding_ctx(): + return {name: self._engine._jitted_kernels[name].lower(*calls[0]).compile() for name, calls in self.calls.items()} + + +# Optimized HLO carries a source-location index -- ` ""` header lines, plus the +# `metadata={...}` and `stack_frame_id=N` referring into it -- which names where the Python was, +# not what the program does, and differs between two lowerings of the same kernel. Same +# normalization as `tests/integration/aot_identical_test.py` and `hlo_diff_test.py`. +_SOURCE_ID_LINE = re.compile(r'^\s*\d+\s+(?:"[^"]*"|\{[^}]*\})\s*$') +_METADATA = re.compile(r"metadata=\{[^}]*\}") +_STACK_FRAME = re.compile(r"stack_frame_id=\d+") + + +def _normalize(hlo: str) -> str: + """Strips source-location bookkeeping from optimized HLO, leaving the program.""" + lines = [] + for line in hlo.splitlines(): + if _SOURCE_ID_LINE.match(line): + continue + line = _METADATA.sub("metadata={}", line) + lines.append(_STACK_FRAME.sub("stack_frame_id=0", line)) + return "\n".join(lines) + + +class EngineAotParityTest(parameterized.TestCase): + """A live engine and an abstract one compile to the same bytes.""" + + __test__ = False # collected only via the subprocess entry point at the top of this file. + + def _live(self, cfg, mesh, micro_batches: int = 2): + """Runs one full optimizer step and returns the engine plus its dispatched kernels. + + Two micro-batches: the accumulating kernel is traced lazily, so a single-micro-batch run + would leave `fwd_bwd_accum` with nothing on the live side to compare against. + """ + engine = maxtext_engine.MaxTextTrainingEngine(cfg, mesh=mesh) + engine.compile(_batch(cfg, 0)) + kernels = _LiveKernels(engine) + for micro in range(micro_batches): + engine.fwd_bwd(_batch(cfg, micro)) + engine.update() + return engine, kernels + + @parameterized.named_parameters( + # The gradient all-reduce is deferred to the update, so the reduced/unreduced tags have to + # survive onto the abstract gradients too. + ("data_parallel", {}), + # Zero-1 shards the moments by moving real arrays on the live path and by restating avals + # on the abstract one. + ("zero1", {"shard_optimizer_over_data": "True"}), + ("fsdp", {"ici_data_parallelism": 1, "ici_fsdp_parallelism": _REQUIRED_DEVICES}), + # The hard case for the abstract path: JAX propagates a layout through `zeros_like` only on + # Explicit axes, so the moments' shardings have to be observed on a stand-in mesh. + ("auto_shard_mode", {"shard_mode": "auto", "ici_data_parallelism": 1, "ici_fsdp_parallelism": _REQUIRED_DEVICES}), + # Gradients narrower than the parameters, which the AOT path must read off the + # forward/backward kernel's outputs rather than assume. + ("bfloat16_grads", {"grad_dtype": "bfloat16"}), + ) + def test_aot_compiles_the_same_hlo_the_trainer_runs(self, overrides): + cfg, mesh = _config_and_mesh(**overrides) + + _, kernels = self._live(cfg, mesh) + live_compiled = kernels.compiled() + aot_compiled = maxtext_engine_compile.AbstractMaxTextEngine(cfg, mesh).compile_kernels(_abstract(_batch(cfg, 0))) + + self.assertEqual(sorted(live_compiled), sorted(aot_compiled)) + for name in live_compiled: + # Optimized HLO, which is what the cost and memory analyses are measured on and what the + # hardware actually runs. + self.assertEqual( + _normalize(live_compiled[name].as_text()), + _normalize(aot_compiled[name].as_text()), + f"{name}: the AOT executable differs from the one training ran", + ) + + def test_the_comparison_can_fail(self): + """The vacuity guard for every assertion above. + + Width rather than sequence length, which would be the obvious knob and the wrong one: + `_update_kernel` sees gradients and moments, never a sequence, so a shorter sequence leaves + its HLO byte-identical and the guard would silently test nothing on that kernel. + """ + cfg, mesh = _config_and_mesh() + wider = _config(base_mlp_dim=256) + + compiled = maxtext_engine_compile.AbstractMaxTextEngine(cfg, mesh).compile_kernels(_abstract(_batch(cfg, 0))) + compiled_wider = maxtext_engine_compile.AbstractMaxTextEngine(wider, mesh).compile_kernels( + _abstract(_batch(wider, 0)) + ) + + for name, narrow in compiled.items(): + self.assertNotEqual( + _normalize(narrow.as_text()), + _normalize(compiled_wider[name].as_text()), + f"{name}: doubling the MLP width changed nothing that survives normalization", + ) + + def test_recompiling_reproduces_the_kernels_training_ran(self): + """`engine.compile_kernels()` is not a second code path -- on a live engine it re-derives its own. + + This is what licenses `_LiveKernels`: without it every comparison here would be between two + compilations and none against training. + """ + cfg, mesh = _config_and_mesh() + + engine, kernels = self._live(cfg, mesh) + dispatched = kernels.compiled() + recompiled = engine.compile_kernels(_batch(cfg, 0)) + + for name in dispatched: + self.assertEqual(_normalize(dispatched[name].as_text()), _normalize(recompiled[name].as_text()), name) + + def test_compile_covers_the_eval_kernel_too(self): + """`compile()` is a promise that nothing after it stalls on XLA, and eval is part of a step. + + The eval kernel is compiled off the same dummy batch as the training ones, so a first + `eval_step` of that shape has to dispatch through the executable rather than replace it. + """ + cfg, mesh = _config_and_mesh() + + engine = maxtext_engine.MaxTextTrainingEngine(cfg, mesh=mesh) + engine.compile(_batch(cfg, 0)) + compiled = engine._compiled_eval + + self.assertIsInstance(compiled, jax.stages.Compiled) + with engine.eval_context(): + engine.eval_step(_batch(cfg, 1)) + self.assertIs(engine._compiled_eval, compiled, "the first eval_step recompiled instead of dispatching") + + def test_the_abstract_train_state_matches_the_live_one_leaf_for_leaf(self): + """Where a divergence would come from, stated directly rather than through the HLO. + + Same tree, paths, shapes, dtypes and shardings -- including the adamw moments and the + `count` scalars, whose eager shardings come from two different places on the two paths. + """ + cfg, mesh = _config_and_mesh() + + live = maxtext_engine.MaxTextTrainingEngine(cfg, mesh=mesh) + live.compile(_batch(cfg, 0)) + abstract = maxtext_engine_compile.AbstractMaxTextEngine(cfg, mesh) + + live_leaves = jax.tree_util.tree_flatten_with_path(nnx.split(live.state)[1])[0] + abstract_leaves = jax.tree_util.tree_flatten_with_path(nnx.split(abstract.state)[1])[0] + + self.assertNotEmpty(live_leaves) + self.assertEqual(len(live_leaves), len(abstract_leaves)) + for (live_path, live_leaf), (abstract_path, abstract_leaf) in zip(live_leaves, abstract_leaves): + where = jax.tree_util.keystr(live_path) + self.assertEqual(where, jax.tree_util.keystr(abstract_path)) + self.assertIsInstance(abstract_leaf, jax.ShapeDtypeStruct, f"{where} was materialized") + self.assertEqual(live_leaf.shape, abstract_leaf.shape, where) + self.assertEqual(live_leaf.dtype, abstract_leaf.dtype, where) + self.assertEqual(live._mesh_sharding(live_leaf), abstract._mesh_sharding(abstract_leaf), where) + + def test_every_step_after_the_first_runs_the_same_kernels(self): + """Otherwise an AOT report describes a program that runs once and is then replaced. + + `nnx.Optimizer` builds optax's `count` and its own `step` with `jnp.zeros` under no mesh, so + they reach the first update uncommitted and come back from it committed -- a second argument + signature, and a second compile of the largest kernel in the engine. `_place_state_on_mesh` + settles them up front; this pins it. + """ + cfg, mesh = _config_and_mesh() + + engine, kernels = self._live(cfg, mesh) + for micro in range(2): + engine.fwd_bwd(_batch(cfg, micro)) + engine.update() + + for name, calls in kernels.calls.items(): + self.assertLen(calls, 2, f"{name} should have been dispatched once per step") + first, second = (jax.tree_util.tree_flatten_with_path(call)[0] for call in calls) + differing = [jax.tree_util.keystr(path) for (path, a), (_, b) in zip(first, second) if a != b] + self.assertEmpty(differing, f"{name} is dispatched with a different signature on the second step") + + def test_zero1_and_the_deferred_all_reduce_engage_on_both_paths(self): + """Guards the two parameterizations that would otherwise pass by both doing nothing.""" + cfg, mesh = _config_and_mesh(shard_optimizer_over_data=True) + + # `compile()`, not a full step: both flags are set while the kernels are traced. + live = maxtext_engine.MaxTextTrainingEngine(cfg, mesh=mesh) + live.compile(_batch(cfg, 0)) + abstract = maxtext_engine_compile.AbstractMaxTextEngine(cfg, mesh) + abstract.compile_kernels(_abstract(_batch(cfg, 0))) + + for engine, label in ((live, "live"), (abstract, "abstract")): + self.assertIsNotNone(engine._zero1_params_shardings, f"Zero-1 never engaged on the {label} engine") + self.assertIsNotNone(engine._unreduced_grad_shardings, f"the deferral never engaged on the {label} engine") + + def test_the_hlo_dump_filters_match_the_names_xla_gives_the_kernels(self): + """Those filters are a claim about names XLA derives from the jitted callables. + + Nothing checks the claim at runtime: a filter that matches nothing leaves an empty dump, not + an error, so a renamed kernel would go unnoticed until someone went looking for HLO. + """ + cfg, mesh = _config_and_mesh() + compiled = maxtext_engine_compile.AbstractMaxTextEngine(cfg, mesh).compile_kernels(_abstract(_batch(cfg, 0))) + + self.assertLen(compiled, len(maxtext_engine_compile.KERNEL_NAMES)) + for kernel in compiled.values(): + module = re.search(r"^HloModule (\S+?),", kernel.as_text(), re.MULTILINE) + self.assertIsNotNone(module, "the executable has no module name to match against") + self.assertRegex(module.group(1), maxtext_engine_compile.HLO_DUMP_DEFAULTS["dump_hlo_local_module_name"]) + self.assertIn(maxtext_engine_compile.HLO_DUMP_DEFAULTS["dump_hlo_module_name"], module.group(1)) + + +class AbstractMaxTextEngineTest(absltest.TestCase): + """What an engine built without weights will and will not do.""" + + __test__ = False # collected only via the subprocess entry point at the top of this file. + + def setUp(self): + super().setUp() + self.cfg, self.mesh = _config_and_mesh() + + def test_refuses_to_run_anything(self): + """Silence here would be worse than an error: `update()` would report a step count.""" + engine = maxtext_engine_compile.AbstractMaxTextEngine(self.cfg, self.mesh) + + for operation, call in ( + ("fwd_bwd", lambda: engine.fwd_bwd(_batch(self.cfg, 0))), + ("update", engine.update), + ("save_checkpoint", lambda: engine.save_checkpoint({"step": 0})), + ("restore_checkpoint", engine.restore_checkpoint), + ): + with self.subTest(operation=operation): + with self.assertRaisesRegex(RuntimeError, "has only shapes"): + call() + + def test_requires_a_mesh(self): + with self.assertRaisesRegex(ValueError, "requires a mesh"): + maxtext_engine_compile.AbstractMaxTextEngine(self.cfg, None) + + def test_compiling_needs_something_to_compile_against(self): + engine = maxtext_engine_compile.AbstractMaxTextEngine(self.cfg, self.mesh) + + with self.assertRaisesRegex(ValueError, "needs a dummy payload"): + engine.compile_kernels(None) + + +@pytest.mark.tpu_backend +class Qwen3TopologyTest(absltest.TestCase): + """The worked example: qwen3-0.6b compiled for four v6e chips this host does not have. + + Needs libtpu, which knows a v6e's shape, but no TPU: the mesh is a topology description and + nothing is executed on it. That is the case the script exists for, and the one the parity + tests above miss, since they run on a mesh of real (if simulated) devices. + """ + + # Compiling qwen3-0.6b for a v6e-4 is the slowest thing in this file, and the two tests that + # read the report only read it, so they share one. + _compiled: dict | None = None + + def setUp(self): + super().setUp() + # `maxtext_engine_compile.main` sets this process-wide; put it back so the tests above cannot + # be reordered into a different RNG implementation. + previous = jax.config.jax_default_prng_impl + self.addCleanup(jax.config.update, "jax_default_prng_impl", previous) + + def _qwen3_config(self) -> pyconfig.HyperParameters: + """qwen3-0.6b at a sequence length short enough to compile inside a test.""" + return pyconfig.initialize( + [ + "", + get_test_config_path("base.yml"), + "model_name=qwen3-0.6b", + "run_name=engine_aot_qwen3_test", + "compile_topology=v6e-4", + "compile_topology_num_slices=1", + "ici_fsdp_parallelism=4", + "per_device_batch_size=1", + "max_target_length=512", + "attention=flash", + "enable_checkpointing=false", + ] + ) + + def _compiled_kernels(self) -> dict: + compiled = Qwen3TopologyTest._compiled + if compiled is None: + cfg = self._qwen3_config() + compiled = maxtext_engine_compile.compile_engine_kernels(cfg, pre_train_compile.get_topology_mesh(cfg)) + Qwen3TopologyTest._compiled = compiled + return compiled + + def test_compiles_all_three_kernels_with_no_weights_and_no_hardware(self): + compiled = self._compiled_kernels() + + self.assertEqual(sorted(compiled), sorted(maxtext_engine_compile.KERNEL_NAMES)) + for name, executable in compiled.items(): + memory = executable.memory_analysis() + with self.subTest(kernel=name): + # A report of zero would mean the kernel was lowered with its arguments closed over as + # constants rather than passed. + self.assertGreater(memory.argument_size_in_bytes, 0) + self.assertGreater(executable.cost_analysis()["flops"], 0) + + def test_reports_the_forward_backward_peak_rather_than_the_updates(self): + """A sanity check on the numbers, not just on their existence. + + Activations dominate a training step and the optimizer's arithmetic is elementwise, so a + report where the update needs the most scratch describes something other than this model -- + most likely a batch that never reached the kernel. + """ + compiled = self._compiled_kernels() + + fwd_bwd = compiled["fwd_bwd"].memory_analysis() + update = compiled["update"].memory_analysis() + self.assertGreater(fwd_bwd.temp_size_in_bytes, update.temp_size_in_bytes) + # Two adamw moments per parameter on top of the parameters themselves, so the update's + # arguments outweigh a single micro-batch's. + self.assertGreater(update.argument_size_in_bytes, fwd_bwd.argument_size_in_bytes) + + def test_the_whole_script_writes_one_executable_per_kernel(self): + output = os.path.join(self.create_tempdir().full_path, "engine.pickle") + + maxtext_engine_compile.main(_script_argv("run_name=engine_aot_save_test", f"compiled_trainstep_file={output}")) + + for name in maxtext_engine_compile.KERNEL_NAMES: + written = maxtext_engine_compile.kernel_save_path(output, name) + self.assertTrue(os.path.exists(written), f"{name} was not written to {written}") + self.assertGreater(os.path.getsize(written), 0, written) + + def test_a_dump_that_matched_nothing_says_which_filter_was_wrong(self): + """`upload_dump` deletes what it uploaded, and cannot be handed a directory XLA skipped. + + `jit_train_step` is `train.py`'s name for its fused step and the config default, and it is + exactly what none of the engine's three kernels are called. + """ + local_dir = os.path.join(self.create_tempdir().full_path, "xla_dump") + + with self.assertRaisesRegex(FileNotFoundError, "matched none of the engine's kernels"): + maxtext_engine_compile.main( + _script_argv( + "run_name=engine_aot_dump_test", + "dump_hlo=true", + f"dump_hlo_local_dir={local_dir}", + "dump_hlo_local_module_name=jit_train_step", + ) + ) + + def test_diloco_is_refused_rather_than_silently_reported_on(self): + """The engine has no outer step, so these numbers would describe a different run.""" + with self.assertRaisesRegex(NotImplementedError, "enable_diloco"): + maxtext_engine_compile.main(_script_argv("run_name=engine_aot_diloco_test", "enable_diloco=true")) + + +class CompileHelpersTest(absltest.TestCase): + """The pieces of the script that decide what gets compiled, and under which name.""" + + def test_the_shaped_batch_is_one_micro_batch_not_the_global_one(self): + cfg = _config(gradient_accumulation_steps=4) + shaped = maxtext_engine_compile.get_shaped_micro_batch(cfg) + + self.assertNotEmpty(shaped) + for key, aval in shaped.items(): + self.assertEqual(aval.shape[0], int(cfg.micro_batch_size_to_train_on), key) + + def test_the_shaped_batch_matches_the_shapes_a_driver_feeds(self): + """`train.loss_fn` slices its batch, so a mismatch here is a recompile, not an error.""" + cfg = _config() + shaped = maxtext_engine_compile.get_shaped_micro_batch(cfg) + driven = _batch(cfg, 0) + + self.assertContainsSubset(driven.keys(), shaped.keys()) + for key, array in driven.items(): + self.assertEqual(shaped[key].shape, array.shape, key) + self.assertEqual(shaped[key].dtype, array.dtype, key) + + def test_an_abstract_batch_is_traced_rather_than_closed_over(self): + """`_is_jax_dynamic` decides this, and a batch it classifies static is not passed at all.""" + batch = _batch(_config(), 0) + dynamic, static = maxtext_engine._split_static_and_dynamic(_abstract(batch)) + + self.assertEmpty(static) + self.assertEqual(sorted(dynamic), sorted(batch)) + + def test_each_kernel_is_saved_to_its_own_file(self): + paths = [ + maxtext_engine_compile.kernel_save_path("/tmp/engine.pickle", name) + for name in maxtext_engine_compile.KERNEL_NAMES + ] + + self.assertEqual(len(set(paths)), len(maxtext_engine_compile.KERNEL_NAMES)) + self.assertEqual(paths[0], "/tmp/engine_fwd_bwd.pickle") + + def test_a_requested_dump_is_pointed_at_the_kernels(self): + argv = maxtext_engine_compile.with_engine_hlo_dump_defaults(["", "base.yml", "dump_hlo=True"]) + + self.assertEqual(argv[:3], ["", "base.yml", "dump_hlo=True"]) + self.assertContainsSubset([f"{key}={value}" for key, value in maxtext_engine_compile.HLO_DUMP_DEFAULTS.items()], argv) + + def test_a_run_that_asked_for_no_dump_is_left_alone(self): + """The regex reaches `XLA_FLAGS` either way, so widening it here would dump on every run.""" + argv = ["", "base.yml", "compile_topology=v6e-4"] + + self.assertEqual(maxtext_engine_compile.with_engine_hlo_dump_defaults(argv), argv) + + def test_an_explicit_filter_is_not_overridden(self): + argv = maxtext_engine_compile.with_engine_hlo_dump_defaults( + ["", "base.yml", "dump_hlo=true", "dump_hlo_module_name=first"] + ) + + self.assertIn("dump_hlo_module_name=first", argv) + self.assertNotIn("dump_hlo_module_name=kernel", argv) + self.assertIn("dump_hlo_local_module_name=jit_.*kernel", argv) + + +_SUITE = (EngineAotParityTest, AbstractMaxTextEngineTest) + + +if __name__ == "__main__": + if jax.device_count() < _REQUIRED_DEVICES: + raise SystemExit( + f"needs {_REQUIRED_DEVICES} devices, got {jax.device_count()}; run this through pytest, which sets " + f"XLA_FLAGS=--xla_force_host_platform_device_count={_REQUIRED_DEVICES}" + ) + _loader = unittest.defaultTestLoader + _result = unittest.TextTestRunner(verbosity=2).run( + unittest.TestSuite(_loader.loadTestsFromTestCase(cls) for cls in _SUITE) + ) + if not _result.wasSuccessful(): + sys.exit(1) + print(f"{_SENTINEL} ran={_result.testsRun - len(_result.skipped)}")