Match Tunix peft_trainer_v2 performance in MaxTextTrainingEngine

.gitignore

@@ -156,3 +156,6 @@ gha-creds-*.json
 # vscode workspace
 maxtext.code-workspace
 maxtext_output/
+
+# Benchmark output from tests/end_to_end/tpu/perf_parity (base_output_directory).
+tests/end_to_end/tpu/perf_parity/maxtext_out/

docs/reference.md

@@ -49,6 +49,13 @@ High-level overview of MaxText design, JAX/XLA choices, and how components inter
 
 Key concepts including checkpointing strategies, quantization, tiling, and Mixture of Experts (MoE) configuration.
 ```
+
+```{grid-item-card} ⚖️ Training Engine vs. Tunix
+:link: reference/training_engine_tunix_parity
+:link-type: doc
+
+Measured numerical parity, step time, and HBM of `MaxTextTrainingEngine` against Tunix `PeftTrainer` v2.
+```
 ````
 
 ## 📚 API Reference
@@ -64,5 +71,6 @@ reference/performance_metrics
 reference/models
 reference/architecture
 reference/core_concepts
+reference/training_engine_tunix_parity
 reference/api.rst
 ```

tests/end_to_end/tpu/perf_parity/RESULTS-qwen35-35b-20260902.md

@@ -0,0 +1,253 @@
+# PR #5060 on qwen3.5-35b-a3b: `MaxTextTrainingEngine` vs tunix `PeftTrainer v2`
+
+Both trainers drive the **same MaxText model**, so what is varied is the trainer and nothing
+else. Run 2026-09-02 against PR #5060 head `0d9db2747` ("Match Tunix peft_trainer_v2
+performance in MaxTextTrainingEngine").
+
+## Summary
+
+1. **At GA=1 the two trainers are at parity** — 2314.3 ms/step engine against 2303.0 ms/step
+   `PeftTrainer`, both at ~99.7% device utilization. At 2.3 s of device work per step every
+   host-side difference between them hides.
+2. **`PeftTrainer v2` cannot run gradient accumulation on this model at all.** GA=2, 4 and 8
+   each OOM with a byte-identical 161.94 G, so GA=1 is the only depth where a head-to-head
+   number exists. Above it the comparison is a capability difference, not a ratio.
+3. **PR #5060 is what makes this model runnable under tensor parallelism.** The pre-PR engine
+   does not merely run slower at tp=2 — it dies with a shape mismatch in the MoE. At tp=1,
+   where both revisions run, the PR is worth 1.79x, and the gain is device-side.
+
+## Environment
+
+| | |
+| --- | --- |
+| Host | 8 x TPU7x (v7-8 Ironwood, 4 chips / 8 JAX devices), single process |
+| JAX | 0.11.1 |
+| Model | `qwen3.5-35b-a3b` — 40 layers, 256 routed + 1 shared expert, emb 2048, 16 query / 2 KV heads, head_dim 256, vocab 248320 |
+| Mesh | fsdp=4 x tp=2, scanned (`--tp 2 --scan`) |
+| Shape | micro-batch 8 x seq 1024, f32 compute and weights, `remat_policy=none` |
+| Optimizer | `optax.sgd(1e-5)`, constant schedule, no clipping |
+| Wall clock | 23 steps, median of the last 19, `--no-trace` |
+| Device time | `XLA Modules` line per core, from a separate `--steps 6` traced run |
+
+`--tp 8` is not available for 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.
+
+## Reproducing
+
+Everything in §1, §2 and §5 comes from one wrapper:
+
+```bash
+cd tests/end_to_end/tpu/perf_parity
+./run_qwen3_5_35b_a3b.sh                     # ~90 min; prints the step times and trace paths
+```
+
+Or the individual arms. Note that only `engine_profile.py` and `peft_trainer_profile.py` can
+run this model — `qwen3_0p6b_tunix_profile.py` is locked to qwen3-0.6b, which is why the
+comparison here is engine vs `PeftTrainer` over the same MaxText model:
+
+```bash
+# Wall clock. --no-trace, because tracing charges per dispatch.
+python engine_profile.py       --model qwen3.5-35b-a3b --tp 2 --scan --no-trace --ga 1
+python peft_trainer_profile.py --model qwen3.5-35b-a3b --tp 2 --scan --no-trace --ga 1
+
+# Device time. --steps 6 keeps each xplane near 1 GiB at this shape.
+PERF_PARITY_PROFILE_ROOT=/tmp/traces \
+  python engine_profile.py --model qwen3.5-35b-a3b --tp 2 --scan --steps 6 --ga 1
+
+# Then read TPU-busy and launches off the xplane. Takes the .xplane.pb itself, not the
+# run directory. See the first measurement note in §5 before quoting anything from it.
+python xplane_device_summary.py --steps 3 /tmp/traces/<arm>/plugins/profile/<ts>/<host>.xplane.pb
+
+# Revision arms. PR #5060 touches only src/maxtext/training_engine/{maxtext_engine,
+# metrics,inflight_throttler}.py, so checking that one directory out at an older revision
+# gives a true A/B with the harness and all of main held at PR head.
+git checkout 2fadab37f -- src/maxtext/training_engine/   # 98e6886e8^, host-path fix reverted
+git checkout cbb2a57dd -- src/maxtext/training_engine/   # 59f49ac90^, pre-PR
+git checkout pr5060-head -- src/maxtext/training_engine/ # restore
+```
+
+## 1. Headline: trainer vs trainer, GA=1
+
+GA=1 because it is the only depth `PeftTrainer v2` survives on this model — see §2.
+
+| Arm (MaxText qwen3.5-35b-a3b on both sides) | Median ms/step | Mean | Max | TPU-busy/step | Util | Loop | Process |
+| --- | --- | --- | --- | --- | --- | --- | --- |
+| MaxText model + **engine**, PR head `0d9db2747` | 2314.3 | 2314.2 | 2314.8 | 2308.1 ms | **99.7%** | 56.2 s | 71.8 s |
+| MaxText model + **`PeftTrainer v2`** | **2303.0** | 2303.0 | 2303.6 | 2299.4 ms | **99.8%** | 55.8 s | 69.3 s |
+| MaxText model + engine, `98e6886e8^` | 2321.8 | 2338.9 | **2642.9** | — | — | 56.9 s | 72.1 s |
+| MaxText model + engine, `59f49ac90^` (pre-PR) | **crash** | — | — | — | — | — | — |
+
+`PeftTrainer` leads by 11.3 ms, 0.5%. Its single fused `jit__train_step` costs one dispatch
+per step where the engine's split `fwd_bwd` + `update` costs two, and the extra
+`jit__update_kernel` is 16.46 ms of that gap. Neither side has meaningful exposed host time.
+
+Traces:
+
+```text
+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
+```
+
+## 2. Gradient accumulation: `PeftTrainer v2` cannot do it here
+
+| GA | engine @ PR head | per micro-batch | `PeftTrainer v2` |
+| --- | --- | --- | --- |
+| 1 | 2314.3 ms | 2314.3 | 2303.0 ms |
+| 2 | 4613.6 ms | 2306.8 | **OOM** — 161.94 G vs 94.74 G |
+| 4 | 9211.0 ms | 2302.8 | **OOM** — 161.94 G vs 94.74 G |
+| 8 | 18405.3 ms | **2300.7** | **OOM** — 161.94 G vs 94.74 G |
+
+```text
+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 — that would scale with N. The cost is GA-independent, and the
+source says why. `PeftTrainer._is_single_microstep()` is true only when
+`gradient_accumulation_steps == 1`, and only on that path does `GradientAccumulator` get
+`allocate_grads=False`:
+
+```python
+else:
+  # When every update consumes exactly one micro-batch, `set()` overwrites
+  # the whole tree before anything reads it, so the initial zeros are dead
+  # on arrival. Skipping them avoids writing a full copy of the parameter
+  # tree (~3.5 GiB per device for gemma4-e2b at 12 layers in fp32).
+  self.grads = nnx.data({})
+```
+
+Any GA>1 takes the other branch and allocates `jnp.zeros_like(x, dtype=jnp.float32)` over the
+whole parameter tree, persistent. At 12-layer gemma4-e2b that is the ~3.5 GiB the comment
+cites; at 35B it does not fit. No shape knob moves it, because the allocation is
+parameter-shaped rather than activation-shaped.
+
+The engine has no such wall: it scales linearly, and its per-micro cost *improves* from 2314.3
+to 2300.7 ms as the single update amortizes over more micro-batches.
+
+Trace (engine, GA=8):
+
+```text
+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
+```
+
+## 3. What PR #5060 buys: tensor parallelism works at all
+
+At tp=2 the pre-PR engine does not run. It fails while tracing the MoE:
+
+```text
+TypeError: add got incompatible shapes for broadcasting: (8, 1024, 1024), (8, 1024, 2048).
+  src/maxtext/models/qwen3.py:1208
+    final_output = routed_output + jax.nn.sigmoid(shared_gate_output) * shared_expert_output
+```
+
+1024 is the TP-sharded width of the 2048 embedding: the routed output came out sharded while
+the shared-expert output stayed full. The cause is `_sharding_ctx`, which enters
+`jax.set_mesh(mesh)` and `nn_partitioning.axis_rules(config.logical_axis_rules)` around
+kernel tracing. It is introduced by 59f49ac90 and appears **0 times pre-PR and 7 times at
+head**. Without it MaxText's logical sharding constraints are silent no-ops.
+
+**Control.** To show this is the axis rules and not an artifact of pairing an old engine with
+current main, the same pre-PR engine was run at tp=1, where no logical constraint has to fire:
+
+| GA=1, `--steps 6` | tp=1 (fsdp=8) | tp=2 (fsdp=4) |
+| --- | --- | --- |
+| engine @ `59f49ac90^` | 4232.7 ms | **crash** |
+| engine @ PR head | 2358.3 ms | 2314.0 ms |
+
+It runs clean at tp=1, so the rig is sound and the tp=2 failure is specific to sharding.
+
+**And at tp=1 the PR is worth 1.79x, device-side:**
+
+| tp=1 (fsdp=8), GA=1 | Median ms/step | TPU-busy/step | Util | fwd/bwd kernel | Launches/core |
+| --- | --- | --- | --- | --- | --- |
+| engine @ `59f49ac90^` | 4232.7 | 4137.1 ms | 97.7% | `jit_kernel` 4120.25 ms/exec | 121 |
+| engine @ PR head | **2358.3** | **2352.2 ms** | 99.7% | `jit_first_kernel` **2336.11** ms/exec | 85 |
+| speedup | **1.79x** | 1.76x | | **1.76x** | |
+
+The compiled kernel itself nearly halves. This is not recovered host latency — it is XLA
+partitioning the model properly once the axis rules are live.
+
+Traces:
+
+```text
+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
+```
+
+## 4. The host-path fix (`98e6886e8`) at this scale: no median, all tail
+
+| A/B on the engine files alone | Median | Mean | Min | Max | Spread |
+| --- | --- | --- | --- | --- | --- |
+| GA=1, `98e6886e8^` | 2321.8 | 2338.9 | 2320.4 | **2642.9** | 322.5 ms |
+| GA=1, PR head | **2314.3** | 2314.2 | 2313.4 | 2314.8 | **1.4 ms** |
+| speedup | 1.003x | 1.011x | | | |
+| GA=8, `98e6886e8^` | 18423.0 | 18423.7 | 18420.5 | 18432.1 | 11.6 ms |
+| GA=8, PR head | **18405.3** | 18405.4 | 18403.3 | 18407.7 | 4.4 ms |
+| speedup | 1.001x | 1.001x | | | |
+
+The median barely moves, which is the expected result and not a disappointing one: at 99.7%
+utilization there is no exposed host time left to remove. What the fix does here is remove the
+tail. The worst step drops 2642.9 -> 2314.8 ms and the mean converges onto the median. Each
+`nnx.split` allocated a large short-lived object graph twice per step, and the GC pauses that
+followed landed on whichever step was unlucky.
+
+This independently reproduces the 2641.8 -> 2314.6 reported in §9 of
+`docs/reference/training_engine_tunix_parity.md`.
+
+## 5. Device-side detail
+
+Per-core figures from the `--steps 6` traced runs. Utilization is TPU-busy over the untraced
+median for the same shape.
+
+| Arm | TPU-busy/step | Launches/core/step | Dominant modules (ms/exec) |
+| --- | --- | --- | --- |
+| engine, tp=2, GA=1 | 2308.1 ms | 14.2 | `jit_first_kernel` 2291.33, `jit__update_kernel` 16.46 |
+| `PeftTrainer`, tp=2, GA=1 | 2299.4 ms | 1.0 | `jit__train_step` 2299.5 |
+| engine, tp=2, GA=8 | 18385.4 ms | 6.5 | `jit_accum_kernel` 2296.81 x7, `jit_first_kernel` 2291.3, `jit__update_kernel` 16.46 |
+| engine, tp=1, GA=1 | 2352.2 ms | 14.2 | `jit_first_kernel` 2336.11, `jit__update_kernel` 15.91 |
+| engine @ pre-PR, tp=1, GA=1 | 4137.1 ms | 20.2 | `jit_kernel` 4120.25, `jit__update_kernel` 16.04 |
+
+The GA=8 step is reconstructed as `1 x jit_first_kernel + 7 x jit_accum_kernel + 1 x
+jit__update_kernel` = 18385.4 ms against an 18405.3 ms wall step, i.e. **99.9%**. The
+`jit_first_kernel` module hash is `6948636727452340735` in both the GA=1 and GA=8 tp=2 traces
+— the same compiled program, so its clean GA=1 per-execution cost is the right one to use.
+
+### Three measurement notes
+
+- **Read TPU-busy off the `XLA Modules` line, not `XLA Ops`.** On this scanned MoE the Ops
+  line sums to roughly twice both the module time and the wall step (39491.5 vs 19846.4 ms on
+  the GA=8 trace), because ops inside the scan are emitted underneath a fusion that already
+  covers them. A module execution is unambiguously "the core was running a compiled program".
+- **The GA=8 device trace truncates.** Only 9 of the 48 `fwd_bwd` executions the run performed
+  were captured, so totals-over-steps understates that row badly. Per-execution cost is
+  unaffected by how many executions were recorded, so the step is built from those.
+- **Tracing overhead is negligible on this model**, unlike qwen3-0.6b. Traced against untraced
+  medians: 2316.0/2314.3, 2304.2/2303.0, 18408.9/18405.3, 2359.8/2358.3 — at most 0.08%. A
+  2.3 s step with ~14 dispatches gives the per-dispatch profiler cost nothing to bite on.
+
+## Profiles
+
+All five xplanes, full paths:
+
+```text
+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
+```
+
+| Trace | Arm | Size |
+| --- | --- | --- |
+| `head/…-engine-scan-fsdp4tp2/` | engine, tp=2, GA=1 | 1225.8 MiB |
+| `head/…-maxtext-scan-fsdp4tp2/` | `PeftTrainer v2`, tp=2, GA=1 | 1598.8 MiB |
+| `head/…-engine-scan-ga8-fsdp4tp2/` | engine, tp=2, GA=8 | 1848.8 MiB |
+| `head/…-engine-scan/` | engine, tp=1, GA=1 | 1112.6 MiB |
+| `pre-pr/…-engine-scan/` | engine @ `59f49ac90^`, tp=1, GA=1 | 1494.0 MiB |
+
+Compilation lands inside each trace window — the harness calls `engine.compile()` inside
+`jax.profiler.trace` so its cost sits where `PeftTrainer`'s does. It contributes no `XLA
+Modules` events, so the device figures above are unaffected.