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.

This is the one behaviour change in the PR, as opposed to a timing change, and it follows directly from the constraints becoming real: micro_batch_size_to_train_on must now be a multiple of data x fsdp. A batch those devices cannot split is 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, giving unusable gradients under a finite loss. MaxText's splash kernel already asserts on exactly this ratio ("Batch dimension should be shardable among the devices in data and fsdp axis", attention_op.py); dot_product does not, so there it surfaces as a NaN rather than an error. Callers relying on an undersized micro-batch are the group affected; in practice such a batch was already leaving devices idle. This is how it first showed up -- the GRPO integration test hard-coded a batch of 2 and ran on a 4-device host.

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.

Rebased onto 2fbedac89. 202a89ab8 landed the always-on gradient norm on main while this was in review; the branch now keeps main's placement — before clipping, which is where Tunix's optax.global_norm sits — rather than this branch's post-clipping one, and adds the float32 reduction and the throttler's use of it on top. Every measurement below runs unclipped, so no number changes. Main's LoRA rejection and perplexity metric are carried through unchanged.

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.

On a second model the same fix is worth almost nothing, and that is the honest framing. Re-run on qwen3.5-35b-a3b — 2 KV heads, so --tp 8 is rejected and unscanned it does not fit in HBM, leaving --tp 2 --scan as the shape that runs on 8 devices:

ms/step, untraced, seq 1024, GA=1	median	mean	max
changes (1)-(6)	2322.5	2339.6	2641.8
+ (7)-(9)	2314.0	2314.0	2314.6
speedup	1.004x	1.011x
same model + Tunix PeftTrainer v2	2303.1	2303.4	2307.1

8.5 ms on the median, against 71.8 ms on qwen3-0.6b. Two things set what the fix is worth, and neither is model size:

• How big the NNX graph is. The cost removed is two graph traversals per step, and scanning collapses the decoder stack into one stacked node set — so the larger model has the smaller graph, 70 parameter leaves against qwen3-0.6b's 310 unscanned, and its nnx.split costs 5.0/4.7 ms against 21.3/22.3 ms.
• How long the device is busy. Host work that fits inside the device step is free. qwen3-0.6b is ~82 ms of TPU-busy, so ~70 ms of host work was mostly exposed; qwen3.5-35b-a3b is 2.3 s, so nearly all of it hides.

The tail still improves where the median does not: the 35b's worst step drops 2641.8 → 2314.6 ms and its mean converges onto its median — 327 ms of jitter removed from a model whose median moved 8.5 ms. And it reaches 1.005x of PeftTrainer driving the identical model, from 1.008x. So: a large win on unscanned models with short steps, a jitter fix everywhere else, and nothing that regresses.

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 tables in "The host step path" come from the three tests/end_to_end/tpu/perf_parity/qwen3_*_profile.py arms, each run with --no-trace and the shape flags in the table headers (--model qwen3-0.6b --tp 8, or --model qwen3.5-35b-a3b --tp 2 --scan; the tunix arm refuses --model/--scan, having only the one architecture); 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, whose §9 carries the (7)-(9) A/B on both models.
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.

Trainer parity: MaxTextTrainingEngine vs tunix PeftTrainer v2

Reproduced from tests/end_to_end/tpu/perf_parity/README.md on this branch, which supersedes the
2026-08-31 performance and HBM numbers this section used to quote. Those, and the (7)-(9)
host-path A/B, remain in docs/reference/training_engine_tunix_parity.md.

Consolidated conclusions from two studies run on a v7-8 (Ironwood) host, 2026-09-02 and
2026-09-03. Both trainers drive the same MaxText model in every arm, so what varies is
the trainer, the mesh and the sharding mode — never the implementation of the network.

Two models, chosen because they stress opposite things:

	qwen3.5-35b-a3b	qwen3-0.6b
Kind	sparse MoE, 256 routed + 1 shared expert, 40 layers	dense, 28 layers
Step time	~0.7–2.3 s — device-bound, host cost invisible	~50–470 ms — host and collective cost visible
What it tests	mesh choice (TP vs EP), MoE kernels, GA capability	sharding mode, optimizer-state sharding, host overhead
Detailed report	RESULTS-qwen35-35b-ep-20260902.md, RESULTS-qwen35-35b-20260902.md	RESULTS-qwen3-0p6b-zero1-20260903.md

Every command needed to reproduce every number is in Appendix A.

Conclusions

1. On raw throughput the two trainers are at parity, and the gap is a fixed dispatch, not a
   scaling one. On the 35b model they land within 0.5% of each other at every mesh
   (670.3 vs 666.6 ms at ep=8; 2313.8 vs 2303.1 at tp=2), both above 99% device utilization.
   PeftTrainer's lead is its single fused jit__train_step against the engine's split
   fwd_bwd + update — a constant ~16 ms, which is 0.7% of a 2.3 s step and a third of a
   61 ms one.

