Match Tunix peft_trainer_v2 performance in MaxTextTrainingEngine

[NuojCheng]

Description

Onboards the HBM and throughput work from tunix#1934 into src/maxtext/training_engine, closing the gap between MaxTextTrainingEngine and tunix/experimental/train/peft_trainer_v2.py.

On llama3.1-8b, ici_fsdp_parallelism=8, per_device_batch_size=1, max_target_length=1024, 4 gradient-accumulation micro-batches, 8x TPU7x:

	before	after	native train.py
update (GA=4)	3501 ms	1303 ms (2.69x)	1006 ms
peak HBM/device (GA=4)	27.97 GiB	15.03 GiB (1.86x less)	—
tokens/s (GA=4)	9 360	25 144	32 568
update (GA=1)	—	338.6 ms (318.6 fwd/bwd + 20.0 update)	325 ms
resident HBM/device	11.26 GiB	11.29 GiB	11.26 GiB

Six changes to the compiled step, in order of impact. Three further changes to the host step path are in their own section below.

1. Trace the kernels under the mesh and the logical axis rules (_sharding_ctx). The engine called jax.jit on its fwd/bwd and update kernels outside nn_partitioning.axis_rules, and those rules live in a context variable. A kernel traced without them sees an empty rule set, so every sharding.maybe_shard_with_logical inside the MaxText layers silently becomes a no-op and XLA is left to guess how to partition activations and gradients. It guesses badly: the same fwd/bwd measured 1012 ms traced outside the context against 581 ms inside it, with no numerical difference. train.py has always wrapped its own jit this way, which is why the standalone trainer never hit this. The context is entered around the call, not around jax.jit(...), since jit is lazy and the rules must be live when tracing actually happens.

2. Donate buffers. _compiled_update donates the train state, matching what maxtext_utils.get_functional_train_with_signature already does for the standalone trainer (donate_argnums = 0), and the accumulating fwd/bwd kernel donates the gradient accumulator. Params are deliberately not donated: micro_grads has the same shard-shape, dtype and sharding as the params, and JAX matches donations by shape rather than by position (jax/_src/interpreters/mlir.py:_set_up_aliases), so donating them would alias the weights straight into the gradient output. Gradients are not donated on the update either — every parameter-shaped output is already claimed by the incoming state, so the donation would have nothing to alias to and JAX would only warn.

3. Accumulate gradients inside the jit, with two kernels. The first micro-batch of an update has nothing to add to, so it returns its own gradients and the engine adopts them; every later micro-batch folds into that buffer in place and donates it. Tunix v2 calls the same split "non-persistent vs persistent" mode. The Python-side jax.tree.map(jnp.add, ...) it replaces 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.

4. Normalize as sum/sum rather than as a mean of per-micro-batch means. Gradients now accumulate unreduced (no 1/denominator per micro-batch) and update() divides once by the summed denominator, so the optimizer sees sum(grads)/sum(denom). The two forms agree only when every micro-batch carries the same token count; under sequence packing or a ragged RL rollout the mean of means silently overweights short micro-batches. This is what MaxText's own pre-train path already does (gradient_accumulation.py: accumulate xent_sum, divide once by the summed total_weights). It is also one fewer full pass over the gradient tree per micro-batch. The denominator is carried in checkpoint metadata, and rebuilt from the cached per-micro-batch losses for checkpoints written before it was tracked, so an intra-step save/restore round-trips.

5. Hand the inflight throttler one scalar off the updated weights instead of the whole train state. The queue keeps its entries alive until popped, so the old behaviour pinned three parameter trees (params + two optimizer slots) per entry; and after (2) those very buffers are donated by a later step, so an entry popped after that would raise Array has been deleted out of jax.block_until_ready. The marker is produced by the same executable as the weight update and reads from its result, so its readiness still means the update landed. Tunix v2 tracks the update's gradient norm for the same reason (_last_update_grad_norm); MaxText only computes a norm when clipping or spike-skipping is on, hence a slice instead.

