Match Tunix peft_trainer_v2 performance in MaxTextTrainingEngine

[NuojCheng]

Description

MaxTextTrainingEngine was slower than Tunix's PeftTrainer v2 on the same model, the same loss function and the same optax transformation, and the gap turned out to be entirely host-side. This PR closes it, and fixes two correctness problems the measurement turned up along the way.

Five files, no new scripts:

file	why
src/maxtext/training_engine/maxtext_engine.py	the step path, the accumulation normalization, the gradient norm
src/maxtext/training_engine/inflight_throttler.py	queue a scalar per step instead of the whole train state
src/maxtext/training_engine/metrics.py	deferred metrics write
tests/post_training/unit/maxtext_engine_test.py	coverage for all of the above
tests/post_training/integration/maxtext_engine_grpo_loss_test.py	size its batch to the mesh, per the behaviour change below

Performance

• Trace the accumulation and update kernels under nn_partitioning.axis_rules, so gradient accumulation is fused into the compiled program instead of running as an eager jax.tree.map(add) over every gradient leaf on every micro-batch.

  This is the one behaviour change in the PR, as opposed to a timing change. Those rules live in a context variable, so a kernel traced without them saw an empty rule set and every maybe_shard_with_logical inside the MaxText layers was silently a no-op; making them real means micro_batch_size_to_train_on must now be a multiple of data x fsdp. A batch those devices cannot split gets 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, and in practice such a batch was already leaving devices idle. This is how it showed up: the GRPO integration test hard-coded a batch of 2 and ran on a 4-device host, hence the fifth file above.

• Carry a pure nnx.State mirror across steps in place of a per-step nnx.split of the module graph. The split was two graph traversals per step, scaling with node count — 21.3 ms on an unscanned qwen3-0.6b's 310 parameter leaves — and it allocated a large short-lived object graph twice per step, whose GC pauses landed on whichever step was unlucky (a 464 ms worst step against a 161.6 ms median).

• Defer the metrics write until after the next dispatch, so its device-to-host read overlaps live work rather than blocking on it.

Correctness

• Accumulate gradients unreduced and divide once by the summed denominator, rather than averaging per-micro-batch means. The two agree only when every micro-batch carries the same number of tokens; on ragged batches, mean-of-means is simply not the gradient of the batch loss. The denominator rides along in checkpoint metadata so a mid-step resume can finish the division, and is backfilled from the restored per-micro-batch losses for checkpoints written before it was tracked.

• Compute the gradient norm on every update, not only under skip_step_on_spikes. learning/grad_norm is a metric trainers/pre_train/train.py reports every step and metrics._METRICS_TO_LOG already lists, but with spike-skipping off — base.yml's default — it read NaN. Tunix computes the same quantity unconditionally inside peft_trainer_v2._update_step. 202a89ab8 landed the same fix on main while this was in review; rebasing onto it, I kept main's placement — before clipping — over this branch's post-clipping one, because that is where Tunix's optax.global_norm sits (its clipping, when a caller configures any, is a later link in the optax chain). In train.py's vocabulary that makes it raw_grad_norm rather than learning/grad_norm. What this PR still adds on top of main's version is the float32 reduction — a sum of squares over bf16 leaves overflows on production-size models — and the norm doubling as the throttler's handle. It agrees with Tunix bit for bit:

  	MaxText gradient_norm	Tunix grad_norm
  GA=1, batch 8	59.77279281616211	59.77279281616211
  GA=8, batch 8	18.25168800354004	18.25168800354004

  It costs ~0.6 ms per update — 0.7% of a GA=1 step, 0.1% of a GA=8 one, and below the noise floor on device (568.0 vs 570.7 ms TPU-busy per optimizer step at GA=8), because XLA fuses the reduction into an update kernel that already streams every gradient leaf.

  The norm also gives InflightThrottler a scalar to block on. It previously queued the train state itself, which pinned three parameter trees alive for as long as the entry sat in the queue, and — once the update kernel donates its state argument — could raise Array has been deleted out of jax.block_until_ready when an entry was popped after a later step had reused those buffers.

Results

qwen3-0.6b, batch 8 × seq 1024, f32, optax.sgd(1e-5), no clipping, measured as an A/B on these four files alone — the "before" arm is git checkout origin/main -- src/maxtext/training_engine/ on this same commit, so nothing else differs. 4 × v6e, fsdp=4, unscanned, median of 19 steps after warmup, no profiler attached:

	before (main)	after (this PR)	speedup