2. Under explicit sharding the engine wins decisively, and the win is the engine's, not the
   mesh mode's. This is the largest trainer-attributable difference measured. Setting
   shard_mode=explicit on qwen3-0.6b at GA=8 makes the engine 1.50x faster (371.7 →
   247.4 ms) and PeftTrainer 1.31x slower (356.8 → 465.7 ms) — the same flag, the same
   mesh, opposite signs. Same-mesh, same-mode, GA=8: 247.4 ms against 465.7 ms, 1.88x.
   Under Explicit axes the engine's cotangents come out unreduced, so the data-parallel
   all-reduce moves out of every micro-batch into update(); tunix's step has no such path
   and simply pays the collectives literally where GSPMD had been optimizing them.

3. The decisive difference is capability, not speed: PeftTrainer v2 cannot do gradient
   accumulation on the 35b model at the tunix revision MaxText pins. GA=2, 4 and 8 each OOM
   with a byte-identical 161.94 G against 94.74 G available. The engine has no such wall and
   its per-micro-batch cost improves with depth. Tunix 44a35eeaf fixes it with one line
   (GradientAccumulator.reset() using v[...] * 0 instead of jnp.zeros_like, which loses
   sharding inside a traced function and materialises the full 129 G tree). Recommend
   bumping src/dependencies/extra_deps/post_train_github_deps.txt past that commit.

4. Mesh choice dominates trainer choice by an order of magnitude. Moving the 35b model
   from fsdp=4 x tp=2 to ep=8 --ring-of-experts --ragged-sort is worth 3.45x on both
   trainers alike. Nothing trainer-side in this study is worth more than 1.9x. Tune the mesh
   before arguing about trainers.

5. For memory on small models, FSDP beats Zero-1 and it is not close. 2.04 G against
   8.70 G at qwen3-0.6b/GA=8. Zero-1 shards the two Adam moments; FSDP shards the parameters
   and gradients too, and at 0.6 B those are most of the residual. Zero-1's advantage is that
   it adds no collectives to forward and backward — not that it saves more bytes. The engine
   forbids combining them.

6. Neither trainer currently offers both speed and memory. At qwen3-0.6b/GA=8 the engine's
   explicit path is 1.36x faster than PeftTrainer on FSDP but uses 6.2x the HBM — 4.3x even
   with Zero-1 on top. That is the open gap this comparison leaves.

Environment

Host	8 x TPU7x (v7-8 Ironwood, 4 chips / 8 JAX devices), single process
JAX	0.11.1
Shape	micro-batch 8 x seq 1024, f32 compute and weights, remat_policy=none
Wall clock	23 steps, median of the last 19, --no-trace
Device time	XLA Modules line, per-execution, from a separate --steps 6 traced run
Peak HBM	peak_bytes_in_use off the TPU allocator, max over the 8 devices, --no-trace runs
Revisions	35b: PR #5060 at 1b75c2479. 0.6b: 18f4a2332 (PR #5060 scripts + PR #5099 engine)

remat_policy=none, scan_layers matched, and dtype=float32 on both arms throughout —
tunix's ModelConfig defaults to f32 and its RematConfig.NONE does not rematerialize, so
leaving MaxText's defaults in place would have MaxText winning on numerics rather than on
implementation. attention: autoselected is deliberately not equalized; the chosen kernel
is part of what is compared, and it is logged.

---

1. qwen3.5-35b-a3b — MoE, mesh-bound

40 layers, 256 routed + 1 shared expert, emb 2048, 16 query / 2 KV heads, head_dim 256,
vocab 248320. Scanned. optax.sgd(1e-5), constant schedule, no clipping.

--tp 8 is illegal on this model: _validate_kv_head_sharding requires
num_kv_heads % tp == 0 and there are 2 KV heads, so fsdp=4 x tp=2 is the widest legal
tensor-parallel shape on 8 devices.

1.1 Trainer vs trainer, same mesh

Mesh	GA	Engine	PeftTrainer v2	Δ
ep=8 + ring + ragged	1	670.3 ms	666.6 ms	Peft 0.55%
ep=8 + ring + ragged	8	5343.0 ms	5278.0 ms ¹	Peft 1.2%
fsdp=4 x tp=2	1	2313.8 ms	2303.1 ms	Peft 0.46%
fsdp=4 x tp=2	8	18405.3 ms	OOM ²	engine only

¹ tunix 07dbe293 (head). OOM at the pinned c4ec573 — see §1.3.
² OOM at the pinned revision; not run at head.