6. Bound the metrics history to the last 128 steps (MetricsRecorder). _metrics_buffer grew one buffer of live device arrays per train step and nothing on the engine's step path ever removed one: get_step_metrics returns 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() never clears at all, and since save_checkpoint serializes the whole retained history, checkpoint size and save latency grew linearly in steps too. Eviction is audible (one warning per window) rather than a silent drop. Tunix v2 keeps exactly one prior step (_prev_buffered_train_metrics); a window is kept here so batched readers of get_metrics_history still work.

Residual gap at GA=4, and why it is not closed here

Per-micro-batch timing shows the first (318.7 ms) and accumulating (319-322 ms) kernels cost the same, so in-kernel accumulation costs about 1 ms, and at GA=1 the engine is within 4% of native. The remaining 1.29x at GA=4 is entirely native's jax.lax.scan over micro-batches letting XLA pipeline FSDP collectives across micro-batch boundaries — native's per-micro-batch cost drops from ~305 ms standalone to ~245 ms inside the scan. Traces confirm it: the engine spends 247 ms/update in unfused dynamic-update-slice on f32[32,512,14336]/f32[32,14336,512] (layout {2,0,1:T(8,128)}, each fed by a retiling %copy), where native has a single convert_dynamic-update-slice_fusion at 9.4 ms. Closing it needs a batched multi-micro-batch entry point, which peft_trainer_v2 also does not have — it drives one micro-batch per call — so the engine now matches Tunix's design.

The host step path: three further changes

The head-to-head below found the engine host-bound, not device-bound. Three changes address that; none of them touches a kernel.

Measured as a straight A/B on the two engine source files — this commit against its parent — on qwen3-0.6b, ici_tensor_parallelism=8 (all 8 devices), micro-batch 8, f32, optax.sgd(1e-5), no clipping, scan_layers=False, 23 steps, last 19 after warmup, 8x TPU7x. This is the third arm of the head-to-head harness, so the loss function and the arithmetic are literally Tunix's. --no-trace: see the note on why below.

ms/step, untraced	seq 1024, GA=1	seq 4096, GA=1	seq 1024, GA=4
changes (1)-(6) — median / mean / max	160.1 / 207.4 / 464.1	444.3 / 476.7 / 765.0	410.8 / 518.2 / 709.1
+ (7)-(9) — median / mean / max	89.8 / 90.0 / 93.4	373.9 / 373.9 / 374.7	333.8 / 349.2 / 639.4
speedup on the median	1.78x	1.19x	1.23x
speedup on the mean	2.30x	1.28x	1.48x

The saving is a fixed ~70 ms per optimizer step, not a percentage: 70.3 ms at seq 1024 and 70.4 ms at seq 4096, where the device work grew 4x. That is what "host-side" means — it shows up as 1.78x at 1024 tokens and 1.19x at 4096 for the same absolute win. It splits across the two dispatches as fwd_bwd 37.5 → 16.3 ms and update 122.2 → 73.8 ms, summing to 69.6 of the 70.3.

Turning the three on one at a time, same shape, so none of them is credited with another's win:

seq 1024, GA=1, untraced	ms/step	removed
changes (1)-(6)	160.1	—
+ (7) pure-state cache	111.5	48.6 ms
+ (8) dead mean_loss	108.2	3.3 ms
+ (9) deferred metric write	89.8	18.4 ms

(7) is the bulk of it and (8) is nearly free to skip but small — worth keeping as one fewer eager dispatch, not worth arguing about. Note (7) recovers 48.6 ms where the two nnx.split calls cost 71.7 ms of wall time in isolation: the rest was already overlapping device work, so removing it buys nothing. Measured, not assumed — this is why the section quotes an A/B rather than a microbenchmark.

