recipes(397b v7x): the cp-as-ep-dense-no-cp mesh rule joins TRAINER_TUNED_FLAGS

[parambole]

Draft — stacked on #2913 (parambole/rl-trainer-4600 at bee63532f), companion to AI-Hypercomputer/maxtext#5783 (20c93dc4f on AI-Hypercomputer/maxtext#5758).

Summary

Adds custom_mesh_and_rule=cp-as-ep-dense-no-cp (defined in AI-Hypercomputer/maxtext#5783) to the TRAINER_TUNED_FLAGS=true block of tunix/experimental/examples/recipes/mlperf_397b_256_v7x.sh, after the recipe's default custom_mesh_and_rule=cp-as-ep.

• When TRAINER_TUNED_FLAGS=false (the default): Every rendered manifest and flag string is byte-identical to Qwen3.5-397B RL trainer on 128 v7x: pad-first sequence packing, fused sharded log-prob head VJP, and the tuned trainer recipe #2913 (custom_mesh_and_rule=cp-as-ep).
• When TRAINER_TUNED_FLAGS=true: MaxText keeps the last value of a repeated key, so the trainer selects cp-as-ep-dense-no-cp. Under cp-as-ep-dense-no-cp, the dense weights (attention q/k/v/out, shared-expert wi_0/wi_1/wo and gate, and the frozen MoE router; logical axes embed, embed_attn, q_lora, kv_lora) are sharded 64-way over fsdp x expert instead of 256-way over fsdp x context x expert and replicated across context=4. Their forward and backward weight all-gathers run in four disjoint 64-device groups without crossing context ranks, while their gradients are still reduce-scattered over all 256 devices and then all-gathered over context. The GDN projections (embed_gdn), routed experts, embeddings, and activations remain sharded identically to cp-as-ep.

---

Results on 128 TPU v7x Chips (Qwen3.5-397B-A17B RL Trainer Bench)

Setup: AI-Hypercomputer/maxtext#5701 bench driver, 128 TPU v7x chips (256 devices, mesh fsdp=32 x context=4 x expert=2), 9 micro-batches per optimizer step (294,912 tokens / step/chip), ragged_buffer_factor=2.75 with dropless fallback, 0 dropped tokens across all 45 micro-batch audits per timing run. Step time is the mean of steps 1–4; regular micro-batch time is the median of micro-batches 3–9.

Configuration	MaxText / Tunix Code	custom_mesh_and_rule	Pods / Runs	Step Time (s/step, Steps 1–4)	Regular Micro-Batch (s/mb)	Throughput (tok/s/chip reg / ruler)	6-Gate Numerics (num_readout.py)
Locked Baseline (#5758 + #2913)	20c877aa1 + bee63532f	cp-as-ep (256-way dense weights)	4 runs in 3 pods (bprf-v14, bl3-v14, bslow-v14)	82.200, 82.158, 82.111, 82.269 (mean 82.185 s)	8.942–8.959 s	3,588 / 3,655	PASS 6/6 (nxprf-l3: loss 2.69e-7, logp 1.87e-6, last-mb 5.70e-6, ent 5.07e-7, grad_norm 1.11e-4 vs tol 2e-3)
This PR (#5783 + cp-as-ep-dense-no-cp)	20c93dc4f + this PR	cp-as-ep-dense-no-cp (64-way dense weights)	2 pods (parambole-bl2-v14, parambole-bl2r-v14)	81.261 s (Pod 1), 81.245 s (Pod 2; step 2 81.231 s) (mean 81.253 s, -0.932 s/step)	8.838 s (-110 ms/mb)	3,630 / 3,698 (+42 tok/s/chip)	PASS 6/6 on BOTH pods (nxl2 Pod 1 N3=4.58e-6, nxl2r Pod 2 N3=2.22e-5 vs replicate spread 1.92e-5, tol 2e-3; loss = 0.0304453575 bit-identical at all 3 steps)

Memory (128-Chip AOT & Runtime)

• AOT (--xla_tpu_max_hbm_size_mib=84992): fwd_bwd HBM usage is 81.05 GiB under cp-as-ep-dense-no-cp vs 78.28 GiB under cp-as-ep (+2.77 GiB; train-kernel device peak 109.04 GiB vs 104.46 GiB).
• TPU Runtime (xprof get_kpi_metrics): Peak HBM is 87.66 GiB (bl2-v14 / bl2r-v14) vs 84.79 GiB (bprf-v14). Both pods completed all 5 steps cleanly without OOM.

---

Testing

• Launcher unit test (tests/experimental/examples/deepswe_dist/mlperf_submission_launcher_test.py): 41/41 passed (0 skipped). Verifies:
  • TRAINER_TUNED_FLAGS=true sets custom_mesh_and_rule=cp-as-ep-dense-no-cp on the trainer JobSet while leaving the other five renders byte-identical.
  • TRAINER_TUNED_FLAGS=false (default) and the other 397b_v7x recipes keep custom_mesh_and_rule=cp-as-ep and do not pass cp-as-ep-dense-no-cp.
• Shell syntax & diff check: bash -n tunix/experimental/examples/recipes/mlperf_397b_256_v7x.sh (rc=0), git diff --check (rc=0).