Device utilization is 99.2–99.8% on every row. At 2.3 s of device work per step every
host-side difference between the trainers hides; the 0.5% that remains is the engine's second
dispatch (jit__update_kernel, 16.0 ms at ep=8 and 16.46 ms at tp=2).

Per-execution device cost, ep=8, GA=1:

Arm	Modules	Total
engine	jit_first_kernel 648.9 + jit__update_kernel 16.0	664.9 ms
PeftTrainer v2	jit__train_step 663.0	663.0 ms

1.2 The mesh is worth 3.45x, and the two trainers track each other exactly

Every row --scan --seq 1024 --ga 1 --no-trace, median of 19 steps.

Mesh and flags	Engine	PeftTrainer v2	vs tp=2	vs row above
--tp 2 (fsdp=4 x tp=2)	2313.8	2303.1	1.00x	—
--tp 2 --ep 4 --ring-of-experts --ragged-sort	1726.9	1704.3	1.34x	—
--ep 8 (no MoE flags)	1063.8	1059.7	2.17x	2.17x
--ep 8 --ragged-sort	1047.6	1043.4	2.21x	1.02x
--ep 8 --ring-of-experts	720.1	717.4	3.21x	1.48x
--ep 8 --ring-of-experts --ragged-sort	670.3	666.6	3.45x	1.07x

• The mesh is the larger half. A bare ep=8 with no kernel flags is already 2.17x. TP
  splits a 2048-wide embedding and 2 KV heads and pays all-reduces for it; EP splits 256
  experts, which is what this model actually has a lot of.
• Ragged sort is a multiplier on ring-of-experts, not independently useful. 1.02x alone,
  1.07x on top of the ring. layers/moe.py selects the ring_ragged_sort kernels only under
  use_ragged_sort and use_ring_of_experts.
• Splitting the mesh across both axes is the worst of both. --tp 2 --ep 4 with both
  flags on (1726.9 ms) is slower than a bare --ep 8 with none (1063.8 ms).

Wall and device speedups agree to two decimals (3.45x / 3.47x), which is what says the gain is
work removed rather than overhead rearranged.

1.3 Gradient accumulation: a capability difference

At ep=8 + ring + ragged, median ms/step:

GA	Engine	per micro	PeftTrainer, tunix c4ec573 (pinned)	PeftTrainer, tunix 07dbe293 (head)	per micro
1	670.3	670.3	666.6	667.3	667.3
2	1338.8	669.4	OOM	1339.6	669.8
4	2672.3	668.1	OOM	2651.8	663.0
8	5343.0	667.9	OOM	5278.0	659.8

jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Ran out of memory on HBM, the total memory
required for HLO temporaries (161.94G) exceeds available HBM (94.74G).
HLO module: jit__update_step.

The figure is byte-identical at GA=2, 4 and 8, which rules out the obvious reading that N
micro-batches are live at once. The allocation is parameter-shaped, not depth-shaped:
PeftTrainer._is_single_microstep() is true only at gradient_accumulation_steps == 1, and
only that path skips allocating the accumulator.

The cause is one line, tested directly. Reverting only tunix 44a35eeaf on top of head —
leaving the other 208 commits in place — brings the failure back byte for byte. reset() runs
inside the traced _update_step, and there jnp.zeros_like does not carry its operand's
sharding, so XLA materialises the full unsharded tree:

129.12 G  full parameter tree, unsharded, from the traced zeros_like
+ 16.14 G  the sharded accumulator itself
+ 16.14 G  one more sharded parameter-tree copy
= 161.40 G   measured: 161.41 G