The tail collapses too, and that is arguably the bigger result. The baseline's mean sits 30% above its median with a worst step of 464 ms against a 160 ms median; after the fixes mean and median agree to 0.2% and the worst step is 93.4 ms against 89.8. The jitter was the two nnx.split calls: each allocates a large short-lived object graph, twice per step, and the GC pauses that follow land on whichever step is unlucky. Removing the allocation removes the pauses.

Against the other two arms at the same shapes, untraced — the middle row is the control that isolates trainer from model, since it is the same MaxText model under Tunix's trainer:

ms/step, median	seq 1024	seq 4096
engine, changes (1)-(6)	160.1	444.3
engine, + (7)-(9)	89.8	373.9
MaxText model + Tunix PeftTrainer v2	83.4	367.0
Tunix model + Tunix PeftTrainer v2	69.1	330.3

Against PeftTrainer driving the identical model the engine goes from 1.92x slower to 1.08x at 1024 tokens, and from 1.21x to 1.019x at 4096. The residue is a near-constant 6.4–6.9 ms/step: the two nnx.update publish calls from (7), which are kept deliberately. The remaining distance to the Tunix model (14.3 ms at 1024, 36.7 ms at 4096) is a model difference, not a trainer one, and is what the rest of this section is about.

A caution on the numbers this PR quoted earlier. The first version of this table was traced, and read 283.2 → 92.3 ms/step, i.e. 3.07x. That overstated the win. The profiler charges per dispatch and the engine dispatches dozens of tiny eager ops per step where PeftTrainer dispatches two, so tracing taxes the baseline hardest and change (8) gets credited for removing tracing overhead as well as real work. The honest figure is the untraced 1.78x. --no-trace was added to all three arms in this PR so the distinction is not re-litigated by hand next time.