GA=8, per optimizer step	4443.3 ms	613.0 ms	7.25x
GA=8, per micro-batch	555.4 ms	76.6 ms	7.25x
GA=8, 23-step loop	135.8 s	24.2 s	5.61x
GA=8, process wall clock	161.3 s	49.8 s	3.24x
GA=1, per step	516.8 ms	84.5 ms	6.12x
GA=1, 23-step loop	21.1 s	8.7 s	2.43x
GA=1, process wall clock	46.6 s	34.1 s	1.37x

Gradient accumulation multiplies the win, because the removed nnx.split sat in fwd_bwd, which runs once per micro-batch rather than once per optimizer step. The process-wall-clock rows are lower than the rest because ~25 s of interpreter start, imports, config and engine build is fixed cost this PR does not touch.

The saving is a fixed quantity of host time, so what it is worth also depends on how big the NNX module graph is — scan_layers collapses a decoder stack into one stacked node set, so a scanned qwen3.5-35b-a3b (70 parameter leaves) sees 2322.5 → 2314.0 ms/step, a 1.004x median, while its worst step still drops 2641.8 → 2314.6 ms as the GC pauses go away — and on how much device time there is to hide host work behind.

Against Tunix's trainer at GA=8, driving the same MaxText model, so only the trainer differs (measured on an earlier base, where this PR's arm read 589.6 ms rather than 613.0):

GA=8, qwen3-0.6b	ms/optimizer step	ms/micro-batch	TPU-busy/step	utilization
MaxTextTrainingEngine	589.6	73.7	568.0 ms	96%
PeftTrainer, MaxText model	1225.9	153.2	617.7 ms	50%

Two independent checks that the win is host-only: TPU-busy time across the GA=8 A/B is 3408.12 vs 3408.13 ms over the six optimizer steps each trace covers — agreement to one part in 340,000 — and jit_accum_kernel carries the same XLA module hash (13160937502384084184) on both sides.

Full write-up, methodology and xplane traces: #5060.

Tests

tests/post_training/unit/maxtext_engine_test.py, extended here to cover:

• the cached pure-state path end to end, on a model with non-Param state and over two steps, which is where a stale or wrongly-partitioned cache would show up;
• unreduced accumulation and its denominator, including a mid-step checkpoint round-trip;
• the throttler queueing one scalar rather than the state's leaves;
• gradient_norm being recorded with skip_step_on_spikes off — the configuration that used to produce nothing.

JAX_PLATFORMS=cpu pytest tests/post_training/unit/maxtext_engine_test.py \
  tests/post_training/unit/maxtext_engine_constructor_test.py \
  tests/post_training/unit/maxtext_engine_e2e_test.py \
  tests/post_training/unit/router_replay_engine_test.py
# 59 passed

On TPU, tests/post_training/integration/maxtext_engine_grpo_loss_test.py passes against the real GCS checkpoint (1 passed in 51.67 s), and the numerical parity harness in #5060 shows identical loss and gradients agreeing to rel_l2 3.0e-4 against Tunix at GA=1. Two things in that test are now load-bearing: its batch comes from mesh.shape["data"] * mesh.shape["fsdp"] rather than a hard-coded 2, for the reason above; and matmul_precision=highest is pinned, because its 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 meant to police.

Checklist

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

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

[gemini-code-assist]

Code Review

This pull request introduces significant performance optimizations and fixes for gradient accumulation and training engine overhead in MaxText. Key changes include caching the pure-state representation of the NNX model to avoid costly graph re-splitting, accumulating unreduced gradients to divide by the accumulated denominator once during updates, bounding the metrics history to prevent HBM bloat, and properly handling PartitionSpecs with reduced/unreduced tags across layers and sharding helpers. The review feedback highlights potential AttributeError issues in moe.py and normalizations.py when partition specs are None in single-device or unsharded environments, providing actionable code suggestions to guard these accesses.

[gemini-code-assist]

If self.mesh is None (e.g., in single-device or local test environments), w0_pspec will be None. In this case, accessing w0_pspec.partitions will raise an AttributeError. We should guard this block with a check to ensure w0_pspec is not None before accessing its partitions.