Once fixed, the two trainers are at parity at every depth and PeftTrainer pulls 1.2% ahead
by GA=8, because it amortizes its update slightly better (667.3 → 659.8 per micro against the
engine's 670.3 → 667.9).

---

2. qwen3-0.6b — dense, sharding-bound

28 layers, emb 1024, 16 query / 8 KV heads, head_dim 128, mlp 3072, vocab 151936, tied
embeddings; ~0.6 B params ≈ 2.4 G in f32. Unscanned. adamw (b1=0.9, b2=0.95, eps=1e-8,
wd=0.1, lr 1e-5 constant, no clipping), matched term for term across both trainers — base.yml's
constants, not optax's defaults.

adamw and --dp 8 rather than the sgd/--fsdp 8 the 35b arms use, because Zero-1 shards
parameter-shaped optimizer state over the data axis: it is vacuous under SGD and mutually
exclusive with FSDP.

2.1 Trainer vs trainer, all six arms

Median ms/step and peak HBM per device (of 101.72 G).

Arm	shard_mode	GA=1 ms	GA=1 HBM	GA=8 ms	GA=8 HBM
engine, dp=8	auto	61.0	9.95 G	371.7	12.73 G
engine, dp=8	explicit	62.1	9.96 G	247.4	12.74 G
engine, dp=8 + Zero-1	explicit	66.1	7.76 G	250.6	8.70 G
PeftTrainer, dp=8	auto	56.0	10.00 G	356.8	12.35 G
PeftTrainer, dp=8	explicit	64.0	10.47 G	465.7	12.39 G
PeftTrainer, fsdp=8	auto	46.4	1.74 G	337.4	2.04 G

The same-mode comparison is the one that matters, and it flips sign with the mode:

dp=8, GA=8	Engine	PeftTrainer v2	Winner
shard_mode=auto	371.7 ms	356.8 ms	Peft, 1.04x
shard_mode=explicit	247.4 ms	465.7 ms	engine, 1.88x

2.2 Where the time goes

Per-execution device cost, --steps 6 traces. A GA=8 step is one first_kernel, seven
accum_kernel and one _update_kernel on the engine; eight _fwd_bwd_step and one
_update_step on PeftTrainer.

Arm, GA=8	per micro-batch	update	Step total
engine, auto	43.57 ms	5.95	354.19
engine, explicit	25.03 ms	27.81	226.19
engine, explicit + Zero-1	25.07 ms	30.71	229.51
PeftTrainer, auto	42.73 ms	6.82	348.66
PeftTrainer, explicit	56.04 ms	6.89	455.21
PeftTrainer, fsdp=8	40.79 ms	1.69	328.01

The all-reduce is moved, not removed, and the same ~20 ms shows up three times. Going
auto → explicit on the engine at GA=8: the first kernel drops 20.08 ms, each of seven
accumulates drops 18.54 ms, and the update rises 21.86 ms. That is one all-reduce of the 2.4 G
f32 gradient tree over 8 devices, which this host does in about 20 ms — paid eight times under
auto, once under explicit. Extrapolated from these kernels it is a 1.78 ms loss at GA=1
(measured at 4.63 ms in the GA=1 trace, which runs slower throughout) and pays from GA=2 on.

PeftTrainer gets no such benefit; Explicit costs it 31%. Its per-micro-batch step goes
42.73 → 56.04 ms while its update is unchanged (6.82 → 6.89). Nothing moves out of the
micro-batch — the collectives stay where they were and get more expensive, because Explicit
axes stop GSPMD from rearranging them.

Zero-1 costs +2.9 ms, entirely in update(), and this is provable rather than inferred.
The Zero-1 arm's jit_first_kernel and jit_accum_kernel carry identical HLO program hashes
to the explicit control's (1445966014043162594 and 614701509375139023), at both GA depths.
Only jit__update_kernel differs (8707565451094719151), by +2.90 ms at GA=8 and +2.85 ms at
GA=1 — the price of replacing the gradient all-reduce with a reduce-scatter plus an all-gather
of the updated parameters.

FSDP does not supply a deferral either. It cuts the per-micro cost only 42.73 → 40.79 ms:
sharding the parameters swaps the gradient all-reduce for a parameter all-gather plus a
gradient reduce-scatter, about the same traffic, still once per micro-batch. The engine's
explicit path runs the same micro-batch in 25.03 ms, 1.63x faster, by not running the
collective there at all.

2.3 Where the memory goes

At ~0.6 B params in f32: parameters 2.4 G, Adam mu + nu 4.8 G, gradients 2.4 G.

	GA=1	GA=8	What is sharded
engine, dp=8 auto	9.95 G	12.73 G	nothing — params, moments and grads all replicated
engine, + Zero-1	7.76 G	8.70 G	the two moments, over data (4.8 G → 0.6 G)
saving	2.20 G (22%)	4.04 G (32%)
PeftTrainer, fsdp=8	1.74 G	2.04 G	params, grads and moments, over fsdp

Zero-1's predicted saving is 4.8 − 0.6 = 4.2 G; it delivers 4.04 G of that at GA=8 but only
2.20 G at GA=1, because with no accumulator live the high-water mark is set partly by the
activation peak instead. Zero-1's saving is capped by whatever else is at the peak, which
is why it looks better the deeper the accumulation goes.

FSDP's 2.04 G is not a better-tuned version of the same idea — it shards a strictly larger set
of tensors.

---

3. Which to use

Situation	Use	Why
Large MoE, GA=1	either	Parity to 0.5%; pick on ergonomics. Spend the effort on --ep 8 --ring-of-experts --ragged-sort instead — 3.45x.
Large MoE, GA>1, tunix at MaxText's pin	engine	PeftTrainer OOMs outright. Capability, not speed.
Large MoE, GA>1, tunix past 44a35eeaf	either	Parity; PeftTrainer 1.2% ahead at GA=8.
Small dense, GA=1	PeftTrainer	46.4 ms on FSDP against the engine's 61.0. The engine's ~19 ms/step of host-side NNX graph work is a third of a step this size.
Small dense, GA≥2, throughput first	engine + --explicit	247.4 ms against 337.4 (Peft/FSDP) and 465.7 (Peft/explicit). The deferred all-reduce is engine-only.
Small dense, memory first	PeftTrainer + FSDP	2.04 G against Zero-1's 8.70 G.
Need optimizer-state sharding without FSDP	engine + --zero1	1.3% for 32% of HBM at GA=8. Requires --explicit, which brings the deferral with it.

Do not read the small-model host overhead as a general engine defect: on the 35b model the
same NNX graph work is 3.1 ms against a 714 ms step, 0.4%.

---

4. Traces

All 24 .xplane.pb files, full paths. -engine is MaxTextTrainingEngine; -maxtext is the
PeftTrainer v2 arm (the MaxText model under tunix's trainer).

qwen3-0.6b, Zero-1 and explicit sharding (2026-09-03)

gs://chengnuojin-xprof/zero1-qwen3-0p6b-20260903/qwen3-0.6b-engine-dp8-adamw/plugins/profile/run/t1v-n-c9d27794-w-0.xplane.pb
gs://chengnuojin-xprof/zero1-qwen3-0p6b-20260903/qwen3-0.6b-engine-dp8-adamw-explicit/plugins/profile/run/t1v-n-c9d27794-w-0.xplane.pb
gs://chengnuojin-xprof/zero1-qwen3-0p6b-20260903/qwen3-0.6b-engine-dp8-adamw-zero1/plugins/profile/run/t1v-n-c9d27794-w-0.xplane.pb
gs://chengnuojin-xprof/zero1-qwen3-0p6b-20260903/qwen3-0.6b-engine-ga8-dp8-adamw/plugins/profile/run/t1v-n-c9d27794-w-0.xplane.pb
gs://chengnuojin-xprof/zero1-qwen3-0p6b-20260903/qwen3-0.6b-engine-ga8-dp8-adamw-explicit/plugins/profile/run/t1v-n-c9d27794-w-0.xplane.pb
gs://chengnuojin-xprof/zero1-qwen3-0p6b-20260903/qwen3-0.6b-engine-ga8-dp8-adamw-zero1/plugins/profile/run/t1v-n-c9d27794-w-0.xplane.pb
gs://chengnuojin-xprof/zero1-qwen3-0p6b-20260903/qwen3-0.6b-maxtext-dp8-adamw/plugins/profile/run/t1v-n-c9d27794-w-0.xplane.pb
gs://chengnuojin-xprof/zero1-qwen3-0p6b-20260903/qwen3-0.6b-maxtext-ga8-dp8-adamw/plugins/profile/run/t1v-n-c9d27794-w-0.xplane.pb
gs://chengnuojin-xprof/zero1-qwen3-0p6b-20260903/qwen3-0.6b-maxtext-ga8-dp8-adamw-explicit/plugins/profile/run/t1v-n-c9d27794-w-0.xplane.pb
gs://chengnuojin-xprof/zero1-qwen3-0p6b-20260903/qwen3-0.6b-maxtext-adamw/plugins/profile/run/t1v-n-c9d27794-w-0.xplane.pb
gs://chengnuojin-xprof/zero1-qwen3-0p6b-20260903/qwen3-0.6b-maxtext-ga8-adamw/plugins/profile/run/t1v-n-c9d27794-w-0.xplane.pb

The last two paths are the fsdp=8 arm despite not saying so. RunSpec.tag() only annotates
non-default mesh fills and --fsdp 8 is the default on 8 devices, so the FSDP arm gets no
mesh suffix while the --dp 8 arms get -dp8. Cross-check on device time if in doubt: the
fsdp arm's GA=1 jit__train_step is 41.07 ms, the dp arm's 50.83 ms.

qwen3.5-35b-a3b, expert parallelism (2026-09-02)

gs://chengnuojin-xprof/ep-parity-qwen35-35b-20260902/qwen3.5-35b-a3b-engine-scan-fsdp1ep8-roe-rsort/plugins/profile/2026_09_02_16_34_54/t1v-n-c9d27794-w-0.xplane.pb
gs://chengnuojin-xprof/ep-parity-qwen35-35b-20260902/qwen3.5-35b-a3b-maxtext-scan-fsdp1ep8-roe-rsort/plugins/profile/2026_09_02_16_37_15/t1v-n-c9d27794-w-0.xplane.pb
gs://chengnuojin-xprof/ep-parity-qwen35-35b-20260902/qwen3.5-35b-a3b-engine-scan-fsdp4tp2/plugins/profile/2026_09_02_16_39_03/t1v-n-c9d27794-w-0.xplane.pb
gs://chengnuojin-xprof/ep-parity-qwen35-35b-20260902/qwen3.5-35b-a3b-maxtext-scan-fsdp4tp2/plugins/profile/2026_09_02_16_40_50/t1v-n-c9d27794-w-0.xplane.pb

qwen3.5-35b-a3b, tensor parallelism and the PR #5060 A/B (2026-09-02)

gs://chengnuojin-xprof/pr5060-qwen35-35b-20260902/head/qwen3.5-35b-a3b-engine-scan-fsdp4tp2/plugins/profile/2026_09_02_01_25_39/t1v-n-c9d27794-w-0.xplane.pb
gs://chengnuojin-xprof/pr5060-qwen35-35b-20260902/head/qwen3.5-35b-a3b-maxtext-scan-fsdp4tp2/plugins/profile/2026_09_02_01_27_21/t1v-n-c9d27794-w-0.xplane.pb
gs://chengnuojin-xprof/pr5060-qwen35-35b-20260902/head/qwen3.5-35b-a3b-engine-scan-ga8-fsdp4tp2/plugins/profile/2026_09_02_01_30_53/t1v-n-c9d27794-w-0.xplane.pb
gs://chengnuojin-xprof/pr5060-qwen35-35b-20260902/head/qwen3.5-35b-a3b-engine-scan/plugins/profile/2026_09_02_01_32_50/t1v-n-c9d27794-w-0.xplane.pb
gs://chengnuojin-xprof/pr5060-qwen35-35b-20260902/pre-pr/qwen3.5-35b-a3b-engine-scan/plugins/profile/2026_09_02_03_59_58/t1v-n-c9d27794-w-0.xplane.pb

qwen3.5-35b-a3b, engine GA anatomy (2026-09-02)

gs://chengnuojin-xprof/engine-ga-anatomy-20260902/ga2-roe-rsort-engine/plugins/profile/2026_09_02_22_59_47/t1v-n-c9d27794-w-0.xplane.pb
gs://chengnuojin-xprof/engine-ga-anatomy-20260902/ga2-roe-rsort-peft/plugins/profile/2026_09_02_23_05_04/t1v-n-c9d27794-w-0.xplane.pb
gs://chengnuojin-xprof/engine-ga-anatomy-20260902/ga1-roe-engine/plugins/profile/2026_09_02_22_19_55/t1v-n-c9d27794-w-0.xplane.pb
gs://chengnuojin-xprof/engine-ga-anatomy-20260902/ga1-roe-peft/plugins/profile/2026_09_02_22_18_23/t1v-n-c9d27794-w-0.xplane.pb

Read any of them with:

python xplane_device_summary.py --steps 3 <path>.xplane.pb   # per-module device time
python xplane_host_summary.py             <path>.xplane.pb   # host-side dispatch

---

Appendix A: Reproducing

A.0 Prerequisites

A TPU VM with 8 devices (these numbers are from a v7-8; other 8-device generations will
give different absolute times but the same comparisons). Then:

# 1. Fetch PR #5060.
git clone https://github.com/AI-Hypercomputer/maxtext.git && cd maxtext
git fetch origin pull/5060/head:pr5060 && git checkout pr5060

# 2. Install MaxText (see docs/install_maxtext.md) plus the post-training extras --
#    the PeftTrainer arms need tunix, which lives in the post-train extra deps.
install_tpu_post_train_extra_deps

# 3. Everything below runs from this directory. It matters: the arms import
#    perf_parity_common as a sibling module, and a working directory containing a
#    tunix/ checkout would shadow the installed package.
cd tests/end_to_end/tpu/perf_parity

The engine-side Zero-1 measured in §2 is PR #5099, not #5060. On a #5060-only checkout
every engine arm prints zero1: UNSUPPORTED or zero1: DECLINED (...) and the Zero-1 rows
become a second copy of the baseline — the runner greps for that line so it cannot pass
unnoticed. The other five arms of §2 and all of §1 reproduce on #5060 alone.

A.1 One command per study

./run_qwen3_0p6b_zero1.sh   [outdir]   # §2, six arms x GA{1,8}, ~20 min (~10 min --no-trace)
./run_qwen3_5_35b_a3b.sh    [outdir]   # §1, ~90 min
./run_qwen3_0p6b.sh         [outdir]   # the sgd/tp qwen3-0.6b shape, for reference

Each prints its step times, peak HBM and trace paths at the end, and writes one log per arm
into the output directory. They run serially — every arm wants the whole host, and two
concurrent runs will fight over the 8 devices.

A.2 Individual arms — §1, qwen3.5-35b-a3b

SHAPE="--model qwen3.5-35b-a3b --scan --seq 1024 --ga 1"

# 1.1 headline, both trainers at both meshes
python engine_profile.py       $SHAPE --ep 8 --ring-of-experts --ragged-sort --no-trace
python peft_trainer_profile.py $SHAPE --ep 8 --ring-of-experts --ragged-sort --no-trace
python engine_profile.py       $SHAPE --tp 2 --no-trace
python peft_trainer_profile.py $SHAPE --tp 2 --no-trace

# 1.2 mesh ablation
python engine_profile.py $SHAPE --tp 2 --ep 4 --ring-of-experts --ragged-sort --no-trace
python engine_profile.py $SHAPE --ep 8 --no-trace
python engine_profile.py $SHAPE --ep 8 --ragged-sort --no-trace
python engine_profile.py $SHAPE --ep 8 --ring-of-experts --no-trace

# 1.3 gradient accumulation
for GA in 2 4 8; do
  python engine_profile.py       ${SHAPE/--ga 1/--ga $GA} --ep 8 --ring-of-experts --ragged-sort --no-trace
  python peft_trainer_profile.py ${SHAPE/--ga 1/--ga $GA} --ep 8 --ring-of-experts --ragged-sort --no-trace
done

# Device time.
PERF_PARITY_PROFILE_ROOT=/tmp/traces \
  python engine_profile.py $SHAPE --ep 8 --ring-of-experts --ragged-sort --steps 6

--ring-of-experts and --ragged-sort are rejected at --ep 1, mirroring MaxText, which
infers the EP rank from logical_axis_rules rather than from ici_expert_parallelism and
raises "When EP rank is 1, use_ring_of_experts must be False".

The GA table's PeftTrainer head column needs tunix past 44a35eeaf. Installing it without
touching the repo's pin:

pip install --no-deps "google-tunix @ https://github.com/google/tunix/archive/07dbe293.zip"

A.3 Individual arms — §2, qwen3-0.6b

SHAPE="--model qwen3-0.6b --seq 1024 --opt adamw"

for GA in 1 8; do
  python engine_profile.py       $SHAPE --dp 8            --no-trace --ga $GA  # auto
  python engine_profile.py       $SHAPE --dp 8 --explicit --no-trace --ga $GA  # explicit
  python engine_profile.py       $SHAPE --dp 8 --zero1    --no-trace --ga $GA  # + Zero-1
  python peft_trainer_profile.py $SHAPE --dp 8            --no-trace --ga $GA
  python peft_trainer_profile.py $SHAPE --dp 8 --explicit --no-trace --ga $GA
  python peft_trainer_profile.py $SHAPE --fsdp 8          --no-trace --ga $GA
done

# Device time. Trace the explicit control too -- shard_mode is not only a layout choice
# here, so the control compiles different kernels from the baseline.
PERF_PARITY_PROFILE_ROOT=/tmp/traces \
  python engine_profile.py $SHAPE --dp 8 --zero1 --steps 6 --ga 8

--zero1 implies --explicit. RunSpec rejects the three combinations the engine would
reject anyway — Zero-1 with FSDP, with data of 1, and with sgd — rather than letting them
run and quietly measure the baseline.

A.4 Flags

Flag	Default	Meaning
--model	qwen3-0.6b	any MaxText model_name; only qwen3_0p6b_tunix_profile.py is model-locked
--dp / --fsdp / --tp / --ep	1 / fills devices / 1 / 1	mesh axes; the product must equal the device count
--ga	1	micro-batches per optimizer step
--seq, --steps	1024, 23	tokens per example; optimizer steps including warmup
--scan	off	MaxText's scanned decoder (its production default; tunix uses a Python loop)
--opt	sgd	sgd or adamw, matched across arms
--explicit	off	shard_mode=explicit — Explicit rather than Auto mesh axes
--zero1	off	shard_optimizer_over_data; implies --explicit; engine arm only
--ring-of-experts, --ragged-sort	off	MoE kernels; require --ep > 1
--no-trace	traces on	skip xprof — required for any wall-clock number
--devices N	all	use only the first N local devices
PERF_PARITY_PROFILE_ROOT	./profiles	where traces are written

---

Appendix B: Measurement notes

These are the traps that produce plausible wrong numbers rather than errors. Each was hit.

Wall clock must come from --no-trace runs. The profiler charges per dispatch, and the two
trainers dispatch very differently — at qwen3-0.6b/GA=8 the traces record 147.3 module
launches per step on the engine against 18.0 on PeftTrainer. Both run nine substantial
kernels; the engine's other ~138 are small eager dispatches, each a separate host round trip.
Traced GA=8 wall clock runs 345–586 ms against 247–372 untraced, and inflates the arms
unevenly — a traced A/B reverses the GA=8 ranking outright. On the 35b model the same
overhead is at most 0.08%, because a 2.3 s step with ~14 dispatches gives it nothing to bite on.

Per-kernel device times are comparable within a GA setting, not across one. The same HLO
program hash measures 1.8–4.6 ms slower at GA=1 than at GA=8 in all three engine arms. Every
subtraction in §2.2 is taken within one GA column.

Tracing lowers the engine's peak HBM by ~2.4 G, and only the engine's. 12.73 G untraced
against 10.36 G traced at GA=8; both PeftTrainer arms are identical either way. Step count is
not the variable — re-running untraced at --steps 6 reproduces the 23-step figures exactly.
The likely cause is dispatch depth (the engine runs further ahead of the device, and the
profiler's synchronization drains the queue), consistent with the engine-only incidence and the
~2.4 G size, but not verified directly. Memory figures come from --no-trace runs.

Read TPU-busy off the XLA Modules line, not XLA Ops. On a scanned MoE the Ops line sums
to roughly twice both the module time and the wall step, because ops inside the scan are
emitted underneath a fusion that already covers them.

Take the maximum per-execution duration, never a mean or a total-over-steps. Ragged sort
runs on SparseCore and floods the trace buffer (~1.3 M events against ~18 k at tp=2), dropping
most of the XLA Modules line; a clipped event then averages in as a fast step. maybe_trace
now passes base.yml's enable_tpu_profiling_options to jax.profiler.trace, which caps the
capture and keeps every execution; --no-tpu-profiling-options restores the raw behaviour.

A bare jax.sharding.Mesh(create_device_mesh(...), axes) silently disables Zero-1.
_zero1_active requires all mesh.axis_types to be Explicit, and that constructor leaves
them Auto. The arms build their mesh with maxtext_utils.get_mesh_from_config(...), and
every engine arm prints a zero1: ACTIVE | DECLINED (reason) | UNSUPPORTED line — a declined
run is the baseline wearing the feature's name, which is invisible in a step time.

The engine's own fwd_bwd / update split does not decompose the step. fwd_bwd dispatches
asynchronously and the blocking wait lands inside update(), so at GA=1 the whole step appears
on the update side. Use the device figures.

enable_checkpointing=False does not stop the engine writing a final checkpoint. close()
ends in save_checkpoint(..., force=True) and the manager arms itself on checkpoint_dir being
non-empty. On the 35b model in f32 that is a ~140 G write per run, which filled this host's disk
and killed the first sweep. engine_profile.py now disarms the manager when the config says
checkpointing is off; the write lands after the timed loop either way.

---

GRPO integration test: tests/post_training/integration/maxtext_engine_grpo_loss_test.py — 1 passed in 51.67 s on the real checkpoint. Note it never calls engine.compile(), so it measures the eager path (237 ms/update vs 14 ms). Two things in it are now load-bearing: the batch is derived from the mesh rather than hard-coded, for the reason under change (1); and matmul_precision=highest is pinned, because the KL assertion compares log-probs from two code paths — Tunix's compute_per_token_logps for the reference against the engine's sharded forward for the policy — which at bf16 disagree by ~2e-2 on values near -12.6, and low_var_kl squares that into a KL of ~1.2e-4, an order of magnitude above the tolerance the assertion is trying to police.

xprof traces for the (7)-(9) A/B, engine arm only, base/ = this fix's parent commit and head/ = the fix, at the two shapes in "The host step path":

gs://chengnuojin-xprof/engine-hostpath-20260901/base/qwen3-0.6b-engine-fsdp1tp8/
gs://chengnuojin-xprof/engine-hostpath-20260901/head/qwen3-0.6b-engine-fsdp1tp8/
gs://chengnuojin-xprof/engine-hostpath-20260901/base/qwen3.5-35b-a3b-engine-scan-fsdp4tp2/
gs://chengnuojin-xprof/engine-hostpath-20260901/head/qwen3.5-35b-a3b-engine-scan-fsdp4tp2/

Traced, that A/B reads 271.4 → 92.8 ms and 2350.0 → 2315.2 ms. Compilation is inside these windows — the engine arm calls compile() under the trace on purpose — so read the steady-state region, not the whole trace.

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

✅ All modified and coverable lines are covered by tests.

📢 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.