7. Carry the pure state across steps instead of re-splitting the module graph. fwd_bwd called nnx.split(model, nnx.Param, ...) and update called nnx.split(state), once each per step. Those are full traversals of the NNX graph of an unrolled 28-layer qwen3-0.6b, and they are not cheap host work. peft_trainer_v2 pays for its equivalent once, via nnx.cached_partial, which is not reusable here: that cache is consulted only inside an nnx transform's SplitContext.split (flax/nnx/graphlib.py:1824), and the engine's kernels are bare jax.jit. (Tunix's own fwd_bwd is not cached either — it is handed zero cached_args. Only its fused train_step benefits, which is the real reason the two designs diverge here.)

The engine now seeds a pure-pytree mirror of the model and of the train state at compile() time and republishes it from each kernel's output with nnx.split_state/nnx.merge_state, which walk the 490-leaf State rather than the graph. Timed on the real state with the device idle:

per call	nnx.split (graph)	nnx.split_state/merge_state (pure)
median	21.3 ms / 22.3 ms	0.83 ms / 1.05 ms
sustained mean, GC included	35.8 ms / 35.9 ms	0.84 ms

Two orders of magnitude on the median. The median understates what a step pays, because the splits allocate a large short-lived object graph twice per step and the GC pauses that follow land on whichever step is unlucky — hence the sustained row, and hence the baseline's 464 ms worst step above. In the loop the cache recovers 48.6 of those 71.7 ms; the remainder was already overlapping device work. Three things make the cache safe rather than merely fast:

• The two nnx.update calls stay. They cost ~12 ms/step combined and they are the publish barrier: engine.model, save_checkpoint and prepare_weight_sync all read the live NNX objects, while the cache is a detached snapshot of kernel output from the first update() onward. Dropping them would silently ship stale weights to an RL rollout, and no test on CPU or TPU would catch it.
• Every reconstruction goes through nnx.State.raw_mapping. State stores its children as plain dicts but wraps them in a State on __getitem__, so rebuilding from the wrapped views yields a tree that is equal key-for-key and leaf-for-leaf yet is a different pytree, one node deeper at every level — which jax.jit rejects as an in_shardings prefix mismatch naming neither the cause nor the site. _check_pure_state_reusable compares full treedefs for the same reason.
• The cache self-disables, with one warning, on any structural surprise — a state whose pure form has no model entry, an update output that does not repartition into the same parameters, or a fwd_bwd that returns a wider non-parameter state than the model was split into. That last one is not hypothetical: record_max_logits, distill_beta > 0 and multi-token prediction all sow nnx.Intermediates that come back in rest, and adopting the wider tree would leave the cache disagreeing with the rest_shardings the kernel was compiled against. The model/optimizer/state setters and restore_checkpoint invalidate it outright. The fallback is the old behaviour, not a wrong answer.

8. Do not compute mean_loss when nothing reads it. _update_kernel uses it only inside its skip_step_on_spikes branch, and that flag is read off the config at trace time and defaults off — so XLA had already dropped the argument, and update_in_shardings already declares None for that position. Producing it was not free: WeightedMetric.compute() is seven eager XLA launches (the eps and min_denom clamps plus a safe divide), i.e. seven dispatches per step feeding an input the executable does not contain.

9. Defer the metric write past the next dispatch. InflightThrottler.wait_for_next blocked on the popped computation and then ran write_metrics inline, which reduces each WeightedMetric on device and pulls the result to host with np.asarray. Those reduction ops are dispatched behind whatever is already queued, so running them there — before the caller dispatches the step it had just made room for — stalls the host on the entire backlog with nothing new running. A cProfile of the step path put 54.6 ms/step in blocking jax/_src/array.py:_value against an idle device; cProfile inflates that, and the clean A/B above credits this change with 18.4 ms/step. The write now happens at the top of the following add_computation, immediately after the caller dispatches, so the transfer hides behind live work. Metrics carry their own step id (MetricsBuffer.id), so nothing downstream can observe the one-dispatch delay, and wait_for_all flushes before returning so draining for a checkpoint or shutdown is unchanged. Tunix's throttler sidesteps this by not logging at all.

Not included

• Fusing fwd/bwd and update into one executable (Tunix's train_step, which it builds only at gradient_accumulation_steps == 1). MaxText clips on a global norm, so the whole gradient tree is live at the optimizer step in both paths and fusing does not shorten its peak liveness. The structural cost is real — two dispatches also mean two wait_for_next() calls against one 2-deep queue, so the engine pipelines one optimizer step deep where Tunix pipelines two — but it is now bounded: after (7)-(9) the engine runs 6.4 ms/step behind PeftTrainer driving the identical model at 1024 tokens and 6.9 ms behind it at 4096, so that is the size of the remaining prize. It is also unreachable through Tunix's TrainerWorker, which exposes only fwd_bwd and update, and maxtext_engine_test.py::test_update_with_inflight_throttling pins the two-dispatch protocol entry by entry. A fused path belongs as a third method with that test rewritten deliberately, not folded into these two.
• Moving the learning-rate schedule off device. It is one more eager dispatch per step, but reading it on the host means changing the metrics.py contract for marginal gain.
• A pre-existing bug, left for a separate change: diff_wrapper differentiates only aux["xent_sum"], so z_loss, mtp_loss, indexer_loss and moe_lb_loss are dropped from the gradient, while train.py adds all of them into the differentiated objective. For an MoE model the load-balancing loss currently contributes nothing through the engine path. Fixing it interacts with the sum/sum normalization above (those terms are already per-token-normalized), so it deserves its own PR rather than a silent change here.

Tests

All run on a v7x-8 VM (4 Ironwood chips, 8 JAX devices).

End-to-end parity, all 6 verifications passed (~65 min), for changes (1)-(6):

bash tests/end_to_end/tpu/test_training_engine_parity.sh

[1/6] verify_parity_with_train_py (Eager)                        PASSED
[2/6] verify_auxiliary_metrics_and_telemetry_parity (Eager)      PASSED
[3/6] verify_gradient_accumulation_parity (Eager)                PASSED
[4/6] verify_parity_with_train_py (JIT, llama3.1-8b)             PASSED
[5/6] verify_auxiliary_metrics_and_telemetry_parity (JIT, 8b)    PASSED
[6/6] verify_gradient_accumulation_parity (JIT, llama3.1-8b)     PASSED

Step 3 and step 6 are the ones that matter for change (4): they compare engine weights against the lax.scan baseline after each optimizer update, over 5 micro-batches with mask_prob=0.3, i.e. with per-micro-batch denominators that genuinely differ.

Re-verified for changes (7)-(9), with one caveat. Steps 1-3 still pass. Steps 4-6 no longer start on this VM: the llama3.1-8b arm reaches create_device_mesh with config.ici_parallelism is None and dies with AttributeError: 'NoneType' object has no attribute 'copy' (maxtext_utils.py:2202) before any engine code runs. That reproduces identically with all of (7)-(9) stashed, so it is unrelated to this PR — but it does mean the llama arms above have not been re-run. All six verifications were instead run against model_name=default, which exercises the same code paths including the compiled and donated update:

python3 tests/end_to_end/tpu/compare_training_engine.py model_name=default test_suite=eager_all  # [1-3] PASSED
python3 tests/end_to_end/tpu/compare_training_engine.py model_name=default test_suite=jit_all    # [4-6] PASSED

Unit tests (the engine tests are cpu_only, so they skip silently on a TPU host without the prefix):

JAX_PLATFORMS=cpu python -m pytest \
  tests/post_training/unit/maxtext_engine_test.py \
  tests/post_training/unit/maxtext_engine_e2e_test.py \
  tests/post_training/unit/maxtext_engine_constructor_test.py \
  tests/post_training/unit/router_replay_engine_test.py \
  tests/post_training/unit/tunix_adapter_test.py            # 57 passed
JAX_PLATFORMS=cpu python -m pytest tests/unit/grpo_nnx_test.py             # 11 passed
JAX_PLATFORMS=cpu python -m pytest tests/post_training/unit/metric_logger_abort_test.py  # 8 passed

Test expectations that pinned the old per-micro-batch scaling were updated to the unreduced accumulation, and new tests cover the checkpointed denominator, the scalar throttler marker, and the bounded metrics history.

For (7), maxtext_engine_test.py gains test_compiled_steps_publish_weights_and_non_param_state — the first CPU test to combine a model with non-Param state (nnx.BatchStat, mutated by the loss function), a real compile(), and two full fwd_bwd + update rounds. It asserts the cache survives both rounds, that the weights moved, and that the mutation reached engine.model — the last of which fails with 1.0 != 2.0 when the publish in _publish_model_rest is disabled, so the test is not vacuous.

Benchmark reproduction. The llama numbers come from driving the engine directly (engine.compile, then ga x fwd_bwd + one update, timed with block_until_ready and device.memory_stats()), against the same config run through train.py for the native column. The qwen3-0.6b tables in "The host step path" come from the three tests/end_to_end/tpu/perf_parity/qwen3_*_profile.py arms, each run with --tp 8 --no-trace and the shape flags in the table headers; the engine arm's _report_nnx_graph_cost additionally prints the per-step graph cost split into what a step still pays and what the cache now saves. The (1)-(6) rows are the same arms run against this commit's parent with only src/maxtext/training_engine/{maxtext_engine,inflight_throttler}.py reverted, so nothing but the engine differs.

pyink and pylint (10.00/10) are clean on all changed files.

Head-to-head against Tunix peft_trainer_v2

Added tests/end_to_end/tpu/compare_tunix_trainer.py and wrote the results up in
docs/reference/training_engine_tunix_parity.md.
It drives MaxTextTrainingEngine and peft_trainer_v2.PeftTrainer over the same model, the same algo_core.grpo_loss_fn, the same micro-batches and the same optax transformation, plus an independently computed sum-of-grads / sum-of-denoms reference gradient.

Qwen3-0.6B from gs://maxtext-model-checkpoints/qwen3-0.6b/2025-10-27/scanned/0/items, fsdp=8, batch 8, f32, gradient_clipping_threshold=0.0 (MaxText clips in its update kernel and Tunix never clips; leaving it on would mask the normalization difference being measured).

Numerics. At GA=1 the two are equivalent — identical loss -0.26746895909309387, gradients within rel_l2 2.98e-4, weight deltas agreeing to nine digits. With GA=4 and ragged micro-batches (denominators 64/16/40/8) they diverge:

	exact reference	MaxText	Tunix v2
accumulated denominator	128.0	128.0	4.0
gradient L2	40.1691955	40.1839846	52.7629719
rel_l2 vs. exact reference	—	0.006206	0.887693

MaxText matches the exact gradient to the same ~6e-3 jit-vs-eager noise floor it hits at GA=1; Tunix lands on mean-of-means, per its own # TODO(b/491970038): update denom for sequence packing. in _fwd_bwd_step. Trainer-vs-trainer is rel_l2 0.676. Worth noting the weight deltas still agree to four digits, because Adam normalizes per-element magnitude — the error is in gradient direction only, so it will not surface as a step-size anomaly. This is change (4) in this PR, measured against the external implementation.

Performance and HBM. The first version of this section quoted host wall clock at 16 tokens per example, from a loop that opened its xplane trace before any warmup. Both were wrong to lean on: a ~570 ms CompileAndLoad(jit__update_kernel) landed inside every MaxText trace, and 16 tokens puts the comparison nowhere near roofline. Re-measured with _WARMUP=2 untimed updates outside the trace context, one trainer per process, and at two shapes. Wall clock is from untraced runs; TPU-busy is total /device:TPU:0 XLA-module time from the traces. All of this section predates changes (7)-(9), which move the MaxText column only.

GA=1, medians over 8 iterations, untraced, one trainer per process	MaxText	Tunix v2
total/update, 16 tokens	34.7 ms	64.0 ms	1.85x
total/update, 1024 tokens	59.7 ms	74.7 ms	1.25x
total/update, 2048 tokens	84.5 ms	102.1 ms	1.21x
total/update, 4096 tokens	189.1 ms	203.9 ms	1.08x
total/update, 8192 tokens	479.6 ms	493.2 ms	1.03x
peak HBM/device, GA=1 (flat 16 → 8192 tokens)	1.77–1.79 GiB	8.64–8.65 GiB	4.9x
peak HBM/device, GA=2	1.81 GiB	4.79 GiB	2.6x
metrics recorded/step	22	2

On step time alone the two trainers are equivalent at production sequence lengths. MaxText's lead is a short-sequence effect that amortizes away: fwd_bwd goes 2.75x → 1.46x → 1.34x → 1.11x → 1.04x across the sweep, and at 8192 tokens the totals are within run-to-run noise. Do not quote the small-shape numbers as a general result.

The mechanism only holds at small shapes. Tunix's nnx.jit carries no out_shardings, so its grads return P() at 2.22 GiB/device against MaxText's sharded 0.278 GiB (params and optimizer slots shard identically), and the resulting extra memory traffic is roughly constant: +14.7 GiB at 16 tokens, +13.0 at 2048 — but −0.7 at 4096 and −12.4 at 8192, where MaxText accesses more and is still marginally ahead. By then both are compute-bound on identical arithmetic (1080.7 vs 1082.4 GFLOP at 1024 tokens), so the times converge.

Peak HBM is the durable difference. It is flat across the whole sweep — a 512x sequence increase moves neither trainer — so unlike step time it does not amortize. The persistent accumulator from tunix#1934 improves it but does not close it: at GA=2 Tunix's accumulator is sharded and peak drops 8.65 → 4.79 GiB, while bytes-accessed is unchanged at 22.67 GiB and fwd_bwd stays at 71.4 ms. Nothing further to port — the explicit in_shardings/out_shardings this PR keeps on _fwd_bwd_kernel are what v2 lacks.

update is the one phase Tunix wins, by ~6 ms at both shapes. MaxText's 14.0 ms is 0.71 ms on device, ~2.8 ms in MetricsRecorder._record_metric (measured by no-op'ing it), ~1.7 ms in nnx.split, and ~8.8 ms of other host work. Tunix is 8.2 ms wall clock against 0.86 ms on device, so both are host-bound; MaxText is just more so. This paragraph is what changes (7)-(9) target — the nnx.split line is now cached away, and the "other host work" turned out to be dominated by the metric fetch in (9). The scan-free qwen3-0.6b arm quantifies it end to end in the table above.

xprof's step statistics are not usable for MaxText, which is worth knowing before comparing these traces to any others. Over 8 steps MaxText executes 424 XLA modules where Tunix executes 16: _record_metric calls jnp.atleast_1d/jnp.append once per metric per micro-step, ~51 tiny eager dispatches per step. Step detection splits on module boundaries, so MaxText's reported step time is computed over 424 fake steps with a 0.012 ms median, and the device row renders as a picket fence next to Tunix's two clean bars. Read total TPU-busy time instead. The same dispatches make MaxText pay disproportionate per-dispatch tracing overhead — traced, its update reads 27.0 ms against an untraced 14.0 ms, while Tunix goes 8.2 → 18.5 ms. Change (8) removes 7 of the ~51 per step; the picket fence remains, and shrinking it properly means batching _record_metric, which is a metrics.py change rather than an engine one.

GRPO integration test: tests/post_training/integration/maxtext_engine_grpo_loss_test.py — 1 passed in 50.32 s on the real checkpoint. Note it never calls engine.compile(), so it measures the eager path (237 ms/update vs 14 ms).

xprof traces, steady-state loop only, no compilation inside the window, with StepTraceAnnotation boundaries and fwd_bwd/update regions:

gs://chengnuojin-xprof/maxtext-vs-tunix-2026-08-31/maxtext_ga1_bs8_seq16/
gs://chengnuojin-xprof/maxtext-vs-tunix-2026-08-31/tunix_ga1_bs8_seq16/
gs://chengnuojin-xprof/maxtext-vs-tunix-2026-08-31/maxtext_ga1_bs8_seq1024/
gs://chengnuojin-xprof/maxtext-vs-tunix-2026-08-31/tunix_ga1_bs8_seq1024/
gs://chengnuojin-xprof/maxtext-vs-tunix-2026-08-31/maxtext_ga2_bs8_seq1024/
gs://chengnuojin-xprof/maxtext-vs-tunix-2026-08-31/tunix_ga2_bs8_seq1024/
gs://chengnuojin-xprof/maxtext-vs-tunix-2026-08-31/raw_results/

Each holds plugins/profile/<ts>/t1v-n-c9d27794-w-0.xplane.pb.

One bug found while measuring, not fixed here. MaxText logs a different loss than it optimizes: MetricsRecorder._record_metric appends one entry per micro-step, no aggregation_fn is registered for "loss", so MetricsLogger._process_metrics reduces it with np.mean — mean-of-means, the normalization change (4) deliberately avoids. On the ragged GA=4 batch the logged loss is 0.08643750101327896 while the gradient uses -0.03790009766817093; the sign disagrees. Gradients are unaffected. Left for a separate change since it is reporting-only and predates this PR.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

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

[gemini-code-assist]

Code Review

This pull request optimizes gradient accumulation and memory footprint in the MaxText training engine. It defers gradient scaling to the update step by tracking unreduced gradients and their accumulated denominators, introduces a split-kernel compilation strategy to optimize buffer allocations, and limits the metrics history buffer to prevent unbounded memory growth. The review feedback highlights potential issues with non-array leaves in the update marker logic and warns against implicit truthiness checks on the restored denominator that could fail when its value is legitimately zero.

[gemini-code-assist]

Using jnp.size(leaf) in a list comprehension can fail if new_state_pure contains non-array leaves such as None (which is common in optax states). In JAX, None is treated as a leaf, and calling jnp.size(None) will raise a TypeError. It is safer to check if the leaf has a size attribute before accessing it.

Suggested change

	leaves = [leaf for leaf in jax.tree.leaves(new_state_pure) if jnp.size(leaf) > 0]
	leaves = [leaf for leaf in jax.tree.leaves(new_state_pure) if hasattr(leaf, "size") and leaf.size > 0]

[gemini-code-assist]

Using truthiness checks like restored_denominator if restored_denominator else 0.0 can be problematic if restored_denominator is legitimately 0.0 (e.g., a step with zero tokens). It is safer and more explicit to check is not None to handle the case where the denominator is missing versus when it is zero.

Suggested change

	self._accumulated_denominator = jnp.float32(restored_denominator if restored_denominator else 0.0)
	self._accumulated_denominator = jnp.float32(restored_denominator if restored_denominator is not None else 0.0)

References

1. Comparisons to singletons like None should always be done with is or is not, never the equality operators or implicit truthiness when 0 is a valid value. ^(link)

[gemini-code-assist]

Using not restored_denominator will evaluate to True if restored_denominator is 0.0. This will trigger the rebuilding logic even if the denominator was successfully restored as 0.0. Check restored_denominator is None instead to correctly identify if the denominator was missing from the checkpoint.

Suggested change

	if not restored_denominator and rebuilt_losses:
	if restored_denominator is None and rebuilt_losses:

References

1. Comparisons to singletons like None should always be done with is or is not, never the equality operators or implicit truthiness when 0 is a valid value. ^(link)

[codecov]

Codecov Report

❌ Patch coverage is 77.09497% with 41 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/training_engine/maxtext_engine.py	73.37%	29 Missing and 12 partials ⚠️

📢 Thoughts on this report? Let us know!

[A9isha]

Two notes on change (4), both narrow — the normalization itself is right, and the mask_prob=0.3 parity run is the evidence that matters.

1. The scale can be computed once instead of per leaf.

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

has_weights is traced, so XLA can't fold the select away — this is a divide plus a broadcast select over every element of the gradient tree. Hoisting it to a scalar gives the same result for one select and a multiply:

scale = jnp.where(has_weights, 1.0 / safe_denominator, 0.0)
grads = jax.tree.map(lambda g: g * scale.astype(g.dtype), accumulated_grads)

That expression is also exactly utils.WeightedMetric.compute_scale() with eps=None, min_denom=None, so calling it directly on a WeightedMetric built over the summed denominator would work too, and would pick up any eps/min_denom the loss declares — tunix's own _loss_fn sets eps=1e-8. Only matters at denominator < 1, so a scalar float in checkpoint metadata is a reasonable trade either way.

2. The accumulator dtype should probably be pinned to float32.

_fwd_bwd_kernel still casts micro_grads to config.grad_dtype before the in-kernel accumulation. That cast predates this PR, but its meaning changes here: the values being summed are now unreduced, so they're larger than the pre-scaled gradients they replace by roughly the micro-batch's denominator, and they're summed across micro-batches before anything divides them down. A grad_dtype that was wide enough for pre-scaled gradients isn't obviously wide enough for this.

tunix's GradientAccumulator treats these as separate knobs for that reason — accumulator_dtype defaults to jnp.float32 "to prevent low-precision underflow and rounding errors during multi-step accumulation", and get() casts back to the native parameter dtypes on the way out. The equivalent here is to accumulate in fp32 and apply the grad_dtype cast in _update_kernel after the division.

grad_dtype defaults to float32 so this is a no-op for everyone today, and compare_tunix_trainer.py pins grad_dtype=float32, which means the reduced-precision path isn't covered by the parity runs. Happy to be told it's out of scope — but if so it's worth a line in the "Not included" section, since after this change the setting is riskier than it was.