Suggested change

	if weight_gather:
	# Read `.partitions`: a kernel spec carrying a reduced axis refuses direct indexing.
	wi_partitions = w0_pspec.partitions
	# weight_gather implies either exp or embed_moe is sharded.
	if self.config.shard_exp_on_fsdp:
	# wi [Experts, In, Hidden] -> Gather Exp(0)
	wi_gather_axes.extend(get_active_sharding_axes(w0_pspec[0], 0))
	wi_gather_axes.extend(get_active_sharding_axes(wi_partitions[0], 0))
	else:
	# Gather In(1) where embed_moe is sharded.
	wi_gather_axes.extend(get_active_sharding_axes(w0_pspec[1], 1))
	wi_gather_axes.extend(get_active_sharding_axes(wi_partitions[1], 1))
	# Gather Hidden(2)
	wi_gather_axes.extend(get_active_sharding_axes(w0_pspec[2], 2))
	wi_gather_axes.extend(get_active_sharding_axes(wi_partitions[2], 2))
	if weight_gather and w0_pspec is not None:
	# Read `.partitions`: a kernel spec carrying a reduced axis refuses direct indexing.
	wi_partitions = w0_pspec.partitions
	# weight_gather implies either exp or embed_moe is sharded.
	if self.config.shard_exp_on_fsdp:
	# wi [Experts, In, Hidden] -> Gather Exp(0)
	wi_gather_axes.extend(get_active_sharding_axes(wi_partitions[0], 0))
	else:
	# Gather In(1) where embed_moe is sharded.
	wi_gather_axes.extend(get_active_sharding_axes(wi_partitions[1], 1))
	# Gather Hidden(2)
	wi_gather_axes.extend(get_active_sharding_axes(wi_partitions[2], 2))

[gemini-code-assist]

Similarly to w0_pspec, wo_pspec can be None if self.mesh is None. We should guard this block to ensure wo_pspec is not None before accessing its partitions.

Suggested change

	if weight_gather:
	# Read `.partitions`: a kernel spec carrying a reduced axis refuses direct indexing.
	wo_partitions = wo_pspec.partitions
	# weight_gather implies either exp or embed_moe is sharded.
	if self.config.shard_exp_on_fsdp:
	# wo [Experts, Hidden, Out] -> Gather Exp(0)
	wo_gather_axes.extend(get_active_sharding_axes(wo_pspec[0], 0))
	wo_gather_axes.extend(get_active_sharding_axes(wo_partitions[0], 0))
	else:
	# Gather Out(2) where embed_moe is sharded.
	wo_gather_axes.extend(get_active_sharding_axes(wo_pspec[2], 2))
	wo_gather_axes.extend(get_active_sharding_axes(wo_partitions[2], 2))
	# Gather Hidden(1)
	wo_gather_axes.extend(get_active_sharding_axes(wo_pspec[1], 1))
	wo_gather_axes.extend(get_active_sharding_axes(wo_partitions[1], 1))
	if weight_gather and wo_pspec is not None:
	# Read `.partitions`: a kernel spec carrying a reduced axis refuses direct indexing.
	wo_partitions = wo_pspec.partitions
	# weight_gather implies either exp or embed_moe is sharded.
	if self.config.shard_exp_on_fsdp:
	# wo [Experts, Hidden, Out] -> Gather Exp(0)
	wo_gather_axes.extend(get_active_sharding_axes(wo_partitions[0], 0))
	else:
	# Gather Out(2) where embed_moe is sharded.
	wo_gather_axes.extend(get_active_sharding_axes(wo_partitions[2], 2))
	# Gather Hidden(1)
	wo_gather_axes.extend(get_active_sharding_axes(wo_partitions[1], 1))

[gemini-code-assist]

If scale is not sharded (e.g., in single-device mode or if auto-sharding is used and no sharding is assigned yet), scale_spec can be None. In this case, accessing scale_spec.partitions will raise an AttributeError. We should guard against scale_spec is None and return scale directly.

Suggested change

	scale_spec = jax.typeof(scale).sharding.spec
	if scale_spec[-1] == activation_spec[-1]:
	if scale_spec.partitions[-1] == activation_axis:
	return scale
	return jax.sharding.reshard(scale, jax.sharding.PartitionSpec(activation_spec[-1]))
	return jax.sharding.reshard(
	scale,
	jax.sharding.PartitionSpec(activation_axis, unreduced=scale_spec.unreduced, reduced=scale_spec.reduced),
	)
	scale_spec = jax.typeof(scale).sharding.spec
	if scale_spec is None:
	return scale
	if scale_spec.partitions[-1] == activation_axis:
	return scale
	return jax.sharding.reshard(
	scale,
	jax.sharding.PartitionSpec(activation_axis, unreduced=scale_spec.unreduced, reduced=scale_spec.reduced),
	)

[codecov]

Codecov Report

❌ Patch coverage is 78.48837% with 37 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/training_engine/maxtext_engine.py	74.82%	27 Missing and 10 partials ⚠️

📢 Thoughts on this report? Let us know!