From 3b420fd9b2aaac173566e80fd08a0db2eef47c86 Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Wed, 2 Sep 2026 00:54:10 +0000 Subject: [PATCH 1/3] Match Tunix peft_trainer_v2 performance in MaxTextTrainingEngine `MaxTextTrainingEngine` was slower than Tunix's `PeftTrainer` v2 on the same model, the same loss and the same optimizer, and the gap was entirely host-side. Three changes to the step path close it, plus two correctness fixes the measurement turned up. Performance: - Trace the accumulation and update kernels under `nn_partitioning.axis_rules`, so gradient accumulation is fused into the compiled program instead of running as eager `jax.tree.map(add)` over every leaf, every micro-batch. - Carry a pure `nnx.State` mirror across steps in place of a per-step `nnx.split` of the module graph. The split was two traversals per step, scaling with node count -- 21.3 ms on an unscanned qwen3-0.6b's 310 parameter leaves -- and it allocated a large short-lived object graph twice a step, whose GC pauses landed on whichever step was unlucky (a 464 ms worst step against a 161.6 ms median). - Defer the metrics write until after the next dispatch, so its device-to-host read overlaps live work rather than blocking on it. Correctness: - Accumulate gradients unreduced and divide once by the summed denominator, rather than averaging per-micro-batch means. The two agree only when every micro-batch has the same token count; on ragged batches the mean-of-means gradient is simply the wrong one. The denominator rides along in checkpoint metadata so a mid-step resume can finish the division, and is backfilled from the restored per-micro-batch losses for checkpoints written before it was tracked. - Compute the gradient norm on every update, not only under `skip_step_on_spikes`. `learning/grad_norm` is a metric `trainers/pre_train/train.py` reports each step and `metrics._METRICS_TO_LOG` already lists, but with spike-skipping off -- base.yml's default -- it read NaN. Tunix computes the same quantity unconditionally in `peft_trainer_v2._update_step`. Taken after clipping, in float32, it now matches Tunix bit for bit: 59.77279281616211 at GA=1 and 18.25168800354004 at GA=8. The norm also gives the inflight throttler a scalar to block on, so it no longer holds the whole train state alive between steps -- three parameter trees pinned per queued entry, whose buffers a later donating update would delete out from under `jax.block_until_ready`. Measured on qwen3-0.6b, batch 8 x seq 1024, f32, `optax.sgd(1e-5)`, no clipping, as an A/B on these files alone. Per optimizer step, medians of 19 steps after warmup: at GA=1 on 8 x v7x, 161.6 -> 89.8 ms (1.80x); at GA=8 on 4 x v6e, 3458.9 -> 589.6 ms (5.87x), against 1225.9 ms for the same MaxText model under `PeftTrainer` (2.08x, at 96% device utilization against 50%). TPU-busy time is identical across the A/B to one part in 340,000 and the accumulation kernel keeps the same XLA module hash, so all of it is host time. Tests: `tests/post_training/unit/maxtext_engine_test.py` covers the cached pure-state path end to end on a model with non-`Param` state, the unreduced accumulation and its denominator, the denominator's checkpoint round-trip, the throttler's scalar handle, and the gradient norm with spike-skipping off. 59 passed across the engine unit tests. --- .../training_engine/inflight_throttler.py | 36 +- src/maxtext/training_engine/maxtext_engine.py | 477 +++++++++++++++--- src/maxtext/training_engine/metrics.py | 58 ++- .../post_training/unit/maxtext_engine_test.py | 156 +++++- 4 files changed, 643 insertions(+), 84 deletions(-) diff --git a/src/maxtext/training_engine/inflight_throttler.py b/src/maxtext/training_engine/inflight_throttler.py index dac7063f05..949336fbc9 100644 --- a/src/maxtext/training_engine/inflight_throttler.py +++ b/src/maxtext/training_engine/inflight_throttler.py @@ -34,19 +34,43 @@ def __init__(self, config: pyconfig.HyperParameters): """ self._inflight_queue = queue.Queue[Any](maxsize=config.max_inflight_computations) self._metrics_logger = metrics_module.MetricsLogger(config=config) + # Popped by `wait_for_next` but not yet written; see `_flush_pending_metrics`. + self._pending_metrics: abstract_engine.MetricsBuffer | None = None def add_computation(self, computation: Any, metrics: abstract_engine.MetricsBuffer | None) -> None: """Adds an active on-device computation to the queue.""" self._inflight_queue.put((jax.tree.leaves(computation), metrics)) + # The caller has just dispatched, so the device has work queued behind this point and the + # blocking read inside `write_metrics` overlaps it instead of running against an idle + # device. This is the whole reason the write is deferred rather than done in place. + self._flush_pending_metrics() + + def _flush_pending_metrics(self) -> None: + """Writes the buffer stashed by the last `wait_for_next`, if any.""" + if self._pending_metrics is None: + return + metrics, self._pending_metrics = self._pending_metrics, None + self._metrics_logger.write_metrics(metrics) def wait_for_next(self) -> None: - """If the limit is reached, wait for the next computation to finish.""" + """If the limit is reached, wait for the next computation to finish. + + Blocks, but does not log. `write_metrics` reduces each `WeightedMetric` on device and then + pulls the result to host with `np.asarray`, and those reduction ops are dispatched behind + whatever is already in the device's queue -- so doing it here, before the caller dispatches + the step it just made room for, stalls the host on the full backlog with nothing new + running. Deferring to the following `add_computation` costs one dispatch of staleness in + the log and hands the same work a busy device to hide behind. Metrics carry their own step + id (`MetricsBuffer.id`), so nothing downstream can tell the difference. + """ if self._inflight_queue.full(): computation, metrics = self._inflight_queue.get() jax.block_until_ready(computation) - # Write metrics for the completed computation. if metrics is not None: - self._metrics_logger.write_metrics(metrics) + # Never hold two: the engine only attaches metrics to one of its two computations per + # step, but a caller that attached them to both would otherwise silently lose a buffer. + self._flush_pending_metrics() + self._pending_metrics = metrics def wait_for_all(self) -> None: """Wait for all inflight computations to finish and log their metrics.""" @@ -55,7 +79,11 @@ def wait_for_all(self) -> None: jax.block_until_ready(computation) # Write metrics for the completed computation. if metrics is not None: - self._metrics_logger.write_metrics(metrics) + self._flush_pending_metrics() + self._pending_metrics = metrics + # Draining is the one place that must not leave a write outstanding: callers use it to + # reach a quiescent state before checkpointing or shutting down. + self._flush_pending_metrics() def cleanup(self) -> None: """Closes the underlying metrics logger and releases resources.""" diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py index 1b3b06a62d..8ae81663ea 100644 --- a/src/maxtext/training_engine/maxtext_engine.py +++ b/src/maxtext/training_engine/maxtext_engine.py @@ -21,12 +21,14 @@ from __future__ import annotations from collections.abc import Callable, Mapping +import contextlib import dataclasses import os from typing import Any from absl import logging from flax import nnx +from flax.linen import partitioning as nn_partitioning from flax.traverse_util import flatten_dict from flax.traverse_util import unflatten_dict import jax @@ -53,6 +55,17 @@ # same situation. Real buffers are identified by their train step, so this cannot collide. EMPTY_METRICS_BUFFER_ID = -1 +# Where `nnx.split(TrainStateNNX(...))` puts the model's own state. See +# `_check_pure_state_reusable` for why the engine verifies this rather than assuming it. +_MODEL_STATE_KEY = "model" + +_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 " + "`nnx.split` calls cost ~92 ms per step on an unrolled 28-layer qwen3-0.6b, against " + "~2 ms for the pure-state equivalent. Logged once per engine instance." +) + def _is_jax_dynamic(value: Any) -> bool: """Returns True if `value` can cross a `jax.jit` boundary as a traced argument. @@ -358,7 +371,19 @@ def __init__( else: self._model = model_or_model_mesh_pair self._state: Any = None + # Pure-pytree mirror of the model and the train state, carried across steps so the step + # path never re-walks the module graph. `None` means "not cached", which is also how the + # fast path is switched off; see `_refresh_pure_state`. + self._params_pure: Any = None + self._rest_pure: Any = None + self._state_pure: Any = None + self._pure_state_warned: bool = False self._accumulated_grads: Any = None + # Sum of the per-micro-batch loss denominators behind `_accumulated_grads`. Tracked + # alongside them because the gradients are accumulated unreduced: this is the divisor + # `update()` applies once, and it is not `micro_step_count` unless every micro-batch + # happened to carry the same token count. + self._accumulated_denominator: Any = None self._micro_step_count = 0 # Set when this run resumed from an intra-step checkpoint, cleared once the step it # resumed into completes and its finished state has been checkpointed. @@ -391,8 +416,10 @@ def model(self, new_model: Any) -> None: self._model = new_model self._compiled = False self._compiled_fwd_bwd = None + self._compiled_fwd_bwd_accum = None self._compiled_update = None self._model_graphdef = None + self._invalidate_pure_state() @property def optimizer(self) -> Any: @@ -405,8 +432,10 @@ def optimizer(self, new_optimizer: Any) -> None: self._optimizer = new_optimizer self._compiled = False self._compiled_fwd_bwd = None + self._compiled_fwd_bwd_accum = None self._compiled_update = None self._state_graphdef = None + self._invalidate_pure_state() @property def train_step(self) -> int: @@ -431,8 +460,10 @@ def state(self, new_state: Any) -> None: self._state = new_state self._compiled = False self._compiled_fwd_bwd = None + self._compiled_fwd_bwd_accum = None self._compiled_update = None self._state_graphdef = None + self._invalidate_pure_state() @property def micro_step_count(self) -> int: @@ -484,8 +515,205 @@ def with_gen_model_input_fn(self, gen_model_input_fn: Callable[[Any], dict[str, self._compiled = False return self - def _fwd_bwd_kernel(self, params, rest, batch): - """Executes a single forward and backward pass to compute gradients.""" + @contextlib.contextmanager + def _sharding_ctx(self): + """Activates the mesh and logical axis rules the MaxText layers are written against. + + MaxText modules place their own sharding constraints through + `nn_partitioning.get_axis_rules()` (`sharding.maybe_shard_with_logical` and friends). + Those rules live in a context variable, so a kernel traced outside this context sees an + empty rule set: every logical constraint silently becomes a no-op and XLA is left to + guess the partitioning for activations and gradients. It guesses badly -- on + llama3.1-8b/fsdp=8 the same fwd/bwd measured 1012 ms untraced-in-context against 581 ms + inside it, with no numerical difference. `train.py` wraps its own `jax.jit` the same way + (`with jax.set_mesh(mesh), mesh, nn_partitioning.axis_rules(config.logical_axis_rules)`), + which is why the standalone trainer never hit this. Its middle `mesh` is left out here: + `Mesh.__enter__` is deprecated and `jax.set_mesh` already covers it. + + Entered around the *call*, not around `jax.jit(...)`: jit is lazy, so the rules must be + live when the first call triggers tracing. + """ + if self._mesh is None: + yield + return + with jax.set_mesh(self._mesh), nn_partitioning.axis_rules(self._config.logical_axis_rules): + yield + + def _invalidate_pure_state(self) -> None: + """Forgets the cached pure state, so the next step re-reads it from the NNX objects. + + Called from the `model`/`optimizer`/`state` setters and after a checkpoint restore -- + the three ways the live NNX variables can be replaced behind the engine's back. It is + not needed on the step path: the step path is what *produces* the cached values. + """ + self._params_pure = None + self._rest_pure = None + self._state_pure = None + + def _disable_pure_state(self, reason: str) -> None: + """Falls back to re-splitting the module graph on every step, saying so once.""" + self._invalidate_pure_state() + if not self._pure_state_warned: + self._pure_state_warned = True + logging.warning(_PURE_STATE_FALLBACK_WARNING, reason) + + @staticmethod + def _with_model_state(state_pure: Any, model_pure: Any) -> Any: + """Returns `state_pure` with its model subtree replaced by `model_pure`. + + Goes through `raw_mapping` rather than `{**state_pure}` because `nnx.State` stores its + children as plain dicts and only wraps them in a `State` on `__getitem__`. Rebuilding + from the wrapped views produces a tree that is equal key for key and leaf for leaf but + is a *different pytree*, one `State` node deeper at every level -- which `jax.jit` + rejects, at the call site, as an `in_shardings` prefix mismatch that names neither this + function nor the reason. + """ + return nnx.State({**state_pure.raw_mapping, _MODEL_STATE_KEY: model_pure.raw_mapping}) + + def _check_pure_state_reusable(self, state_pure: Any, params_pure: Any, rest_pure: Any) -> str | None: + """Returns why the pure state cannot be carried across steps, or None if it can. + + The step path rebuilds the update kernel's `state_pure` argument by dropping the model's + parameter and non-parameter state back into `state_pure["model"]`, and re-derives the + next step's parameters from the kernel's output the same way. That is only the same + value `nnx.split` would have produced if NNX puts the model's state there, and puts it + there exactly once. + + It does for `TrainStateNNX`: `__init__` assigns `self.model` before `self.optimizer`, so + the model is flattened first and the optimizer's reference to the same module becomes a + graph reference rather than a second copy of the weights -- the pure dict is + `{"model": ..., "optimizer": {"opt_state": ..., "step": ...}}`. But that is a property of + a class this engine does not own, and `engine.state` is a public setter that will accept + anything, so it is checked rather than assumed. + + The check runs the real reconstruction and compares the result against what `nnx.split` + produced, rather than testing some proxy for it. A pytree that differs from the one the + kernels were compiled against fails at the `jax.jit` call site with an `in_shardings` + prefix mismatch, which is a hard error to read back to its cause; catching it here costs + one comparison per compile. + + Returns: + `None` when the fast path is safe, else a short phrase naming what did not line up. + """ + if not isinstance(state_pure, nnx.State) or _MODEL_STATE_KEY not in state_pure: + return f"the train state's pure form has no {_MODEL_STATE_KEY!r} entry" + if not hasattr(state_pure, "raw_mapping"): + return "this version of flax.nnx.State does not expose raw_mapping" + rebuilt = self._with_model_state(state_pure, nnx.merge_state(params_pure, rest_pure)) + if jax.tree.structure(rebuilt) != jax.tree.structure(state_pure): + # The likely cause is a train state that holds the model somewhere other than + # `.model`, or that holds a *second* module sharing its variables, which makes the + # first-flattened copy the only real one and this one a reference. + return f"state[{_MODEL_STATE_KEY!r}] is not the model's own state" + return None + + def _refresh_pure_state(self) -> None: + """Re-reads the model and train state as pure `nnx.State`, and caches both. + + Runs once per compile rather than once per step, which is the point. `nnx.split` walks + the entire module graph: on an unrolled 28-layer qwen3-0.6b it costs 51.6 ms for the + model and 40.8 ms for the train state, so the two calls the step path used to make were + 92 ms of a 283 ms step -- more than the 82 ms the step spent on the TPU. The same + partition applied to an already-flat state (`nnx.split_state`) costs 0.84 ms, because it + walks 490 leaves instead of 1756 graph nodes. + + What this does *not* change is when results are published: `fwd_bwd` and `update` still + `nnx.update` the live NNX objects at exactly the points they always did, so `self.model` + and `self.state` are never stale and nothing outside the step path has to know the cache + exists. Tunix v2 gets the same saving a different way, by holding the graph inside + `nnx.cached_partial` (`peft_trainer_v2.maybe_cache_and_partial`); that hook is specific + to `nnx.jit`, and these kernels are plain `jax.jit` over pure state. + """ + if self._state is None: + self._state = train_state_nnx.TrainStateNNX(self._model, self._optimizer) + model = getattr(self._state, _MODEL_STATE_KEY, self._model) + self._state_graphdef, state_pure = nnx.split(self._state) + self._model_graphdef, params_pure, rest_pure = nnx.split(model, nnx.Param, ...) + + reason = self._check_pure_state_reusable(state_pure, params_pure, rest_pure) + if reason is not None: + self._disable_pure_state(reason) + return + self._params_pure, self._rest_pure, self._state_pure = params_pure, rest_pure, state_pure + + def _read_model_pure(self, model: Any) -> tuple[Any, Any]: + """Returns the model's `(params, rest)` pure state, from the cache when it is live.""" + if self._params_pure is not None: + return self._params_pure, self._rest_pure + self._model_graphdef, params, rest = nnx.split(model, nnx.Param, ...) + return params, rest + + def _read_state_pure(self) -> Any: + """Returns the train state's pure form, from the cache when it is live.""" + if self._state_pure is not None: + return self._state_pure + self._state_graphdef, state_pure = nnx.split(self._state) + return state_pure + + def _publish_model_rest(self, new_rest: Any) -> None: + """Folds a fwd/bwd's updated non-parameter state into the cached train state. + + Required, not an optimization: `update()` passes the whole train state to its kernel, and + before this cache existed the fresh `rest` reached it via `nnx.update(model, new_rest)` + landing in the variables that the following `nnx.split(state)` then read back. Skipping + the fold here would hand the update kernel the *previous* micro-batch's non-parameter + state -- for a model with RNG counters or batch statistics, a silent regression. + + The structure is checked, not assumed. `_fwd_bwd_kernel` re-splits the model it merged, so + anything the forward pass `sow`s as an `nnx.Intermediate` comes back as an extra entry in + `new_rest` -- `record_max_logits`, `distill_beta > 0` and multi-token prediction all do + this. Adopting a wider `rest` would leave the cache disagreeing with the `rest_shardings` + the kernel was compiled against, and the mismatch would surface one call later as a + `jax.jit` in_shardings prefix error naming neither the sow site nor this method. + """ + if self._params_pure is None: + return + if jax.tree.structure(new_rest) != jax.tree.structure(self._rest_pure): + self._disable_pure_state("fwd_bwd returned a wider non-parameter state than the model was split into") + return + self._rest_pure = new_rest + self._state_pure = self._with_model_state(self._state_pure, nnx.merge_state(self._params_pure, new_rest)) + + def _publish_state(self, new_state_pure: Any) -> None: + """Adopts the update kernel's output as the cached state, re-deriving `(params, rest)`. + + The re-derivation is checked rather than trusted. It reads `.type` off the state's + leaves, which survive the round trip through `jax.jit` as part of the treedef -- but if + a future NNX ever flattened them differently, the partition would come back wrong and + the next `nnx.merge(self._model_graphdef, params, rest)` would fail inside a traced + kernel, where the error names a pytree mismatch and not this line. Comparing the treedef + here costs ~1 ms and turns that into a fallback plus one warning. Correctness does not + hinge on it either way: `update()` has already written `new_state_pure` into the live NNX + objects by the time this runs, so dropping the cache loses speed and nothing else. + """ + if self._params_pure is None: + return + if not isinstance(new_state_pure, nnx.State) or _MODEL_STATE_KEY not in new_state_pure: + self._disable_pure_state("the update kernel returned a state with no model entry") + return + params_pure, rest_pure = nnx.split_state(new_state_pure[_MODEL_STATE_KEY], nnx.Param, ...) + if jax.tree.structure(params_pure) != jax.tree.structure(self._params_pure): + self._disable_pure_state("the update kernel's output does not partition into the same parameters") + return + self._params_pure, self._rest_pure, self._state_pure = params_pure, rest_pure, new_state_pure + + def _fwd_bwd_kernel(self, params, rest, batch, acc_grads=None, acc_denom=None): + """Executes a single forward and backward pass and folds the result into the accumulator. + + Args: + params: Pure `nnx.Param` state to differentiate against. + rest: The model's remaining (non-parameter) pure state. + batch: Loss-function inputs for this micro-batch. + acc_grads: Gradients accumulated over earlier micro-batches of this update, or None on + the first micro-batch of an update. Passing None is what lets the first micro-batch + skip allocating -- and zeroing -- a parameter-sized buffer that would be overwritten + before anything read it. + acc_denom: Denominator accumulated alongside `acc_grads`, or None with it. + + Returns: + `(primary_loss, aux_metrics, new_rest, acc_grads, acc_denom)`, where the last two are + this micro-batch folded into the running totals. + """ loss_callable = self._loss_fn if self._loss_fn is not None else maxtext_train.loss_fn def diff_wrapper(p, r, b): @@ -546,13 +774,10 @@ def diff_wrapper(p, r, b): ) grad_func = jax.value_and_grad(diff_wrapper, argnums=0, has_aux=True) - # Every non-raising branch of `diff_wrapper` builds a LossOutput, so `loss_out` is - # always one and the gradient scaling below is unconditional. The value returned by - # `value_and_grad` is the unreduced sum that was differentiated, which - # `loss_out.primary_loss` already carries, so it is discarded here. + # Every non-raising branch of `diff_wrapper` builds a LossOutput, so `loss_out` is always + # one. The value returned by `value_and_grad` is the unreduced sum that was + # differentiated, which `loss_out.primary_loss` already carries, so it is discarded here. (_, (loss_out, new_rest)), micro_grads = grad_func(params, rest, batch) - scale = loss_out.primary_loss.compute_scale() - micro_grads = jax.tree.map(lambda g: g * scale, micro_grads) micro_grads = jax.tree.map( lambda x: ( @@ -563,24 +788,52 @@ def diff_wrapper(p, r, b): micro_grads, ) - return loss_out.primary_loss, loss_out.aux_metrics, new_rest, micro_grads + # The gradients accumulated here are the UNREDUCED ones -- d/dparam of the summed loss, + # with no `1/denominator` applied. `_update_kernel` divides the total by the total + # denominator, so the optimizer sees the global weighted mean `sum(grads)/sum(denom)` + # rather than a mean of per-micro-batch means. The two agree only when every micro-batch + # carries the same token count; under sequence packing or a ragged RL rollout they do + # not, and the mean-of-means silently overweights short micro-batches. It is also one + # fewer full pass over the gradient tree per micro-batch. This mirrors what MaxText's own + # pre-train path already does (`gradient_accumulation.py`: accumulate `xent_sum`, divide + # once by the summed `total_weights`). + denominator = loss_out.primary_loss.denominator.astype(jnp.float32) + if acc_grads is None: + return loss_out.primary_loss, loss_out.aux_metrics, new_rest, micro_grads, denominator + acc_grads = jax.tree.map(jnp.add, acc_grads, micro_grads) + return loss_out.primary_loss, loss_out.aux_metrics, new_rest, acc_grads, acc_denom + denominator + + def _update_kernel(self, state_pure, accumulated_grads, accumulated_denominator, mean_loss): + """Applies accumulated gradients to update the NNX model state. - def _update_kernel(self, state_pure, accumulated_grads, micro_step_count, mean_loss): - """Applies accumulated gradients to update the NNX model state.""" + Returns: + `(new_state_pure, grad_norm, is_skipped)`. `grad_norm` doubles as the throttler's + handle on this update; see the `add_computation` call in `update()`. + """ grad_norm = None is_skipped_val = None if state_pure is not None: - if micro_step_count <= 1: - grads = accumulated_grads - else: - grads = jax.tree.map( - lambda g: g / micro_step_count, - accumulated_grads, - ) - grad_norm = max_utils.l2norm_pytree(grads) + # `accumulated_grads` holds sum_i grad(unreduced_sum_i) and `accumulated_denominator` + # holds sum_i denominator_i, so this one division is the whole normalization. A zero + # total means every micro-batch was empty; yield zeros rather than a NaN, matching + # `WeightedMetric.compute_scale()` and the `has_weights` guard in + # `gradient_accumulation.py`. + has_weights = accumulated_denominator > 0 + safe_denominator = jnp.where(has_weights, accumulated_denominator, 1.0) + grads = jax.tree.map( + lambda g: jnp.where(has_weights, g / safe_denominator.astype(g.dtype), jnp.zeros_like(g)), + accumulated_grads, + ) + # Before clipping, which is where 202a89ab8 put it and where Tunix's own + # `optax.global_norm` sits (`peft_trainer_v2._update_step`): its clipping, if any, is a + # link in the optax chain that runs after. Note this is `raw_grad_norm` in + # `train.py`'s vocabulary rather than its `learning/grad_norm`; with clipping off -- + # base.yml's default, and Tunix, which never clips -- the two coincide. In float32 + # whatever `grad_dtype` is, because a sum of squares over bf16 leaves overflows on + # production-size models, and it is the norm the throttler blocks on. + grad_norm = max_utils.l2norm_pytree(jax.tree.map(lambda g: g.astype(jnp.float32), grads)) if self._config.gradient_clipping_threshold > 0: grads = maxtext_utils.apply_gradient_clipping(grads, None, self._config.gradient_clipping_threshold) - local_state = nnx.merge(self._state_graphdef, state_pure, copy=True) if hasattr(local_state, "apply_gradients"): if self._config.skip_step_on_spikes: @@ -710,43 +963,85 @@ def _compile_for_batch(self, dynamic_batch: Any, static_batch: dict[str, Any]) - `static_batch` is closed over rather than passed, so non-array loss arguments (Tunix's `algo_config`, `pad_id`, `eos_id`) never reach the jit boundary. """ - if self._state is None: - self._state = train_state_nnx.TrainStateNNX(self._model, self._optimizer) - - self._state_graphdef, state_pure = nnx.split(self._state) - self._model_graphdef, params_pure, rest_pure = nnx.split(self._model, nnx.Param, ...) - - def kernel(params, rest, dynamic): + # Re-reads both graphs and both pure states, and is the only place on the step path that + # does: everything after this is maintained as plain pytrees until something invalidates + # the cache. A recompile is exactly when the graph may legitimately have changed shape, + # so it is also the right moment to re-derive the shardings below. + self._refresh_pure_state() + state_pure = self._read_state_pure() + params_pure, rest_pure = self._read_model_pure(getattr(self._state, _MODEL_STATE_KEY, self._model)) + + def first_kernel(params, rest, dynamic): batch = {**dynamic, **static_batch} if isinstance(dynamic, dict) else dynamic return self._fwd_bwd_kernel(params, rest, batch) + def accum_kernel(params, rest, dynamic, acc_grads, acc_denom): + batch = {**dynamic, **static_batch} if isinstance(dynamic, dict) else dynamic + return self._fwd_bwd_kernel(params, rest, batch, acc_grads, acc_denom) + if self._mesh is not None: + replicated = jax.sharding.NamedSharding(self._mesh, jax.sharding.PartitionSpec()) state_mesh_shardings = jax.tree.map(self._mesh_sharding, state_pure) params_shardings = jax.tree.map(self._mesh_sharding, params_pure) rest_shardings = jax.tree.map(self._mesh_sharding, rest_pure) - fwd_bwd_in_shardings = (params_shardings, rest_shardings, self._batch_data_shardings(dynamic_batch)) - fwd_bwd_out_shardings = (None, None, rest_shardings, params_shardings) - update_in_shardings = (state_mesh_shardings, params_shardings, None) + batch_shardings = self._batch_data_shardings(dynamic_batch) + first_in_shardings = (params_shardings, rest_shardings, batch_shardings) + accum_in_shardings = first_in_shardings + (params_shardings, replicated) + fwd_bwd_out_shardings = (None, None, rest_shardings, params_shardings, replicated) + update_in_shardings = (state_mesh_shardings, params_shardings, replicated, None) update_out_shardings = (state_mesh_shardings, None, None) else: - fwd_bwd_in_shardings = None + first_in_shardings = None + accum_in_shardings = None fwd_bwd_out_shardings = None update_in_shardings = None update_out_shardings = None - # 1. JIT Compile Micro FWD/BWD Pass + # 1. JIT Compile Micro FWD/BWD Pass. + # + # Two kernels rather than one. The first micro-batch of an update has nothing to add to, + # so it returns its own gradients and the engine adopts them as the accumulator; every + # later micro-batch folds into that buffer in place. Tunix v2 calls the same split + # "non-persistent vs persistent" mode. It buys two things: the first micro-batch never + # allocates or zeroes a parameter-sized buffer that would be overwritten unread, and the + # accumulating kernel can *donate* the accumulator so XLA writes the sum back into the + # incoming buffer. Before this, accumulation happened in Python + # (`jax.tree.map(jnp.add, ...)`), which materialized the micro-batch gradients as a + # program output *and* allocated a fresh sum -- two extra parameter-sized buffers live + # at once. `jax.jit` is lazy, so the accumulating kernel costs nothing to compile when + # every update consumes a single micro-batch. + # + # `params` is deliberately NOT donated: `micro_grads` has the same shape, dtype and + # sharding, and JAX matches donations by shard-shape rather than by position + # (`jax/_src/interpreters/mlir.py:_set_up_aliases`), so donating it would alias the + # weights straight into the gradient output and destroy them. self._compiled_fwd_bwd = jax.jit( - kernel, - in_shardings=fwd_bwd_in_shardings, + first_kernel, + in_shardings=first_in_shardings, out_shardings=fwd_bwd_out_shardings, ) + self._compiled_fwd_bwd_accum = jax.jit( + accum_kernel, + in_shardings=accum_in_shardings, + out_shardings=fwd_bwd_out_shardings, + donate_argnums=(3, 4), + ) - # 2. JIT Compile Optimizer Update Pass + # 2. JIT Compile Optimizer Update Pass. + # + # `state_pure` (parameters plus optimizer slots) is donated, exactly as + # `maxtext_utils.get_functional_train_with_signature` does for the standalone trainer + # (`donate_argnums = 0`): it is dead the moment the kernel returns, since the engine + # rebinds the state from the output, so aliasing saves holding a second copy of the whole + # train state. The accumulated gradients are *not* donated. Every parameter-shaped output + # is already claimed by the incoming state (weights and optimizer slots alike), so a + # gradient donation has nothing left to alias to and JAX would only warn about it; the + # buffers are freed by `update()` dropping its reference anyway. self._compiled_update = jax.jit( self._update_kernel, in_shardings=update_in_shardings, out_shardings=update_out_shardings, - static_argnums=(2,), + donate_argnums=(0,), ) self._compiled_signature = _batch_signature(dynamic_batch, static_batch) self._compiled = True @@ -800,8 +1095,7 @@ def fwd_bwd(self, payload: abstract_engine.TrainerPayload, **kwargs: Any) -> Non if self._state is None: self._state = train_state_nnx.TrainStateNNX(self._model, self._optimizer) - model = getattr(self._state, "model", self._model) - self._model_graphdef, params, rest = nnx.split(model, nnx.Param, ...) + model = getattr(self._state, _MODEL_STATE_KEY, self._model) if self._compile_requested: dynamic_batch, static_batch = _split_static_and_dynamic(batch) @@ -812,10 +1106,27 @@ def fwd_bwd(self, payload: abstract_engine.TrainerPayload, **kwargs: Any) -> Non signature = _batch_signature(dynamic_batch, static_batch) if not self._compiled or self._needs_recompile(signature): self._compile_for_batch(dynamic_batch, static_batch) - loss, aux, new_rest, micro_grads = self._compiled_fwd_bwd(params, rest, dynamic_batch) + # Read after any recompile, not before: `_compile_for_batch` refreshes the cache, and + # reading first would hand the new kernel a pure state split against the old graph. + params, rest = self._read_model_pure(model) + with self._sharding_ctx(): + if self._accumulated_grads is None: + loss, aux, new_rest, acc_grads, acc_denom = self._compiled_fwd_bwd(params, rest, dynamic_batch) + else: + # `self._accumulated_grads` and `self._accumulated_denominator` are donated by this + # call, so their buffers are gone once it returns. Rebinding both from the outputs + # below is what keeps that safe -- nothing else holds a reference to either. + loss, aux, new_rest, acc_grads, acc_denom = self._compiled_fwd_bwd_accum( + params, rest, dynamic_batch, self._accumulated_grads, self._accumulated_denominator + ) else: - loss, aux, new_rest, micro_grads = self._fwd_bwd_kernel(params, rest, batch) + params, rest = self._read_model_pure(model) + with self._sharding_ctx(): + loss, aux, new_rest, acc_grads, acc_denom = self._fwd_bwd_kernel( + params, rest, batch, self._accumulated_grads, self._accumulated_denominator + ) nnx.update(model, new_rest) + self._publish_model_rest(new_rest) # Don't add metrics to the throttler queue because metrics are logged after # the update step. @@ -831,10 +1142,9 @@ def fwd_bwd(self, payload: abstract_engine.TrainerPayload, **kwargs: Any) -> Non self.record_metrics(key, value) self._cached_losses.append(loss) - if self._accumulated_grads is None: - self._accumulated_grads = micro_grads - else: - self._accumulated_grads = jax.tree.map(jnp.add, self._accumulated_grads, micro_grads) + # Accumulation happened inside the kernel; there is nothing to add here. + self._accumulated_grads = acc_grads + self._accumulated_denominator = acc_denom self._micro_step_count += 1 def update(self, **kwargs: Any) -> int: @@ -859,40 +1169,62 @@ def update(self, **kwargs: Any) -> int: # Wait for previous computations to finish before dispatching the update step to TPU. self._throttler.wait_for_next() - # TODO(mazumdera): The logic below should be pre-compiled. if self._state is None: self._state = train_state_nnx.TrainStateNNX(self._model, self._optimizer) - self._state_graphdef, state_pure = nnx.split(self._state) - - if self._cached_losses: + state_pure = self._read_state_pure() + + # `_update_kernel` reads `mean_loss` only inside its `skip_step_on_spikes` branch, and that + # flag is read off `self._config` at trace time, so with spike-skipping off -- the base.yml + # default -- the value is discarded. Computing it is not free: `WeightedMetric.compute()` is + # seven eager XLA launches (the eps and min_denom clamps plus a safe divide), so this was + # seven dispatches per step feeding an argument the executable does not contain. `None` is + # an empty pytree, which is what `update_in_shardings` already declares for this position. + if not self._config.skip_step_on_spikes: + mean_loss = None + elif self._cached_losses: loss_values = [l.compute() if isinstance(l, abstract_engine.WeightedMetric) else l for l in self._cached_losses] mean_loss = jnp.mean(jnp.stack(loss_values)) if len(loss_values) > 1 else loss_values[0] else: mean_loss = jnp.array(0.0) - if self._compiled and hasattr(self, "_compiled_update"): - new_state_pure, grad_norm, is_skipped = self._compiled_update( - state_pure, self._accumulated_grads, self._micro_step_count, mean_loss - ) - else: - new_state_pure, grad_norm, is_skipped = self._update_kernel( - state_pure, self._accumulated_grads, self._micro_step_count, mean_loss - ) + # `state_pure` aliases the model's and optimizer's live buffers and is donated to the + # compiled kernel. Between the call and the `nnx.update` below, `self._state` is torn: its + # arrays have been deleted and reading one raises "Array has been deleted". Keep those two + # statements adjacent. + with self._sharding_ctx(): + if self._compiled and hasattr(self, "_compiled_update"): + new_state_pure, grad_norm, is_skipped = self._compiled_update( + state_pure, self._accumulated_grads, self._accumulated_denominator, mean_loss + ) + else: + new_state_pure, grad_norm, is_skipped = self._update_kernel( + state_pure, self._accumulated_grads, self._accumulated_denominator, mean_loss + ) nnx.update(self._state, new_state_pure) + self._publish_state(new_state_pure) if grad_norm is not None: self.record_metrics("gradient_norm", grad_norm) if is_skipped is not None: self.record_metrics("step_skipped", is_skipped) - # Add the state to the throttler queue so jax.block_until_ready() waits - # for the optimizer update to complete before logging the metrics. + # Queue something the update produced so jax.block_until_ready() waits for the optimizer + # update to complete before logging the metrics. The gradient norm rather than the state + # itself: the throttler keeps queued computations alive until it pops them, so handing it + # the train state pinned every parameter and optimizer slot -- three parameter trees -- + # for as long as the entry sat there, and once the update kernel donates its state + # argument those buffers are deleted by a later step, so an entry popped after that would + # raise "Array has been deleted" out of `jax.block_until_ready`. The norm is an output of + # the same executable as the weight update, so its readiness still means the update + # landed, and it is a tiny buffer that nothing donates. Tunix v2 uses its own update's + # gradient norm for exactly this (`peft_trainer_v2.py`, `_last_update_grad_norm`). self._throttler.add_computation( - self._state if self._state is not None else self._model, + grad_norm if grad_norm is not None else (self._state if self._state is not None else self._model), self._metrics_recorder.get_step_metrics(self.train_step), ) self._cached_losses.clear() self._accumulated_grads = None + self._accumulated_denominator = None self._micro_step_count = 0 self._train_step += 1 @@ -961,6 +1293,10 @@ def save_checkpoint(self, metadata: Any, **kwargs: Any) -> None: if metadata: # Metadata from Orchestrator custom_metadata["additional_metadata"] = metadata + # The accumulated gradients are stored unreduced, so the divisor `update()` will apply to + # them has to survive the round-trip too. It is a scalar, so metadata is the cheapest home. + if self._micro_step_count > 0 and self._accumulated_denominator is not None: + custom_metadata["accumulated_denominator"] = float(self._accumulated_denominator) ckpt_saved = self._checkpoint_manager.save_checkpoint( step=step, @@ -1005,6 +1341,10 @@ def restore_checkpoint(self, **kwargs: Any) -> Any: return None logging.info("Checkpoint restored from step %d.", restored_step) + # Orbax has just written new arrays into the live NNX variables, which is the one thing + # that makes the cached pure state wrong rather than merely old. Drop it; the next step + # re-reads the restored weights. + self._invalidate_pure_state() if restored_checkpoint_state.accumulated_metrics: buffers = [] @@ -1034,8 +1374,10 @@ def restore_checkpoint(self, **kwargs: Any) -> Any: # Checkpoint with no metadata says nothing about how far into its step it # got, and must not inherit the count from whatever this engine was doing before. self._micro_step_count = 0 + restored_denominator = None if restored_metadata: self._micro_step_count = restored_metadata.get("micro_step_count", 0) + restored_denominator = restored_metadata.get("accumulated_denominator", None) restored_additional_metadata = restored_metadata.get("additional_metadata", None) if self._micro_step_count > 0: @@ -1058,13 +1400,15 @@ def restore_checkpoint(self, **kwargs: Any) -> Any: # above, which the branch above has already discarded. if self._micro_step_count > 0 and restored_checkpoint_state.accumulated_grads: self._accumulated_grads = restored_checkpoint_state.accumulated_grads + self._accumulated_denominator = jnp.float32(restored_denominator if restored_denominator else 0.0) + rebuilt_losses = None if self._metrics_recorder._metrics_buffer: # pylint: disable=protected-access active_buf = self._metrics_recorder.get_step_metrics(restored_step) if active_buf and "loss" in active_buf.weighted_metrics: wm = active_buf.weighted_metrics["loss"] if wm.unreduced_sum.ndim > 0: - self._cached_losses = [ + rebuilt_losses = [ abstract_engine.WeightedMetric( unreduced_sum=wm.unreduced_sum[i], denominator=wm.denominator[i], @@ -1074,7 +1418,18 @@ def restore_checkpoint(self, **kwargs: Any) -> Any: for i in range(wm.unreduced_sum.shape[0]) ] else: - self._cached_losses = [wm] + rebuilt_losses = [wm] + self._cached_losses = rebuilt_losses + + # Checkpoints written before the denominator was tracked carry no value for it. The + # per-micro-batch losses just rebuilt above carry the very denominators that went into + # the saved gradients, so their sum is exactly what was lost. Only those count: any + # `_cached_losses` left over from before the restore belong to a different run. + if not restored_denominator and rebuilt_losses: + denominator = jnp.float32(0.0) + for cached_loss in rebuilt_losses: + denominator = denominator + jnp.sum(cached_loss.denominator).astype(jnp.float32) + self._accumulated_denominator = denominator return restored_additional_metadata diff --git a/src/maxtext/training_engine/metrics.py b/src/maxtext/training_engine/metrics.py index 3e6a2caadf..8e82420074 100644 --- a/src/maxtext/training_engine/metrics.py +++ b/src/maxtext/training_engine/metrics.py @@ -40,6 +40,22 @@ "tflops", ] +# How many completed step buffers stay resident before the oldest are evicted. +# +# Every buffer holds live device arrays -- one leaf per scalar metric, plus one per aux +# metric when the engine runs with `has_aux`, which for MoE/MTP configs is tens of leaves -- +# so nothing in the list is free. Nothing on the engine's own step path removes an entry +# either: `get_step_metrics` hands out the newest by reference, and only +# `get_metrics_history(clear_cache=True)` and `cleanup()` clear. A driver that reads the +# engine's own TensorBoard output rather than calling `get_metrics()` therefore never +# clears, and both HBM and -- because `save_checkpoint` serializes the whole retained +# history -- checkpoint size and save latency grow linearly in steps. +# +# Tunix v2 keeps exactly one prior step (`_prev_buffered_train_metrics` in +# `peft_trainer_v2.py`). This keeps a window instead so batched readers of +# `get_metrics_history` still work, while making the footprint constant in step count. +_DEFAULT_MAX_BUFFERED_STEPS = 128 + class MetricsRecorder: """Synchronous frontend for buffering and aggregating step metrics on-device. @@ -53,8 +69,17 @@ class MetricsRecorder: for processing. """ - def __init__(self): + def __init__(self, max_buffered_steps: int = _DEFAULT_MAX_BUFFERED_STEPS): + """Initializes the recorder. + + Args: + max_buffered_steps: How many completed step buffers to retain; older ones are evicted + as new steps start. Pass 0 or a negative value to retain everything, which is only + safe when the driver drains the history itself. + """ self._metrics_buffer: list[abstract_engine.MetricsBuffer] = [] + self._max_buffered_steps = max_buffered_steps + self._dropped_buffer_count = 0 def buffer_metrics( self, @@ -74,10 +99,36 @@ def buffer_metrics( if not self._metrics_buffer or self._metrics_buffer[-1].id != train_step: new_buffer = abstract_engine.MetricsBuffer(id=train_step, mode="train") self._metrics_buffer.append(new_buffer) + self._evict_old_buffers() # Record the new metric in the buffer for the current step. self._record_metric(name, metric, aggregation_fn=aggregation_fn) + def _evict_old_buffers(self) -> None: + """Drops the oldest step buffers once the retention window is full. + + Only ever runs when a *new* step starts, so the buffer the current step is writing into + and the one `get_step_metrics` is about to hand the throttler are never the ones dropped. + """ + if self._max_buffered_steps <= 0 or len(self._metrics_buffer) <= self._max_buffered_steps: + return + num_dropped = len(self._metrics_buffer) - self._max_buffered_steps + oldest_dropped_id = self._metrics_buffer[0].id + del self._metrics_buffer[:num_dropped] + self._dropped_buffer_count += num_dropped + # Dropping metrics silently is the pattern that produced the fabricated 0.0 in the parity + # harness, so say so -- but once per window, not once per step. + logging.log_every_n( + logging.WARNING, + "Metrics history is full at %d step(s); evicting buffers from step %s onwards " + "(%d dropped so far). Drain it with MetricsRecorder.get_metrics_history() or the " + "engine's get_metrics() if you need every step.", + self._max_buffered_steps, + self._max_buffered_steps, + oldest_dropped_id, + self._dropped_buffer_count, + ) + def _record_metric( self, name: str, @@ -120,13 +171,14 @@ def get_metrics_history(self, clear_cache: bool = True) -> list[abstract_engine. """Returns every cached step buffer and optionally clears the metrics cache. The engine's own `get_metrics` returns only the most recent buffer, per the trainer - contract. This is the accessor that keeps the full history reachable. + contract. This is the accessor that keeps the history reachable -- the last + `max_buffered_steps` of it; see `_DEFAULT_MAX_BUFFERED_STEPS` for why it is a window. Args: clear_cache: Whether to reset cached metrics after retrieval. Returns: - One on-device MetricsBuffer per recorded train step, oldest first. + One on-device MetricsBuffer per retained train step, oldest first. """ metrics_to_return = self._metrics_buffer if clear_cache: diff --git a/tests/post_training/unit/maxtext_engine_test.py b/tests/post_training/unit/maxtext_engine_test.py index af49f81d73..0e230b23a8 100644 --- a/tests/post_training/unit/maxtext_engine_test.py +++ b/tests/post_training/unit/maxtext_engine_test.py @@ -48,6 +48,19 @@ def __init__(self): self.weights = nnx.Param(jnp.array([1.0, 2.0])) +class DummyStatefulNNXModel(nnx.Module): + """A model whose state is not all `nnx.Param`, so `rest` is non-empty. + + `DummyNNXModel` is parameter-only, which makes `nnx.split(model, nnx.Param, ...)` return an + empty `rest` and every publish of non-parameter state a no-op. Real models carry RNG + counters and batch statistics, so one model here has to as well. + """ + + def __init__(self): + self.weights = nnx.Param(jnp.array([1.0, 2.0])) + self.calls = nnx.BatchStat(jnp.array(0.0)) + + @dataclasses.dataclass(kw_only=True) class DummyPayload(abstract_engine.TrainerPayload): token_ids: Any = dataclasses.field(default_factory=lambda: jnp.ones((2, 2))) @@ -162,6 +175,82 @@ def test_max_text_trainer_instantiation_with_pyconfig(self): self.assertIsNone(t._accumulated_grads) self.assertEqual(t.train_step, 2) + def test_compiled_steps_publish_weights_and_non_param_state(self): + """Exercises the cached pure-state path end to end, which nothing else on CPU does. + + Three conditions have to hold at once for the cache to be involved, and no other test + here meets all three: the model must have non-`Param` state (otherwise `_publish_model_rest` + is vacuous), `compile()` must be called (the cache is seeded by `_compile_for_batch`, so + the eager path never touches it), and `update()` must run (only that reaches + `_publish_state`). Two steps rather than one, because the interesting failure is the + second step reading a stale or wrongly-partitioned cache written by the first. + """ + mesh = jax.sharding.Mesh(maxtext_utils.create_device_mesh(self.mock_config), self.mock_config.mesh_axes) + self.mock_from_pretrained.return_value = (DummyStatefulNNXModel(), mesh) + t = maxtext_engine.MaxTextTrainingEngine(self.mock_config) + + def loss_fn(model, *_args, **_kwargs): + # Mutating a non-`Param` variable is what makes `new_rest` differ from the cached + # `rest`; scaling the loss by it makes a stale publish show up as a wrong gradient + # rather than only as a wrong counter. + model.calls.value = model.calls.value + 1.0 + return ( + abstract_engine.WeightedMetric( + unreduced_sum=jnp.sum(model.weights.value) * model.calls.value, + denominator=jnp.array(1.0), + ), + {}, + ) + + t.with_loss_fn(loss_fn) + payload = DummyPayload() + before = np.asarray(t.model.weights.value) + + t.compile(payload) + self.assertIsNotNone(t._params_pure, "compile() did not seed the pure-state cache") + for _ in range(2): + t.fwd_bwd(payload) + t.update() + + self.assertIsNotNone(t._params_pure, "the pure-state cache fell back to re-splitting the graph") + # `nnx.update` is the publish barrier: the live module the engine hands out, checkpoints + # and syncs weights from must track the cache, not lag behind it. + self.assertEqual(float(t.model.calls.value), 2.0) + self.assertGreater(float(np.abs(np.asarray(t.model.weights.value) - before).max()), 0.0) + self.assertEqual(t.train_step, 2) + + def test_gradient_norm_is_over_the_accumulated_normalized_gradient(self): + """Pins down *which* gradient is normed once accumulation no longer means-of-means. + + `test_gradient_norm_is_recorded_every_step` already covers the norm existing with + `skip_step_on_spikes` off, on a single micro-batch. Two micro-batches here, so the value + can only come out right if the norm is taken over the accumulated sum after its single + division by the accumulated denominator: a norm over either micro-batch's own gradient + reads sqrt(2), and one taken before the division reads 4x this. The micro-batches are + uniform, so this does not separate sum/sum from mean-of-means -- ยง3 of the parity write-up + is where that distinction is measured. + """ + self.assertFalse(self.mock_config.skip_step_on_spikes) + t = maxtext_engine.MaxTextTrainingEngine(self.mock_config) + # d(unreduced_sum)/dw is 1.0 per element, so two micro-batches accumulate [2.0, 2.0] + # against a denominator of 8.0 -> [0.25, 0.25], whose l2 norm is sqrt(2 * 0.25**2). + t.with_loss_fn( + lambda model, *_args, **_kwargs: ( + abstract_engine.WeightedMetric(unreduced_sum=jnp.sum(model.weights.value), denominator=jnp.array(4.0)), + {}, + ) + ) + payload = DummyPayload() + t.fwd_bwd(payload) + t.fwd_bwd(payload) + t.update() + + metrics = t.get_metrics(clear_cache=True) + self.assertIn("gradient_norm", metrics.scalar_metrics) + recorded = np.asarray(metrics.scalar_metrics["gradient_norm"]).reshape(-1) + self.assertEqual(recorded.shape, (1,), "one norm per update, not one per micro-batch") + np.testing.assert_allclose(recorded[0], np.sqrt(2 * 0.25**2), rtol=1e-5) + @mock.patch("orbax.checkpoint.CheckpointManager") def test_max_text_trainer_checkpoint_manager_init(self, mock_create_mgr): mock_config = self.setup_config(enable_checkpointing=True) @@ -380,6 +469,7 @@ def test_save_checkpoint_called_after_fwd_bwd_before_update(self): t._micro_step_count = 1 t.train_step = 10 t._accumulated_grads = {"params": {"w": jnp.array([0.5, 0.5])}} + t._accumulated_denominator = jnp.float32(6.0) dummy_metadata = mock.MagicMock() t.save_checkpoint(metadata=dummy_metadata) @@ -388,6 +478,8 @@ def test_save_checkpoint_called_after_fwd_bwd_before_update(self): mock_orbax_mgr.save.assert_called_once() call_kwargs = mock_orbax_mgr.save.call_args.kwargs self.assertEqual(call_kwargs["custom_metadata"]["micro_step_count"], 1) + # The gradients are saved unreduced, so their divisor has to ride along with them. + self.assertEqual(call_kwargs["custom_metadata"]["accumulated_denominator"], 6.0) self.assertEqual(call_kwargs["custom_metadata"]["additional_metadata"], dummy_metadata) args_dict = ( dict(call_kwargs["args"].items()) @@ -560,19 +652,21 @@ def test_update_with_inflight_throttling(self): self.assertEqual(t.train_step, 1) # wait_for_next() in update() sees qsize=2 (full), so it pops # index 0 (loss for micro_step_count=0), leaving qsize=1. - # Then add_computation() queues the updated model state and step 0 metrics. + # Then add_computation() queues the update's gradient norm and step 0 metrics. # Since we removed the trailing wait_for_next() from update(), qsize # remains 2. self.assertEqual(t._throttler._inflight_queue.qsize(), 2) - expected_state_leaves = jax.tree.leaves(t._state if t._state else t._model) for idx, (computation, metrics) in enumerate(t._throttler._inflight_queue.queue): if idx == 0: # Loss for micro_step_count=0. self.assertIsNone(metrics) if idx == 1: - # Metrics for train_step=0. + # Metrics for train_step=0, waited on through the update's gradient norm -- one + # scalar out of the same executable, not the state itself, whose buffers the next + # update donates away. self.assertIsNotNone(metrics) - self.assertEqual(computation, expected_state_leaves) + self.assertLen(computation, 1) + self.assertEqual(jnp.shape(computation[0]), ()) # train_step=1: fwd_bwd + update # Calling fwd_bwd() while queue is full (qsize=2) triggers wait_for_next(), @@ -612,8 +706,10 @@ def _loss_fn(model, *args, **kwargs): self.assertEqual(t._micro_step_count, 1) self.assertIsNotNone(t._accumulated_grads) - # Check that grad is scaled by 1/4.0 - np.testing.assert_allclose(t._accumulated_grads["weights"], jnp.array([2.0, 2.0]), rtol=1e-5) + # Gradients accumulate unreduced: d(unreduced_sum)/dw is 8.0 per element and the 1/4.0 + # from the denominator is applied once, in update(). + np.testing.assert_allclose(t._accumulated_grads["weights"], jnp.array([8.0, 8.0]), rtol=1e-5) + np.testing.assert_allclose(t._accumulated_denominator, 4.0, rtol=1e-5) metrics = t.get_metrics(clear_cache=True) self.assertIn("loss", metrics.weighted_metrics) @@ -646,8 +742,9 @@ def custom_loss(model, *args, **kwargs): t.with_loss_fn(custom_loss, has_aux=True) t.fwd_bwd(payload) - # Check that grad is scaled by 1/4.0 - np.testing.assert_allclose(t._accumulated_grads["weights"], jnp.array([2.0, 2.0]), rtol=1e-5) + # Unreduced, with the 1/4.0 deferred to update(). + np.testing.assert_allclose(t._accumulated_grads["weights"], jnp.array([8.0, 8.0]), rtol=1e-5) + np.testing.assert_allclose(t._accumulated_denominator, 4.0, rtol=1e-5) metrics = t.get_metrics(clear_cache=True) self.assertIn("loss", metrics.weighted_metrics) self.assertIn("aux_stat", metrics.scalar_metrics) @@ -671,7 +768,7 @@ def _loss_fn(model, **kwargs): # Arrived by keyword, under the names the adapter chose. self.assertEqual(sorted(seen), ["alpha", "beta"]) - np.testing.assert_allclose(t._accumulated_grads["weights"], jnp.array([2.0, 2.0]), rtol=1e-5) + np.testing.assert_allclose(t._accumulated_grads["weights"], jnp.array([8.0, 8.0]), rtol=1e-5) def test_without_gen_model_input_fn_the_maxtext_convention_is_kept(self): """With no adapter, the loss still gets MaxText's positional signature.""" @@ -742,7 +839,7 @@ def test_compiled_path_closes_over_non_array_loss_arguments(self): compiled.fwd_bwd(DummyPayload()) np.testing.assert_allclose(compiled._accumulated_grads["weights"], eager._accumulated_grads["weights"], rtol=1e-5) - np.testing.assert_allclose(compiled._accumulated_grads["weights"], jnp.array([2.0, 2.0]), rtol=1e-5) + np.testing.assert_allclose(compiled._accumulated_grads["weights"], jnp.array([8.0, 8.0]), rtol=1e-5) def test_compile_without_dummy_data_defers_to_first_fwd_bwd(self): """`compile(None)` cannot know input shapes, so it defers instead of failing. @@ -761,7 +858,7 @@ def test_compile_without_dummy_data_defers_to_first_fwd_bwd(self): # Deferred, not abandoned: the first real batch supplies the shapes. t.fwd_bwd(DummyPayload()) self.assertTrue(t._compiled) - np.testing.assert_allclose(t._accumulated_grads["weights"], jnp.array([2.0, 2.0]), rtol=1e-5) + np.testing.assert_allclose(t._accumulated_grads["weights"], jnp.array([8.0, 8.0]), rtol=1e-5) def test_compiled_kernel_is_rebuilt_when_a_static_loss_argument_changes(self): """A changed non-traced loss argument must reach the loss, not the stale closure. @@ -787,14 +884,14 @@ def _loss_fn(model, tokens, algo_config): t.compile(DummyPayload()) t.fwd_bwd(DummyPayload()) - np.testing.assert_allclose(t._accumulated_grads["weights"], jnp.array([2.0, 2.0]), rtol=1e-5) + np.testing.assert_allclose(t._accumulated_grads["weights"], jnp.array([8.0, 8.0]), rtol=1e-5) # Same payload shapes, different static value: only the static half of the signature # can catch this. holder[0] = types.SimpleNamespace(scale=3.0) t._accumulated_grads = None t.fwd_bwd(DummyPayload()) - np.testing.assert_allclose(t._accumulated_grads["weights"], jnp.array([6.0, 6.0]), rtol=1e-5) + np.testing.assert_allclose(t._accumulated_grads["weights"], jnp.array([24.0, 24.0]), rtol=1e-5) def test_uncomparable_static_arguments_warn_once(self): """An uncomparable static argument recompiles every step, and says so once. @@ -1047,6 +1144,33 @@ def test_get_metrics_returns_newest_and_history_stays_reachable(self): self.assertEqual(newest.id, 2) self.assertIn("dropping 2 older buffer", "".join(logs.output)) + def test_metrics_history_is_bounded(self): + """The step history is a window, so a driver that never drains it cannot grow forever. + + Each retained buffer pins live device arrays and `save_checkpoint` serializes the whole + history, so an unbounded list would cost HBM and checkpoint latency linear in step count. + """ + recorder = metrics_module.MetricsRecorder(max_buffered_steps=4) + for step in range(10): + recorder.buffer_metrics(train_step=step, name="loss", metric=jnp.array(float(step))) + + history = recorder.get_metrics_history(clear_cache=False) + self.assertLen(history, 4) + # The newest steps are the ones kept, and the step currently being written is never evicted. + self.assertEqual([b.id for b in history], [6, 7, 8, 9]) + self.assertEqual(recorder.get_step_metrics(9).id, 9) + self.assertEqual(recorder._dropped_buffer_count, 6) + + # Opting out is possible for drivers that drain the history themselves. + unbounded = metrics_module.MetricsRecorder(max_buffered_steps=0) + for step in range(10): + unbounded.buffer_metrics(train_step=step, name="loss", metric=jnp.array(float(step))) + self.assertLen(unbounded.get_metrics_history(clear_cache=False), 10) + + # The engine's own recorder is bounded by default. + t = maxtext_engine.MaxTextTrainingEngine(self.mock_config) + self.assertGreater(t._metrics_recorder._max_buffered_steps, 0) + def test_has_aux_false_drops_tuple_aux(self): """`has_aux=False` suppresses aux recording; `has_aux=True` keeps it. @@ -1122,8 +1246,8 @@ def _loss_fn(model, *args, **kwargs): t.with_loss_fn(_loss_fn) t.fwd_bwd(payload) - # d(unreduced_sum)/dw is 8.0 per element, scaled by compute_scale() = 1/4.0. - np.testing.assert_allclose(t._accumulated_grads["weights"], jnp.array([2.0, 2.0]), rtol=1e-5) + # d(unreduced_sum)/dw is 8.0 per element; the 1/4.0 is applied once, in update(). + np.testing.assert_allclose(t._accumulated_grads["weights"], jnp.array([8.0, 8.0]), rtol=1e-5) metrics = t.get_metrics(clear_cache=True) self.assertIn("loss", metrics.weighted_metrics) self.assertAlmostEqual(float(metrics.weighted_metrics["loss"].compute().item()), 6.0, places=4) @@ -1154,7 +1278,7 @@ def _tunix_loss_fn(model, *args, **kwargs): t.with_loss_fn(_tunix_loss_fn) t.fwd_bwd(payload) - np.testing.assert_allclose(t._accumulated_grads["weights"], jnp.array([2.0, 2.0]), rtol=1e-5) + np.testing.assert_allclose(t._accumulated_grads["weights"], jnp.array([8.0, 8.0]), rtol=1e-5) metrics = t.get_metrics(clear_cache=True) self.assertIn("loss", metrics.weighted_metrics) self.assertIn("metric_a", metrics.weighted_metrics) From 7bab74ef7d648b6e52d194df36960cb1186bb745 Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Wed, 2 Sep 2026 03:59:01 +0000 Subject: [PATCH 2/3] Size the GRPO integration test to the mesh The engine now traces its kernels under `nn_partitioning.axis_rules`, which is what makes MaxText's logical constraints on the activations real. That turns the test's hard-coded batch of 2 into a problem on any host with more than 2 data x fsdp devices: XLA pads the batch out, the padded lanes are all-zero sequences that mask themselves out of attention entirely, and their contribution comes back as NaN on the pad token's embedding row -- a finite loss with unusable gradients, so the "a parameter moved" assertion saw nan. The splash kernel asserts on exactly this ratio; dot_product does not, which is why it surfaced as a NaN rather than an error. Derive the batch from the mesh instead, and pin matmul_precision. The KL assertion compares log-probs from two code paths -- tunix's compute_per_token_logps for the reference and the engine's sharded forward for the policy -- and at bf16 they disagree by ~2e-2 on values near -12.6, which low_var_kl squares into ~1e-4. --- .../maxtext_engine_grpo_loss_test.py | 21 ++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/tests/post_training/integration/maxtext_engine_grpo_loss_test.py b/tests/post_training/integration/maxtext_engine_grpo_loss_test.py index 224ccc29c2..448a714576 100644 --- a/tests/post_training/integration/maxtext_engine_grpo_loss_test.py +++ b/tests/post_training/integration/maxtext_engine_grpo_loss_test.py @@ -90,6 +90,13 @@ def _config(**overrides) -> pyconfig.HyperParameters: "learning_rate=1e-4", "micro_batch_size_to_train_on=2", "max_target_length=64", + # The KL assertion below compares log-probs produced by two different code paths: + # tunix's compute_per_token_logps for the reference, and the engine's own sharded + # forward pass for the policy. At the default bf16 matmul precision those disagree + # by ~2e-2 on log-probs near -12.6, and low_var_kl squares the difference into a + # KL of ~1e-4 -- noise, but large enough to swamp a real regression. fp32 matmuls + # bring it back under 1e-5; the model is tiny enough that the cost is invisible. + "matmul_precision=highest", ] argv.extend(f"{k}={v}" for k, v in overrides.items()) return pyconfig.initialize(argv) @@ -171,6 +178,18 @@ class MaxTextEngineGrpoLossTest(absltest.TestCase): def test_grpo_loss_drives_a_training_step(self): cfg = _config() mesh = jax.sharding.Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes) + # The batch has to divide the data x fsdp device count, and only the mesh knows that + # number, so it cannot be hard-coded in `_config`. The engine traces its kernels under + # `nn_partitioning.axis_rules`, which is what makes MaxText's logical constraints on the + # activations real; a batch those devices cannot split is then padded out by XLA, and the + # padded lanes are all-zero sequences that mask themselves out of attention entirely. + # Their contribution comes back as NaN on the pad token's embedding row -- a finite loss + # and unusable gradients. The splash kernel asserts on exactly this ratio + # (`attention_op.py`, "Batch dimension should be shardable"); dot_product does not, so + # here it surfaced as a NaN instead of an error. + batch = int(mesh.shape["data"] * mesh.shape["fsdp"]) + if cfg.micro_batch_size_to_train_on != batch: + cfg = _config(micro_batch_size_to_train_on=batch) algo_config = _GrpoConfig() engine = maxtext_engine.MaxTextTrainingEngine( @@ -197,7 +216,7 @@ def test_grpo_loss_drives_a_training_step(self): ) self.assertIs(returned, engine) - payload = _train_example(engine.model, algo_config) + payload = _train_example(engine.model, algo_config, batch=batch) before = _param_leaves(engine.model) engine.fwd_bwd(payload) From b3d54c23083755162d304b20143ea8f072edc6bc Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Wed, 2 Sep 2026 04:43:27 +0000 Subject: [PATCH 3/3] Cut the comments down Same code, 131 fewer lines of prose. Keeps the load-bearing ones -- why params are not donated, why the throttler queues a scalar, why the batch has to be shardable -- and drops the measurements and narrative that belong in the write-up rather than the source. --- .../training_engine/inflight_throttler.py | 21 +- src/maxtext/training_engine/maxtext_engine.py | 248 ++++++------------ src/maxtext/training_engine/metrics.py | 27 +- .../maxtext_engine_grpo_loss_test.py | 22 +- .../post_training/unit/maxtext_engine_test.py | 11 +- 5 files changed, 99 insertions(+), 230 deletions(-) diff --git a/src/maxtext/training_engine/inflight_throttler.py b/src/maxtext/training_engine/inflight_throttler.py index 949336fbc9..488040fd2e 100644 --- a/src/maxtext/training_engine/inflight_throttler.py +++ b/src/maxtext/training_engine/inflight_throttler.py @@ -34,15 +34,13 @@ def __init__(self, config: pyconfig.HyperParameters): """ self._inflight_queue = queue.Queue[Any](maxsize=config.max_inflight_computations) self._metrics_logger = metrics_module.MetricsLogger(config=config) - # Popped by `wait_for_next` but not yet written; see `_flush_pending_metrics`. self._pending_metrics: abstract_engine.MetricsBuffer | None = None def add_computation(self, computation: Any, metrics: abstract_engine.MetricsBuffer | None) -> None: """Adds an active on-device computation to the queue.""" self._inflight_queue.put((jax.tree.leaves(computation), metrics)) - # The caller has just dispatched, so the device has work queued behind this point and the - # blocking read inside `write_metrics` overlaps it instead of running against an idle - # device. This is the whole reason the write is deferred rather than done in place. + # Flushed here, not in `wait_for_next`: the caller has just dispatched, so the blocking + # read inside `write_metrics` overlaps that work instead of an idle device. self._flush_pending_metrics() def _flush_pending_metrics(self) -> None: @@ -55,20 +53,14 @@ def _flush_pending_metrics(self) -> None: def wait_for_next(self) -> None: """If the limit is reached, wait for the next computation to finish. - Blocks, but does not log. `write_metrics` reduces each `WeightedMetric` on device and then - pulls the result to host with `np.asarray`, and those reduction ops are dispatched behind - whatever is already in the device's queue -- so doing it here, before the caller dispatches - the step it just made room for, stalls the host on the full backlog with nothing new - running. Deferring to the following `add_computation` costs one dispatch of staleness in - the log and hands the same work a busy device to hide behind. Metrics carry their own step - id (`MetricsBuffer.id`), so nothing downstream can tell the difference. + Blocks, but does not log: the metrics write is stashed for the next `add_computation`. + Buffers carry their own step id, so the extra dispatch of staleness is invisible. """ if self._inflight_queue.full(): computation, metrics = self._inflight_queue.get() jax.block_until_ready(computation) if metrics is not None: - # Never hold two: the engine only attaches metrics to one of its two computations per - # step, but a caller that attached them to both would otherwise silently lose a buffer. + # Never hold two, or a caller attaching metrics to every computation loses a buffer. self._flush_pending_metrics() self._pending_metrics = metrics @@ -81,8 +73,7 @@ def wait_for_all(self) -> None: if metrics is not None: self._flush_pending_metrics() self._pending_metrics = metrics - # Draining is the one place that must not leave a write outstanding: callers use it to - # reach a quiescent state before checkpointing or shutting down. + # A drain must not leave a write outstanding. self._flush_pending_metrics() def cleanup(self) -> None: diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py index 8ae81663ea..e397b3f4a5 100644 --- a/src/maxtext/training_engine/maxtext_engine.py +++ b/src/maxtext/training_engine/maxtext_engine.py @@ -55,8 +55,8 @@ # same situation. Real buffers are identified by their train step, so this cannot collide. EMPTY_METRICS_BUFFER_ID = -1 -# Where `nnx.split(TrainStateNNX(...))` puts the model's own state. See -# `_check_pure_state_reusable` for why the engine verifies this rather than assuming it. +# Where `nnx.split(TrainStateNNX(...))` puts the model's own state; verified, not assumed, +# by `_check_pure_state_reusable`. _MODEL_STATE_KEY = "model" _PURE_STATE_FALLBACK_WARNING = ( @@ -371,18 +371,15 @@ def __init__( else: self._model = model_or_model_mesh_pair self._state: Any = None - # Pure-pytree mirror of the model and the train state, carried across steps so the step - # path never re-walks the module graph. `None` means "not cached", which is also how the - # fast path is switched off; see `_refresh_pure_state`. + # 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". self._params_pure: Any = None self._rest_pure: Any = None self._state_pure: Any = None self._pure_state_warned: bool = False self._accumulated_grads: Any = None - # Sum of the per-micro-batch loss denominators behind `_accumulated_grads`. Tracked - # alongside them because the gradients are accumulated unreduced: this is the divisor - # `update()` applies once, and it is not `micro_step_count` unless every micro-batch - # happened to carry the same token count. + # Summed loss denominators behind `_accumulated_grads`, which are unreduced: this is the + # divisor `update()` applies once. self._accumulated_denominator: Any = None self._micro_step_count = 0 # Set when this run resumed from an intra-step checkpoint, cleared once the step it @@ -519,19 +516,12 @@ def with_gen_model_input_fn(self, gen_model_input_fn: Callable[[Any], dict[str, def _sharding_ctx(self): """Activates the mesh and logical axis rules the MaxText layers are written against. - MaxText modules place their own sharding constraints through - `nn_partitioning.get_axis_rules()` (`sharding.maybe_shard_with_logical` and friends). - Those rules live in a context variable, so a kernel traced outside this context sees an - empty rule set: every logical constraint silently becomes a no-op and XLA is left to - guess the partitioning for activations and gradients. It guesses badly -- on - llama3.1-8b/fsdp=8 the same fwd/bwd measured 1012 ms untraced-in-context against 581 ms - inside it, with no numerical difference. `train.py` wraps its own `jax.jit` the same way - (`with jax.set_mesh(mesh), mesh, nn_partitioning.axis_rules(config.logical_axis_rules)`), - which is why the standalone trainer never hit this. Its middle `mesh` is left out here: - `Mesh.__enter__` is deprecated and `jax.set_mesh` already covers it. - - Entered around the *call*, not around `jax.jit(...)`: jit is lazy, so the rules must be - live when the first call triggers tracing. + The rules live in a context variable, so a kernel traced outside this context sees an + empty rule set, every `maybe_shard_with_logical` becomes a no-op and XLA guesses the + partitioning -- badly: 1012 ms against 581 ms for the same fwd/bwd on llama3.1-8b/fsdp=8. + `train.py` wraps its own `jax.jit` the same way. Entered around the *call*, since jit is + lazy and the rules must be live when tracing happens. Note a batch that data x fsdp + cannot divide gives NaN gradients once the constraints are real. """ if self._mesh is None: yield @@ -542,9 +532,8 @@ def _sharding_ctx(self): def _invalidate_pure_state(self) -> None: """Forgets the cached pure state, so the next step re-reads it from the NNX objects. - Called from the `model`/`optimizer`/`state` setters and after a checkpoint restore -- - the three ways the live NNX variables can be replaced behind the engine's back. It is - not needed on the step path: the step path is what *produces* the cached values. + For the three ways the live NNX variables get replaced behind the engine's back: the + `model`/`optimizer`/`state` setters, and a checkpoint restore. """ self._params_pure = None self._rest_pure = None @@ -561,36 +550,20 @@ def _disable_pure_state(self, reason: str) -> None: def _with_model_state(state_pure: Any, model_pure: Any) -> Any: """Returns `state_pure` with its model subtree replaced by `model_pure`. - Goes through `raw_mapping` rather than `{**state_pure}` because `nnx.State` stores its - children as plain dicts and only wraps them in a `State` on `__getitem__`. Rebuilding - from the wrapped views produces a tree that is equal key for key and leaf for leaf but - is a *different pytree*, one `State` node deeper at every level -- which `jax.jit` - rejects, at the call site, as an `in_shardings` prefix mismatch that names neither this - function nor the reason. + Through `raw_mapping` rather than `{**state_pure}`: `nnx.State` wraps children in a + `State` on `__getitem__`, and rebuilding from those views gives an equal-valued but + deeper pytree that `jax.jit` rejects as an `in_shardings` prefix mismatch. """ return nnx.State({**state_pure.raw_mapping, _MODEL_STATE_KEY: model_pure.raw_mapping}) def _check_pure_state_reusable(self, state_pure: Any, params_pure: Any, rest_pure: Any) -> str | None: """Returns why the pure state cannot be carried across steps, or None if it can. - The step path rebuilds the update kernel's `state_pure` argument by dropping the model's - parameter and non-parameter state back into `state_pure["model"]`, and re-derives the - next step's parameters from the kernel's output the same way. That is only the same - value `nnx.split` would have produced if NNX puts the model's state there, and puts it - there exactly once. - - It does for `TrainStateNNX`: `__init__` assigns `self.model` before `self.optimizer`, so - the model is flattened first and the optimizer's reference to the same module becomes a - graph reference rather than a second copy of the weights -- the pure dict is - `{"model": ..., "optimizer": {"opt_state": ..., "step": ...}}`. But that is a property of - a class this engine does not own, and `engine.state` is a public setter that will accept - anything, so it is checked rather than assumed. - - The check runs the real reconstruction and compares the result against what `nnx.split` - produced, rather than testing some proxy for it. A pytree that differs from the one the - kernels were compiled against fails at the `jax.jit` call site with an `in_shardings` - prefix mismatch, which is a hard error to read back to its cause; catching it here costs - one comparison per compile. + The step path rebuilds the kernel's `state_pure` by dropping the model's state back into + `state_pure["model"]`, which only reproduces `nnx.split` if NNX put it there exactly + once. `TrainStateNNX` does, but `engine.state` is a public setter that accepts anything, + so this runs the real reconstruction and compares treedefs -- once per compile, against + a `jax.jit` in_shardings mismatch that is hard to read back to its cause. Returns: `None` when the fast path is safe, else a short phrase naming what did not line up. @@ -601,28 +574,17 @@ def _check_pure_state_reusable(self, state_pure: Any, params_pure: Any, rest_pur return "this version of flax.nnx.State does not expose raw_mapping" rebuilt = self._with_model_state(state_pure, nnx.merge_state(params_pure, rest_pure)) if jax.tree.structure(rebuilt) != jax.tree.structure(state_pure): - # The likely cause is a train state that holds the model somewhere other than - # `.model`, or that holds a *second* module sharing its variables, which makes the - # first-flattened copy the only real one and this one a reference. return f"state[{_MODEL_STATE_KEY!r}] is not the model's own state" return None def _refresh_pure_state(self) -> None: """Re-reads the model and train state as pure `nnx.State`, and caches both. - Runs once per compile rather than once per step, which is the point. `nnx.split` walks - the entire module graph: on an unrolled 28-layer qwen3-0.6b it costs 51.6 ms for the - model and 40.8 ms for the train state, so the two calls the step path used to make were - 92 ms of a 283 ms step -- more than the 82 ms the step spent on the TPU. The same - partition applied to an already-flat state (`nnx.split_state`) costs 0.84 ms, because it - walks 490 leaves instead of 1756 graph nodes. - - What this does *not* change is when results are published: `fwd_bwd` and `update` still - `nnx.update` the live NNX objects at exactly the points they always did, so `self.model` - and `self.state` are never stale and nothing outside the step path has to know the cache - exists. Tunix v2 gets the same saving a different way, by holding the graph inside - `nnx.cached_partial` (`peft_trainer_v2.maybe_cache_and_partial`); that hook is specific - to `nnx.jit`, and these kernels are plain `jax.jit` over pure state. + Once per compile rather than once per step, which is the point: the two `nnx.split` + calls the step path used to make walked 1756 graph nodes for 92 ms of a 283 ms step on + an unrolled qwen3-0.6b, against 0.84 ms for `nnx.split_state` over the flat state. + Publication is unchanged -- `fwd_bwd` and `update` still `nnx.update` the live objects + where they always did, so `self.model` and `self.state` are never stale. """ if self._state is None: self._state = train_state_nnx.TrainStateNNX(self._model, self._optimizer) @@ -653,18 +615,10 @@ def _read_state_pure(self) -> Any: def _publish_model_rest(self, new_rest: Any) -> None: """Folds a fwd/bwd's updated non-parameter state into the cached train state. - Required, not an optimization: `update()` passes the whole train state to its kernel, and - before this cache existed the fresh `rest` reached it via `nnx.update(model, new_rest)` - landing in the variables that the following `nnx.split(state)` then read back. Skipping - the fold here would hand the update kernel the *previous* micro-batch's non-parameter - state -- for a model with RNG counters or batch statistics, a silent regression. - - The structure is checked, not assumed. `_fwd_bwd_kernel` re-splits the model it merged, so - anything the forward pass `sow`s as an `nnx.Intermediate` comes back as an extra entry in - `new_rest` -- `record_max_logits`, `distill_beta > 0` and multi-token prediction all do - this. Adopting a wider `rest` would leave the cache disagreeing with the `rest_shardings` - the kernel was compiled against, and the mismatch would surface one call later as a - `jax.jit` in_shardings prefix error naming neither the sow site nor this method. + Required, not an optimization: without it `update()` would see the *previous* + micro-batch's RNG counters and batch statistics. The structure is checked because + anything the forward pass `sow`s (`record_max_logits`, `distill_beta`, MTP) widens + `new_rest`, which would disagree with the shardings the kernel was compiled against. """ if self._params_pure is None: return @@ -677,14 +631,10 @@ def _publish_model_rest(self, new_rest: Any) -> None: def _publish_state(self, new_state_pure: Any) -> None: """Adopts the update kernel's output as the cached state, re-deriving `(params, rest)`. - The re-derivation is checked rather than trusted. It reads `.type` off the state's - leaves, which survive the round trip through `jax.jit` as part of the treedef -- but if - a future NNX ever flattened them differently, the partition would come back wrong and - the next `nnx.merge(self._model_graphdef, params, rest)` would fail inside a traced - kernel, where the error names a pytree mismatch and not this line. Comparing the treedef - here costs ~1 ms and turns that into a fallback plus one warning. Correctness does not - hinge on it either way: `update()` has already written `new_state_pure` into the live NNX - objects by the time this runs, so dropping the cache loses speed and nothing else. + The re-derivation reads `.type` off leaves that survive `jax.jit` in the treedef; the + treedef comparison costs ~1 ms and turns a future NNX flattening change into a fallback + plus a warning rather than a pytree error inside a traced kernel. Correctness does not + hinge on it: `update()` has already written the state into the live NNX objects. """ if self._params_pure is None: return @@ -705,9 +655,7 @@ def _fwd_bwd_kernel(self, params, rest, batch, acc_grads=None, acc_denom=None): rest: The model's remaining (non-parameter) pure state. batch: Loss-function inputs for this micro-batch. acc_grads: Gradients accumulated over earlier micro-batches of this update, or None on - the first micro-batch of an update. Passing None is what lets the first micro-batch - skip allocating -- and zeroing -- a parameter-sized buffer that would be overwritten - before anything read it. + the first, which is what lets it skip allocating a parameter-sized buffer. acc_denom: Denominator accumulated alongside `acc_grads`, or None with it. Returns: @@ -788,15 +736,10 @@ def diff_wrapper(p, r, b): micro_grads, ) - # The gradients accumulated here are the UNREDUCED ones -- d/dparam of the summed loss, - # with no `1/denominator` applied. `_update_kernel` divides the total by the total - # denominator, so the optimizer sees the global weighted mean `sum(grads)/sum(denom)` - # rather than a mean of per-micro-batch means. The two agree only when every micro-batch - # carries the same token count; under sequence packing or a ragged RL rollout they do - # not, and the mean-of-means silently overweights short micro-batches. It is also one - # fewer full pass over the gradient tree per micro-batch. This mirrors what MaxText's own - # pre-train path already does (`gradient_accumulation.py`: accumulate `xent_sum`, divide - # once by the summed `total_weights`). + # Accumulated UNREDUCED, with no `1/denominator` applied: `_update_kernel` divides once + # by the total, so the optimizer sees `sum(grads)/sum(denom)` rather than a mean of + # per-micro-batch means, which overweights short micro-batches. Same as + # `gradient_accumulation.py` in the pre-train path. denominator = loss_out.primary_loss.denominator.astype(jnp.float32) if acc_grads is None: return loss_out.primary_loss, loss_out.aux_metrics, new_rest, micro_grads, denominator @@ -813,24 +756,17 @@ def _update_kernel(self, state_pure, accumulated_grads, accumulated_denominator, grad_norm = None is_skipped_val = None if state_pure is not None: - # `accumulated_grads` holds sum_i grad(unreduced_sum_i) and `accumulated_denominator` - # holds sum_i denominator_i, so this one division is the whole normalization. A zero - # total means every micro-batch was empty; yield zeros rather than a NaN, matching - # `WeightedMetric.compute_scale()` and the `has_weights` guard in - # `gradient_accumulation.py`. + # This one division is the whole normalization. A zero total means every micro-batch + # was empty; yield zeros rather than a NaN, as `gradient_accumulation.py` does. has_weights = accumulated_denominator > 0 safe_denominator = jnp.where(has_weights, accumulated_denominator, 1.0) grads = jax.tree.map( lambda g: jnp.where(has_weights, g / safe_denominator.astype(g.dtype), jnp.zeros_like(g)), accumulated_grads, ) - # Before clipping, which is where 202a89ab8 put it and where Tunix's own - # `optax.global_norm` sits (`peft_trainer_v2._update_step`): its clipping, if any, is a - # link in the optax chain that runs after. Note this is `raw_grad_norm` in - # `train.py`'s vocabulary rather than its `learning/grad_norm`; with clipping off -- - # base.yml's default, and Tunix, which never clips -- the two coincide. In float32 - # whatever `grad_dtype` is, because a sum of squares over bf16 leaves overflows on - # production-size models, and it is the norm the throttler blocks on. + # Before clipping, where Tunix's `optax.global_norm` also sits -- `train.py` would call + # this `raw_grad_norm`. In float32 whatever `grad_dtype` is: a sum of squares over bf16 + # overflows on production-size models. grad_norm = max_utils.l2norm_pytree(jax.tree.map(lambda g: g.astype(jnp.float32), grads)) if self._config.gradient_clipping_threshold > 0: grads = maxtext_utils.apply_gradient_clipping(grads, None, self._config.gradient_clipping_threshold) @@ -963,10 +899,8 @@ def _compile_for_batch(self, dynamic_batch: Any, static_batch: dict[str, Any]) - `static_batch` is closed over rather than passed, so non-array loss arguments (Tunix's `algo_config`, `pad_id`, `eos_id`) never reach the jit boundary. """ - # Re-reads both graphs and both pure states, and is the only place on the step path that - # does: everything after this is maintained as plain pytrees until something invalidates - # the cache. A recompile is exactly when the graph may legitimately have changed shape, - # so it is also the right moment to re-derive the shardings below. + # 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() state_pure = self._read_state_pure() params_pure, rest_pure = self._read_model_pure(getattr(self._state, _MODEL_STATE_KEY, self._model)) @@ -999,22 +933,12 @@ def accum_kernel(params, rest, dynamic, acc_grads, acc_denom): # 1. JIT Compile Micro FWD/BWD Pass. # - # Two kernels rather than one. The first micro-batch of an update has nothing to add to, - # so it returns its own gradients and the engine adopts them as the accumulator; every - # later micro-batch folds into that buffer in place. Tunix v2 calls the same split - # "non-persistent vs persistent" mode. It buys two things: the first micro-batch never - # allocates or zeroes a parameter-sized buffer that would be overwritten unread, and the - # accumulating kernel can *donate* the accumulator so XLA writes the sum back into the - # incoming buffer. Before this, accumulation happened in Python - # (`jax.tree.map(jnp.add, ...)`), which materialized the micro-batch gradients as a - # program output *and* allocated a fresh sum -- two extra parameter-sized buffers live - # at once. `jax.jit` is lazy, so the accumulating kernel costs nothing to compile when - # every update consumes a single micro-batch. - # - # `params` is deliberately NOT donated: `micro_grads` has the same shape, dtype and - # sharding, and JAX matches donations by shard-shape rather than by position - # (`jax/_src/interpreters/mlir.py:_set_up_aliases`), so donating it would alias the - # weights straight into the gradient output and destroy them. + # Two kernels: the first micro-batch has no accumulator to add to and allocates none, + # later ones fold in and donate it, so the sum is written back in place instead of + # materializing the micro gradients plus a fresh sum. `jax.jit` is lazy, so the second + # costs nothing when every update takes one micro-batch. `params` is deliberately NOT + # donated: JAX matches donations by shard-shape, not position, so it would alias the + # weights into the gradient output. self._compiled_fwd_bwd = jax.jit( first_kernel, in_shardings=first_in_shardings, @@ -1029,14 +953,10 @@ def accum_kernel(params, rest, dynamic, acc_grads, acc_denom): # 2. JIT Compile Optimizer Update Pass. # - # `state_pure` (parameters plus optimizer slots) is donated, exactly as - # `maxtext_utils.get_functional_train_with_signature` does for the standalone trainer - # (`donate_argnums = 0`): it is dead the moment the kernel returns, since the engine - # rebinds the state from the output, so aliasing saves holding a second copy of the whole - # train state. The accumulated gradients are *not* donated. Every parameter-shaped output - # is already claimed by the incoming state (weights and optimizer slots alike), so a - # gradient donation has nothing left to alias to and JAX would only warn about it; the - # buffers are freed by `update()` dropping its reference anyway. + # `state_pure` is donated, as `get_functional_train_with_signature` does for the + # standalone trainer: the engine rebinds the state from the output, so it is dead on + # return. The gradients are not -- every parameter-shaped output is already claimed by + # the incoming state, so JAX would only warn. self._compiled_update = jax.jit( self._update_kernel, in_shardings=update_in_shardings, @@ -1106,16 +1026,14 @@ def fwd_bwd(self, payload: abstract_engine.TrainerPayload, **kwargs: Any) -> Non signature = _batch_signature(dynamic_batch, static_batch) if not self._compiled or self._needs_recompile(signature): self._compile_for_batch(dynamic_batch, static_batch) - # Read after any recompile, not before: `_compile_for_batch` refreshes the cache, and - # reading first would hand the new kernel a pure state split against the old graph. + # After any recompile, not before: reading first would hand the new kernel a pure + # state split against the old graph. params, rest = self._read_model_pure(model) with self._sharding_ctx(): if self._accumulated_grads is None: loss, aux, new_rest, acc_grads, acc_denom = self._compiled_fwd_bwd(params, rest, dynamic_batch) else: - # `self._accumulated_grads` and `self._accumulated_denominator` are donated by this - # call, so their buffers are gone once it returns. Rebinding both from the outputs - # below is what keeps that safe -- nothing else holds a reference to either. + # Both accumulators are donated here, so they are rebound from the outputs below. loss, aux, new_rest, acc_grads, acc_denom = self._compiled_fwd_bwd_accum( params, rest, dynamic_batch, self._accumulated_grads, self._accumulated_denominator ) @@ -1142,7 +1060,6 @@ def fwd_bwd(self, payload: abstract_engine.TrainerPayload, **kwargs: Any) -> Non self.record_metrics(key, value) self._cached_losses.append(loss) - # Accumulation happened inside the kernel; there is nothing to add here. self._accumulated_grads = acc_grads self._accumulated_denominator = acc_denom self._micro_step_count += 1 @@ -1173,12 +1090,9 @@ def update(self, **kwargs: Any) -> int: self._state = train_state_nnx.TrainStateNNX(self._model, self._optimizer) state_pure = self._read_state_pure() - # `_update_kernel` reads `mean_loss` only inside its `skip_step_on_spikes` branch, and that - # flag is read off `self._config` at trace time, so with spike-skipping off -- the base.yml - # default -- the value is discarded. Computing it is not free: `WeightedMetric.compute()` is - # seven eager XLA launches (the eps and min_denom clamps plus a safe divide), so this was - # seven dispatches per step feeding an argument the executable does not contain. `None` is - # an empty pytree, which is what `update_in_shardings` already declares for this position. + # `_update_kernel` reads `mean_loss` only under `skip_step_on_spikes`, which is traced + # off `self._config`, so otherwise this was seven eager launches per step + # (`WeightedMetric.compute()`) feeding an argument the executable does not contain. if not self._config.skip_step_on_spikes: mean_loss = None elif self._cached_losses: @@ -1186,10 +1100,8 @@ def update(self, **kwargs: Any) -> int: mean_loss = jnp.mean(jnp.stack(loss_values)) if len(loss_values) > 1 else loss_values[0] else: mean_loss = jnp.array(0.0) - # `state_pure` aliases the model's and optimizer's live buffers and is donated to the - # compiled kernel. Between the call and the `nnx.update` below, `self._state` is torn: its - # arrays have been deleted and reading one raises "Array has been deleted". Keep those two - # statements adjacent. + # `state_pure` is donated, so between this call and the `nnx.update` below `self._state` + # is torn -- reading one of its arrays raises "Array has been deleted". Keep them adjacent. with self._sharding_ctx(): if self._compiled and hasattr(self, "_compiled_update"): new_state_pure, grad_norm, is_skipped = self._compiled_update( @@ -1207,16 +1119,11 @@ def update(self, **kwargs: Any) -> int: if is_skipped is not None: self.record_metrics("step_skipped", is_skipped) - # Queue something the update produced so jax.block_until_ready() waits for the optimizer - # update to complete before logging the metrics. The gradient norm rather than the state - # itself: the throttler keeps queued computations alive until it pops them, so handing it - # the train state pinned every parameter and optimizer slot -- three parameter trees -- - # for as long as the entry sat there, and once the update kernel donates its state - # argument those buffers are deleted by a later step, so an entry popped after that would - # raise "Array has been deleted" out of `jax.block_until_ready`. The norm is an output of - # the same executable as the weight update, so its readiness still means the update - # landed, and it is a tiny buffer that nothing donates. Tunix v2 uses its own update's - # gradient norm for exactly this (`peft_trainer_v2.py`, `_last_update_grad_norm`). + # Queue something the update produced, so `jax.block_until_ready` waits for it before the + # metrics are logged. The gradient norm rather than the state: the throttler holds queued + # entries until it pops them, which pinned three parameter trees, and once the state is + # donated a late pop raises "Array has been deleted". The norm comes out of the same + # executable, so its readiness still means the update landed. Tunix v2 does the same. self._throttler.add_computation( grad_norm if grad_norm is not None else (self._state if self._state is not None else self._model), self._metrics_recorder.get_step_metrics(self.train_step), @@ -1293,8 +1200,7 @@ def save_checkpoint(self, metadata: Any, **kwargs: Any) -> None: if metadata: # Metadata from Orchestrator custom_metadata["additional_metadata"] = metadata - # The accumulated gradients are stored unreduced, so the divisor `update()` will apply to - # them has to survive the round-trip too. It is a scalar, so metadata is the cheapest home. + # The gradients are stored unreduced, so their divisor has to survive the round-trip too. if self._micro_step_count > 0 and self._accumulated_denominator is not None: custom_metadata["accumulated_denominator"] = float(self._accumulated_denominator) @@ -1341,9 +1247,8 @@ def restore_checkpoint(self, **kwargs: Any) -> Any: return None logging.info("Checkpoint restored from step %d.", restored_step) - # Orbax has just written new arrays into the live NNX variables, which is the one thing - # that makes the cached pure state wrong rather than merely old. Drop it; the next step - # re-reads the restored weights. + # Orbax has just written new arrays into the live NNX variables, so the cache is wrong + # rather than merely old. self._invalidate_pure_state() if restored_checkpoint_state.accumulated_metrics: @@ -1421,10 +1326,9 @@ def restore_checkpoint(self, **kwargs: Any) -> Any: rebuilt_losses = [wm] self._cached_losses = rebuilt_losses - # Checkpoints written before the denominator was tracked carry no value for it. The - # per-micro-batch losses just rebuilt above carry the very denominators that went into - # the saved gradients, so their sum is exactly what was lost. Only those count: any - # `_cached_losses` left over from before the restore belong to a different run. + # Checkpoints predating the denominator carry no value for it, but the losses rebuilt + # above carry the very denominators that went into the saved gradients. Only those: + # any `_cached_losses` from before the restore belong to a different run. if not restored_denominator and rebuilt_losses: denominator = jnp.float32(0.0) for cached_loss in rebuilt_losses: diff --git a/src/maxtext/training_engine/metrics.py b/src/maxtext/training_engine/metrics.py index 8e82420074..25adc54dc2 100644 --- a/src/maxtext/training_engine/metrics.py +++ b/src/maxtext/training_engine/metrics.py @@ -40,20 +40,9 @@ "tflops", ] -# How many completed step buffers stay resident before the oldest are evicted. -# -# Every buffer holds live device arrays -- one leaf per scalar metric, plus one per aux -# metric when the engine runs with `has_aux`, which for MoE/MTP configs is tens of leaves -- -# so nothing in the list is free. Nothing on the engine's own step path removes an entry -# either: `get_step_metrics` hands out the newest by reference, and only -# `get_metrics_history(clear_cache=True)` and `cleanup()` clear. A driver that reads the -# engine's own TensorBoard output rather than calling `get_metrics()` therefore never -# clears, and both HBM and -- because `save_checkpoint` serializes the whole retained -# history -- checkpoint size and save latency grow linearly in steps. -# -# Tunix v2 keeps exactly one prior step (`_prev_buffered_train_metrics` in -# `peft_trainer_v2.py`). This keeps a window instead so batched readers of -# `get_metrics_history` still work, while making the footprint constant in step count. +# Completed step buffers to keep resident. Each holds live device arrays and nothing on the +# engine's step path removes one, so unbounded history grows HBM, and checkpoint size with +# it. A window rather than Tunix's single prior step, so batched readers still work. _DEFAULT_MAX_BUFFERED_STEPS = 128 @@ -74,8 +63,7 @@ def __init__(self, max_buffered_steps: int = _DEFAULT_MAX_BUFFERED_STEPS): Args: max_buffered_steps: How many completed step buffers to retain; older ones are evicted - as new steps start. Pass 0 or a negative value to retain everything, which is only - safe when the driver drains the history itself. + as new steps start. Zero or less retains everything. """ self._metrics_buffer: list[abstract_engine.MetricsBuffer] = [] self._max_buffered_steps = max_buffered_steps @@ -107,8 +95,7 @@ def buffer_metrics( def _evict_old_buffers(self) -> None: """Drops the oldest step buffers once the retention window is full. - Only ever runs when a *new* step starts, so the buffer the current step is writing into - and the one `get_step_metrics` is about to hand the throttler are never the ones dropped. + Only runs when a *new* step starts, so the current step's buffer is never a candidate. """ if self._max_buffered_steps <= 0 or len(self._metrics_buffer) <= self._max_buffered_steps: return @@ -116,8 +103,6 @@ def _evict_old_buffers(self) -> None: oldest_dropped_id = self._metrics_buffer[0].id del self._metrics_buffer[:num_dropped] self._dropped_buffer_count += num_dropped - # Dropping metrics silently is the pattern that produced the fabricated 0.0 in the parity - # harness, so say so -- but once per window, not once per step. logging.log_every_n( logging.WARNING, "Metrics history is full at %d step(s); evicting buffers from step %s onwards " @@ -172,7 +157,7 @@ def get_metrics_history(self, clear_cache: bool = True) -> list[abstract_engine. The engine's own `get_metrics` returns only the most recent buffer, per the trainer contract. This is the accessor that keeps the history reachable -- the last - `max_buffered_steps` of it; see `_DEFAULT_MAX_BUFFERED_STEPS` for why it is a window. + `max_buffered_steps` of it. Args: clear_cache: Whether to reset cached metrics after retrieval. diff --git a/tests/post_training/integration/maxtext_engine_grpo_loss_test.py b/tests/post_training/integration/maxtext_engine_grpo_loss_test.py index 448a714576..6284277cd6 100644 --- a/tests/post_training/integration/maxtext_engine_grpo_loss_test.py +++ b/tests/post_training/integration/maxtext_engine_grpo_loss_test.py @@ -90,12 +90,9 @@ def _config(**overrides) -> pyconfig.HyperParameters: "learning_rate=1e-4", "micro_batch_size_to_train_on=2", "max_target_length=64", - # The KL assertion below compares log-probs produced by two different code paths: - # tunix's compute_per_token_logps for the reference, and the engine's own sharded - # forward pass for the policy. At the default bf16 matmul precision those disagree - # by ~2e-2 on log-probs near -12.6, and low_var_kl squares the difference into a - # KL of ~1e-4 -- noise, but large enough to swamp a real regression. fp32 matmuls - # bring it back under 1e-5; the model is tiny enough that the cost is invisible. + # The KL assertion compares log-probs from two code paths -- tunix's + # compute_per_token_logps and the engine's sharded forward -- which at bf16 disagree by + # ~2e-2 near -12.6, and low_var_kl squares that into a KL of ~1e-4. fp32 fixes it. "matmul_precision=highest", ] argv.extend(f"{k}={v}" for k, v in overrides.items()) @@ -178,15 +175,10 @@ class MaxTextEngineGrpoLossTest(absltest.TestCase): def test_grpo_loss_drives_a_training_step(self): cfg = _config() mesh = jax.sharding.Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes) - # The batch has to divide the data x fsdp device count, and only the mesh knows that - # number, so it cannot be hard-coded in `_config`. The engine traces its kernels under - # `nn_partitioning.axis_rules`, which is what makes MaxText's logical constraints on the - # activations real; a batch those devices cannot split is then padded out by XLA, and the - # padded lanes are all-zero sequences that mask themselves out of attention entirely. - # Their contribution comes back as NaN on the pad token's embedding row -- a finite loss - # and unusable gradients. The splash kernel asserts on exactly this ratio - # (`attention_op.py`, "Batch dimension should be shardable"); dot_product does not, so - # here it surfaced as a NaN instead of an error. + # The batch must divide data x fsdp, which only the mesh knows. A batch those devices + # cannot split is padded by XLA into all-zero sequences that mask themselves out of + # attention, returning NaN on the pad token's embedding row under a finite loss. Splash + # asserts on this ratio ("Batch dimension should be shardable"); dot_product does not. batch = int(mesh.shape["data"] * mesh.shape["fsdp"]) if cfg.micro_batch_size_to_train_on != batch: cfg = _config(micro_batch_size_to_train_on=batch) diff --git a/tests/post_training/unit/maxtext_engine_test.py b/tests/post_training/unit/maxtext_engine_test.py index 0e230b23a8..c3821b9b8c 100644 --- a/tests/post_training/unit/maxtext_engine_test.py +++ b/tests/post_training/unit/maxtext_engine_test.py @@ -190,9 +190,8 @@ def test_compiled_steps_publish_weights_and_non_param_state(self): t = maxtext_engine.MaxTextTrainingEngine(self.mock_config) def loss_fn(model, *_args, **_kwargs): - # Mutating a non-`Param` variable is what makes `new_rest` differ from the cached - # `rest`; scaling the loss by it makes a stale publish show up as a wrong gradient - # rather than only as a wrong counter. + # Mutating a non-`Param` makes `new_rest` differ from the cached `rest`; scaling the + # loss by it turns a stale publish into a wrong gradient, not just a wrong counter. model.calls.value = model.calls.value + 1.0 return ( abstract_engine.WeightedMetric( @@ -213,8 +212,7 @@ def loss_fn(model, *_args, **_kwargs): t.update() self.assertIsNotNone(t._params_pure, "the pure-state cache fell back to re-splitting the graph") - # `nnx.update` is the publish barrier: the live module the engine hands out, checkpoints - # and syncs weights from must track the cache, not lag behind it. + # `nnx.update` is the publish barrier: the live module must track the cache. self.assertEqual(float(t.model.calls.value), 2.0) self.assertGreater(float(np.abs(np.asarray(t.model.weights.value) - before).max()), 0.0) self.assertEqual(t.train_step, 2) @@ -662,8 +660,7 @@ def test_update_with_inflight_throttling(self): self.assertIsNone(metrics) if idx == 1: # Metrics for train_step=0, waited on through the update's gradient norm -- one - # scalar out of the same executable, not the state itself, whose buffers the next - # update donates away. + # scalar out of the same executable, not the state, whose buffers get donated away. self.assertIsNotNone(metrics) self.assertLen(computation, 1) self.assertEqual(jnp.shape(computation[0]), ())