Repository navigation
Qwen3.5-397B RL trainer on 128 v7x: the tuned recipe turns on the fp8 fake-quant hoist, fused SwiGLU backward, attention savings, the 2048 wi forward tile and the GDN context-parallel m scan, and pins the MaxText that defines them - #2948
Draft
parambole wants to merge 6 commits into
Draft
parambole wants to merge 6 commits into
parambole wants to merge 6 commits into
Conversation
Moves the MaxText pin from atwigg/mlperf at bbd0f6f10 (the squash merge of AI-Hypercomputer/maxtext#5721) to the head of the branch parambole/rl-4600-final-maxtext of AI-Hypercomputer/maxtext, 94f11c09df83e9903c093e4544beba3bb198ff8b. That branch is bbd0f6f10 plus six commits: the SwiGLU backward in the epilogue of the MoE down projection's input-gradient GMM (moe_wo_gmm_fused_dswiglu, two commits), the forward's load-balanced positions and segment ids reused in the backward (context_parallel_attention_load_balance_hoist_ids), the query kernel gathered once per attention layer in the backward (query_proj_kernel_gather_once), and two test-only commits. The three keys default to off, so this commit alone changes no rendered launch: the next two commits pass them under TRAINER_TUNED_FLAGS=true. Comments that named the pin follow it: the recipe's sentence on where the tuned keys are defined names the branch, and the launcher test's comment on _BENCH_MESH_RULE no longer says that the pinned MaxText lacks cp-as-ep-dense-no-cp (it has had it since bbd0f6f10). Tests, in the bench image (dbdcf5f8dcba; CPU, no network, the tree read-only): with this commit alone mlperf_submission_launcher_test.py passes 48 of 48, and the renders of mlperf_gbs1024_submission.sh (TRAINER_TUNED_FLAGS on and off, the 256-chip meshes 32x2x8 and 64x1x8) and of the other 397B v7x recipes are byte-identical to the parent's.
…joins TRAINER_TUNED_FLAGS TRAINER_TUNED_FLAGS=true now also passes attention's load-balance reorder of the positions and segment ids, kept out of the remat replay: the one MaxText change of the measured cell that the block did not pass. In mlperf_397b_256_v7x.sh the block of MaxText keys (part 2) gains context_parallel_attention_load_balance_hoist_ids=true, before the block's custom_mesh_and_rule=cp-as-ep-dense-no-cp: remat_policy=custom saves each attention layer's load-balance-ordered positions and segment ids. They depend only on the batch, so jax.lax.scan keeps one copy outside the layer loop, and the remat replay of every attention layer reads it instead of reordering them again. The values are the same. Compiled for TPU v7x, XLA already moves the reorder's collective-permutes out of the loop without the key; the key moves out the rest of the reorder and the work on its result (per the MaxText commit: an AOT compile of the bench's 128-chip trainer program at an earlier cell has 47 fewer instructions in its while loops with the change, the same collective-permutes and the same memory). MaxText refuses the key without context_parallel_attention_load_balance=true, an even context parallelism of at least 2, remat_policy=custom and scan_layers=true. This recipe's trainer has them at 32x2x4 and at the 256-chip meshes 32x2x8 and 64x1x8 (context 8), so the key goes on every trainer mesh, like the block's other keys. The key is defined in the MaxText change on top of AI-Hypercomputer/maxtext#5721 (branch parambole/rl-4600-final-maxtext), which the previous commit pins; atwigg/mlperf at bbd0f6f10 does not define it. TRAINER_TUNED_FLAGS=false and the other 397B v7x recipes do not pass it. Bench (AI-Hypercomputer/maxtext#5701 driver, Qwen3.5-397B-A17B, 128 TPU v7x, mesh fsdp 32 x context 4 x expert 2, ragged_buffer_factor 2.75, 9 micro-batches per step, mean of steps 1 to 4, no token dropped in any arm), alone, in three comparisons against controls in other pods whose arm commands differ only by the change's switch: 99.504 -> 99.267 s per step (o275x-gdn-ctl -> o275x-s3-x17h), 99.248 -> 99.267 (o275x-mf-ctl -> o275x-s3-x17h) and 91.7513 -> 91.8665 (o275x-stk3-l3-prof -> o275x-s5c-l3-prof), i.e. -0.237, +0.019 and +0.115 s per step. One unchanged configuration spans 0.255 s per step over four pods (o275x-hd), so no effect is resolved. Together with the next commit's two keys and the mesh rule, see the next commit (F2). This is 11c096bb0 (branch parambole/rl-trainer-4600-next of google/tunix), as ported onto #2884, replayed onto #2884's head b16ee52. Two hunks conflicted, both comments that #2884 had just rewritten for its new pin: the recipe's sentence on where the tuned keys are defined and the test's comment on _BENCH_MAXTEXT_KEYS. Each keeps #2884's wording and adds the key. The test's comment on _MAXTEXT_DEFAULTS now says the change is on top of #5721. Tests (mlperf_submission_launcher_test.py): _MAXTEXT_DEFAULTS has the key's default (false), the exact-set test expects it among _TUNED_KEYS, and the switch-off and other-recipe tests check that it is not passed (_BENCH_MAXTEXT_KEYS). In the bench image (dbdcf5f8dcba; CPU, no network) 48 of 48 pass at this commit. Rendered, the submission launch differs from the parent's only by this key in the trainer's MaxText flags, at 32x2x4, 32x2x8 and 64x1x8; with TRAINER_TUNED_FLAGS=false, and the other 397B v7x recipes at their defaults, the renders are byte-identical.
…her_once join TRAINER_TUNED_FLAGS TRAINER_TUNED_FLAGS=true now also turns on the two remaining keys of the MaxText change on top of AI-Hypercomputer/maxtext#5721 (branch parambole/rl-4600-final-maxtext, pinned two commits ago): the SwiGLU backward fused into the down projection's input-gradient GMM, and the query kernel gathered once per attention layer in the backward. In mlperf_397b_256_v7x.sh the block of MaxText keys (part 2) gains: - moe_wo_gmm_fused_dswiglu=true, on every trainer mesh: the input-gradient GMM of the MoE down projection turns its float32 accumulator into d[gate | up] in its epilogue (with moe_router_grad_from_wo_input, also the routing weight's gradient) instead of writing the activation's gradient in bfloat16 for separate passes to read back. The gradients of wi and of everything before it change by bfloat16 rounding. MaxText refuses the key unless that GMM has one n tile (wo_tile_dlhs_mlp_dim = the MoE MLP dim, 1024 in this recipe) and, with moe_wi_gmm_fused_swiglu, wi_tile_fwd_mlp_dim halves to that dim or to a multiple of 128 dividing it (this recipe's 1024 gives two n tiles of 512 per half); it also needs explicit tiles, prefuse_moe_weights, use_tokamax_gmm, use_gmm_v2 and a remat policy that does not save the routed experts' gate or up half. - query_proj_kernel_gather_once=true, at every trainer mesh except the 256-chip 32x2x8 and 64x1x8: the attention query projection keeps the kernel it gathered in the remat replay for its input gradient, so the backward gathers the query kernel once per attention layer instead of twice. The values are the same. MaxText refuses the key unless the backward replays the query projection (remat_policy=custom with query_proj=remat). #2884's block sets query_proj=device at 32x2x8 and 64x1x8, where the backward does not replay the projection, so the key goes in a new default branch of that mesh case instead of on the tuned line: passed there, MaxText would refuse the 256-chip trainer when it starts. Measured on TPU together (F2): the bench cell of this block at 32x2x4 with these two keys, context_parallel_attention_load_balance_hoist_ids and cp-as-ep-dense-no-cp, on the MaxText of AI-Hypercomputer/maxtext#5758 with the change on top of it (4ed82d793, whose src/maxtext is the pinned MaxText's without #5721's three commits of Rohan Bierneni and two model YAMLs). Bench (AI-Hypercomputer/maxtext#5701 driver, Qwen3.5-397B-A17B, 128 TPU v7x, mesh fsdp 32 x context 4 x expert 2, ragged_buffer_factor 2.75, padding first, 9 micro-batches per step, mean of steps 1 to 4, no token dropped), booking r-7a4cc: 80.398 and 80.342 s per step in two pods (o275x-f2-l3-prof in parambole-bf2-book, o275x-f2-l3-r in parambole-bf2r-book), 3,668 and 3,671 tokens/s/chip, regular micro-batch 8.748 and 8.749 s, against 82.373 s per step, 3,580 tokens/s/chip and 8.966 s for the code of #5758 with the keys before these four (o275x-prfc-f1 in parambole-bf1-book). The F2 pods' own control arms did not finish, so the comparison is across pods: -1.974 and -2.030 s per step. Numerics in both pods (nxf2: loss gates off, random advantages, three steps): all six gates pass against that code under the same protocol (nxprf-l3) and against the reference pair (nxref, nxref2); the loss is 0.0304453575 in all three steps, nxprf-l3's to the ten digits logged, and the gradient norm (gate N3) is within 1.97e-4 and 1.23e-4 relative of nxprf-l3's (tolerance 2e-3). Each key's own measurement: - moe_wo_gmm_fused_dswiglu, behind an environment switch on an overlay build with the keys of #2913's recipe, the hoist and padding first, without the fused log-prob head, against a control in another pod: 83.316 -> 82.148 s per step (o275x-c9pfxctl-l3 -> o275x-c9ypfctl-l3), 9.0845 -> 8.9545 s per regular micro-batch. On one TPU v7x chip its kernels at the cell's shapes match a float32 SwiGLU backward to 1.658e-3 relative L2 (parambole-bck-2d4c2), and its numerics run passes all six gates (nxc9ypfckhv). - query_proj_kernel_gather_once: not timed on TPU alone. An AOT compile of the cell with cp-as-ep gathers the query kernel once per attention layer in the backward instead of twice at the same device peak (104.46 GiB), and the MaxText change's test, executed on TPU v5p, gives bit-identical values and gradients with and without it. Not run: with #5721's three commits (#5721's own content with the rule ran at 81.230 and 81.168 s per step in booking r-67724, o275x-rbd-v16 and o275x-rbd-r-v16), at the 256-chip meshes, or end to end. This is b7f7a6570 (branch parambole/rl-trainer-4600-next of google/tunix) without its third value, custom_mesh_and_rule=cp-as-ep-dense-no-cp, which #2884 already passes, ported onto #2884 and replayed onto its head b16ee52. The port keeps #2884's structure: the header (#2884 rewrote it for the 256-chip meshes; it gains one line naming the rule under TRAINER_TUNED_FLAGS), the end of the block's comment, the tuned line (#2884's, with moe_wo_gmm_fused_dswiglu after the hoist key), the mesh case (the new default branch) and, in the test, the defaults (the two keys join the MaxText-change defaults ahead of #2884's per-tensor remat keys). On replay the same two comments as in the previous commit conflicted (the recipe's sentence on where the keys are defined and the test's comment on _BENCH_MAXTEXT_KEYS), resolved the same way, now naming all three keys. The block's comment on the whole cell replaces "has not run on TPU" with the F2 measurement above. Tests, in the bench image (dbdcf5f8dcba; CPU, no network, the tree read-only, stubs for kubectl, gcloud and gsutil first on PATH): - mlperf_submission_launcher_test.py: 48 of 48 pass at this commit, at each of the two commits before it, and at #2884's b16ee52 and 8647d69. - Five deliberately wrong copies of the recipe each fail exactly the tests predicted before the run: moe_wo_gmm_fused_dswiglu or context_parallel_attention_load_balance_hoist_ids dropped from the block (the 3 exact-set tests); the default branch not passing query_proj_kernel_gather_once (those 3, the default-cell test and the two 256-chip mesh tests); query_proj_kernel_gather_once also on the tuned line (the 4 tests of the 256-chip meshes); moe_wo_gmm_fused_dswiglu passed only with TRAINER_TUNED_FLAGS=false (the 3 and the two switch-off tests). This test fails 6 tests on b16ee52's recipe, and b16ee52's test fails 5 on this recipe, both as predicted. - Rendered, the submission launch differs from the parent's only in the trainer's MaxText flags: moe_wo_gmm_fused_dswiglu=true and query_proj_kernel_gather_once=true at 32x2x4, moe_wo_gmm_fused_dswiglu=true at 32x2x8 and 64x1x8. With TRAINER_TUNED_FLAGS=false, and the other 397B v7x recipes at their defaults, the renders are byte-identical. - MaxText config check (the rendered trainer JobSet through run_trainer_node's argument parsing and build_maxtext_config to pyconfig.initialize, which runs the validators, with mocked devices): the pinned MaxText (94f11c09d) accepts the trainer at 32x2x4 (256 devices; all three keys on, query_proj remat), at 32x2x8 and 64x1x8 (512 devices; query_proj_kernel_gather_once off, query_proj device) and with TRAINER_TUNED_FLAGS=false. It refuses the 32x2x4 trainer on the previous pin bbd0f6f10, which does not define context_parallel_attention_load_balance_hoist_ids, and refuses query_proj_kernel_gather_once at 32x2x8 next to the block's query_proj=device ("query_proj_kernel_gather_once=True requires a remat policy that replays the query projection").
…LAGS at 32x2x4 - requirements: pin MaxText to e1cd8bf1d, the new head of AI-Hypercomputer/maxtext branch parambole/rl-4600-final-maxtext (94f11c09d + the fp8_fake_quant_hoist commit + its test). The key defaults to off, so the pin alone changes no rendered launch. - mlperf_397b_256_v7x.sh: TRAINER_TUNED_FLAGS=true passes fp8_fake_quant_hoist=gather8 only at the 128-chip 32x2x4 trainer mesh, the one it was timed at, and only with TRAINER_FP8=true (fp8_moe_fake_quant=true), without which MaxText raises when the trainer traces it. The fp8 fake-quant absmax of every routed-expert weight is computed once for the layer stack, before the scan, and the expert weights' FSDP all-gathers move fp8 codes; the same values (bit-identical on CPU). Measured on the RL trainer bench at 32x2x4, same pod: 80.385 -> 79.397 s/step (3,669 -> 3,714 tokens/s/chip), 0 tokens dropped. - mlperf_submission_launcher_test.py: the bench cell gains exactly this key, once; it is absent with the tuned flags off, at the 256-chip meshes and with TRAINER_FP8=false.
parambole
changed the title
parambole
changed the title
…r_count=2 at 32x2x4 - requirements: pin MaxText to 15285fd74, the new head of AI-Hypercomputer/maxtext branch parambole/rl-4600-final-maxtext (e1cd8bf1d + the wi_tile_fwd_rhs_buffer_count commit + its test). The key defaults to 0 (the kernel default), so the pin alone changes no rendered launch. - mlperf_397b_256_v7x.sh: with TRAINER_TUNED_FLAGS=true, only at the 128-chip 32x2x4 trainer mesh, the one they were timed at, the recipe's own wi_tile_fwd_mlp_dim becomes 2048 (one value in the flags, not an override) and the block passes wi_tile_fwd_rhs_buffer_count=2. The fused SwiGLU wi forward GMM then covers each half of [gate | up] with one N tile instead of two; with its default three weight buffers that tile does not fit VMEM (66.95 of 57.60 MiB), with two it does (50.95 MiB). Buffering only: the same values. Measured on the RL trainer bench at 32x2x4, same pod: 79.342 -> 79.024 and 78.985 s/step (3,717 -> 3,732 and 3,734 tokens/s/chip), 0 tokens dropped. - mlperf_submission_launcher_test.py: the bench cell gains exactly these two values, once each; the tile stays 1024 and the buffer count absent with the tuned flags off and at the 256-chip meshes.
- requirements: pin MaxText to 1bf840cff, the new head of AI-Hypercomputer/maxtext branch parambole/rl-4600-final-maxtext (15285fd74 + the GDN context-parallel scan commits, their key gating and its test). Both new keys default to off, so the pin alone changes no rendered launch. - mlperf_397b_256_v7x.sh: with TRAINER_TUNED_FLAGS=true, only at the 128-chip 32x2x4 trainer mesh, the one it was timed at, the block passes gdn_cp_m_scan=true: with gdn_cp_mode=seq the GDN context-parallel forward folds each rank's chunk transitions with one Pallas scan in VMEM instead of an XLA reduction tree. Measured on the RL trainer bench at 32x2x4, same pod: 79.480 -> 78.822 s/step (3,711 -> 3,742 tokens/s/chip), 78.795 in a replicate in another pod, 0 tokens dropped. Its backward twin gdn_cp_dh_scan stays off: alone it made the step slower (79.403 -> 80.610). - mlperf_submission_launcher_test.py: the bench cell gains exactly this key, once; it is absent with the tuned flags off and at the 256-chip meshes, and gdn_cp_dh_scan is in no render.
parambole
changed the title
parambole
changed the title
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #2884 at
b16ee522f. That commit is pushed here asparambole/rl-4600-final-tunix-basebecause #2884's branch lives in a fork, so this PR's diff is only its six commits; once #2884 merges, retarget toatwigg/mlperf. Companion MaxText PR: AI-Hypercomputer/maxtext#5799, which defines the keys.With
TRAINER_TUNED_FLAGS=true, the Qwen3.5-397B v7x recipe now turns on six MaxText trainer keys from #5799 and, at the 128-chip mesh, sets its wi forward tile to 2048; the MaxText pin moves to #5799's head1bf840cff, which defines the keys. On the RL trainer bench (128 TPU v7x) the settings before the last two take 79.34-79.54 s per step. The last two were each measured without the other: the 2048 wi tile at 78.99-79.14 s per step = 3,727-3,734 tok/s/chip (-0.32 and -0.36 s/step against a control in the same pod), andgdn_cp_m_scanat 78.80-78.82 s per step = 3,742-3,743 tok/s/chip (-0.66 s/step in the same pod). Together they measured 78.35 s per step = 3,764 tok/s/chip (-1.01 s/step against a control in the same pod; under Results). All six numerics gates pass for each lever; a numerics run of the combination was cancelled before it ran. This PR's base (#2884 with MaxTextbbd0f6f10, i.e. #5721's code with the dense mesh rule) measured 81.17-81.23 s (3,631-3,633 tok/s/chip) in an earlier booking, so this PR adds about -2.0 to -2.2 s/step with the tile and -2.4 with the m scan across bookings, each without the other, and -2.8 to -2.9 with both. The locked #5758 code measured 82.33 s (3,582 tok/s/chip).TRAINER_TUNED_FLAGS=false, and the other 397B v7x recipes at their defaults, render byte-identically to before.What the recipe passes to the trainer
flowchart LR classDef pr fill:#dbeafe,stroke:#1d4ed8,stroke-width:2px,color:#0b1f4d classDef same fill:#f3f4f6,stroke:#9ca3af,color:#6b7280 T["TRAINER_TUNED_FLAGS=true"]:::same --> M{"trainer mesh"} M -->|"128 chips, 32x2x4 (default)"| K128["moe_wo_gmm_fused_dswiglu=true<br/>context_parallel_attention_load_balance_hoist_ids=true<br/>query_proj_kernel_gather_once=true<br/>fp8_fake_quant_hoist=gather8 (only with TRAINER_FP8=true, the default)<br/>wi_tile_fwd_mlp_dim=2048 (the recipe's 1024 replaced)<br/>wi_tile_fwd_rhs_buffer_count=2<br/>gdn_cp_m_scan=true"]:::pr M -->|"256 chips, 32x2x8 or 64x1x8"| K256["moe_wo_gmm_fused_dswiglu=true<br/>context_parallel_attention_load_balance_hoist_ids=true<br/>(these meshes keep query_proj=device, so no gather-once)"]:::pr P["requirements/maxtext_requirements.txt"]:::same --> PIN["MaxText pin: 1bf840cff<br/>head of maxtext PR 5799"]:::prcontext_parallel_attention_load_balance_hoist_ids=true moe_wo_gmm_fused_dswiglu=true query_proj_kernel_gather_once=true fp8_fake_quant_hoist=gather8 wi_tile_fwd_rhs_buffer_count=2 gdn_cp_m_scan=true(fp8_fake_quant_hoistonly withTRAINER_FP8=true, the default), andwi_tile_fwd_mlp_dim=2048in place of the recipe's 1024context_parallel_attention_load_balance_hoist_ids=true moe_wo_gmm_fused_dswiglu=truequery_proj_kernel_gather_onceis left out at the two 256-chip meshes: they keepquery_proj=device, and MaxText refuses the key there.fp8_fake_quant_hoistneedsfp8_moe_fake_quant=true, whichTRAINER_FP8=truesets. Without it, MaxText raises when the trainer traces the layer scan, so the recipe leaves the key out whenTRAINER_FP8=false.wi_tile_fwd_mlp_dim=2048withwi_tile_fwd_rhs_buffer_count=2: the fused SwiGLU wi forward GMM covers each half of its [gate | up] output with one N tile instead of two. At 2048 the kernel's default three weight buffers do not fit VMEM (66.95 of 57.60 MiB) and two do (50.95 MiB), so the two values come together. A gate before the recipe's MaxText flags sets the recipe's own tile, so the trainer's flags carry one value of it, not 1024 followed by 2048. Neither depends onTRAINER_FP8.gdn_cp_m_scan=true: withgdn_cp_mode=seq(set by this block), the GDN context-parallel forward folds each rank's chunk transitions with one Pallas scan in VMEM instead of an XLA reduction tree. Its backward twingdn_cp_dh_scanis never passed: it made the step slower on the bench.Results
RL trainer bench (AI-Hypercomputer/maxtext#5701), Qwen3.5-397B-A17B, 128 TPU v7x, mesh fsdp 32 x context 4 x expert 2. The bench cell uses padding first, the fused log-prob head and the staged step boundary. 9 micro-batches per step, mean of steps 1-4, 0 dropped tokens.
parambole-bf2-v18/o275x-prfc-f2bbd0f6f10(#5721's code with the dense rule; this PR's base)parambole-rb4-v16/o275x-rbd-v16,o275x-rbd-r-v16(booking r-67724)nxrbd-v16)parambole-bf2-v18/o275x-f2-l3-prof,o275x-f2-l3-rnxf2)94f11c09d(this PR's pin before commit 4) with this PR's flags before commit 4parambole-bpr-v18/o275x-pr-v18,o275x-pr-r-v18;parambole-v19a/o275x-v19a-ctlnxpr-v18)e1cd8bf1dwithfp8_fake_quant_hoist=gather8(this PR before commit 5)parambole-v19a/o275x-v19a-fq8(same pod aso275x-v19a-ctl);parambole-v19b/o275x-v19b-fq8r(new pod); booking r-b6424:parambole-w21ds/o275x-d1-ctl,parambole-w21cm/o275x-c1-ctl(two more pods)nx-v19b-fq8)15285fd74with the two tile keys off (commit 4's flags), same pod as the next rowparambole-w21ec/o275x-ec-ctl15285fd74, withwi_tile_fwd_mlp_dim=2048andwi_tile_fwd_rhs_buffer_count=2parambole-w21ec/o275x-ec-e6k,o275x-ec-e6kr(same pod aso275x-ec-ctl);parambole-w21e6k/o275x-e6k1,parambole-w21p3/o275x-p3-e6k(two more pods)nx-e6k1)parambole-w21g2a/o275x-g2-ctlgdn_cp_m_scan=trueparambole-w21g2a/o275x-g2-m(same pod aso275x-g2-ctl);parambole-w21gr/o275x-gr-m(new pod)nx-gm-m)gdn_cp_dh_scan=truealone; both scansparambole-w21g3/o275x-g3-dh(same pod aso275x-g3-ctl, 79.403);parambole-w21g2a/o275x-g2-both(same pod aso275x-g2-ctl)This PR's head, the tile and the m scan together, ran in one pod with its own control (job
parambole-w22k1a), from an overlay whose ten files are this branch's at1bf840cff(tgz4b980f9c, md5-checked). A replicate and a numerics arm of the combination were planned inparambole-w22k1bbut cancelled before they ran (chip budget); each lever passed the numerics gates on its own:TPU combined:
parambole-w22k1a/o275x-k1-ctl(this overlay, all keys off) 79.361 s/step (3,716 tok/s/chip),o275x-k1-em(the tile and the m scan, same pod) 78.349 s/step (3,764 tok/s/chip): -1.012 s/step. The prediction written before the run was -1.00 (band -1.30 to -0.70); the two levers' separate same-pod deltas sum to -0.996. Logs: md5 a17b17b3 (control), 8d7b0411 (combined).gdn_cp_dh_scanmade the step slower: +1.207 s/step alone, +0.672 next to the m scan.fp8_fake_quant_hoist=gather8. The tile arms passedwi_tile_fwd_mlp_dim=2048 wi_tile_fwd_rhs_buffer_count=2after the recipe's flags (MaxText keeps the last value: their config logs read 2048 and 2); commit 5 renders the same config with one value of the tile. The m-scan arms passedgdn_cp_m_scan=trueafter the recipe's flags; commit 6 puts it at the end of the block's flags, the same config. The bench adds its ownenable_checkpointing=false.rl/common.pyandrl/packing.py, come from [Qwen3.5-397B MLPerf] Fused sharded logps head, start-padded packing, and 128c/256c trainer recipe defaults #2884. This PR does not change them.Commits
d1f0fa44ebbd0f6f10->94f11c09d, then the head of #5799. The new keys default to off, so this commit changes no rendered launch.d88ed4b48context_parallel_attention_load_balance_hoist_ids=trueon every trainer meshddacc81afmoe_wo_gmm_fused_dswiglu=trueon every trainer mesh, andquery_proj_kernel_gather_once=trueon every mesh except the 256-chip 32x2x8 and 64x1x8ff2f9cbcf94f11c09d->e1cd8bf1d(+fp8_fake_quant_hoist+ its test), andfp8_fake_quant_hoist=gather8only at 32x2x4 withTRAINER_FP8=true6de2ec73de1cd8bf1d->15285fd74(#5799's head: +wi_tile_fwd_rhs_buffer_count+ its test), and at 32x2x4 withTRAINER_TUNED_FLAGS=truethe recipe'swi_tile_fwd_mlp_dimat 2048 withwi_tile_fwd_rhs_buffer_count=2a94b7e98015285fd74->1bf840cff(#5799's head: + the GDN context-parallel scans, their keys and the keys' test), andgdn_cp_m_scan=trueonly at 32x2x4 withTRAINER_TUNED_FLAGS=true;gdn_cp_dh_scanis never passedTesting
mlperf_submission_launcher_test.py, run on CPU in the bench image with stubs forkubectl,gcloudandgsutil.test_fp8_fake_quant_hoist_only_at_the_timed_cell(commit 4),test_wi_tile_2048_with_two_weight_buffers_only_at_the_timed_mesh(commit 5) andtest_gdn_cp_m_scan_only_at_the_timed_mesh(commit 6).TRAINER_FP8guard removed;gdn_cp_m_scanat every mesh;gdn_cp_m_scandropped;gdn_cp_dh_scanpassed next to it;gdn_cp_m_scanpassed also withTRAINER_TUNED_FLAGS=false.wi_tile_fwd_mlp_dim1024 -> 2048 in place and one added token,wi_tile_fwd_rhs_buffer_count=2(also withTRAINER_FP8=false).gdn_cp_m_scan=true(also withTRAINER_FP8=false).TRAINER_TUNED_FLAGS=falseand the other recipes at their defaults are byte-identical.run_trainer_node,build_maxtext_configandpyconfig.initialize.1bf840cffaccepts 32x2x4, 32x2x8, 64x1x8, tuned off,TRAINER_FP8=falseand commit 5's render. At 32x2x4 it resolvesgdn_cp_m_scan=Truewithgdn_cp_dh_scan=False, andwi_tile_fwd_mlp_dim=2048withwi_tile_fwd_rhs_buffer_count=2(also withTRAINER_FP8=false), elsewhere False, 1024 and 0; MaxText's own fp8 mode check returnsgather8.15285fd74gdn_cp_m_scan,e1cd8bf1dwi_tile_fwd_rhs_buffer_count, and94f11c09dfp8_fake_quant_hoist.TRAINER_FP8guard, run withTRAINER_FP8=false, makes MaxText's mode check raise. That is why the guard is there.gdn_cp_dh_scan=true; only the launcher test keeps the losing scan out of the recipe.Not done / known gaps
parambole-w22k1a). No replicate and no numerics run of the combination: both were cancelled before they ran. Each lever passed the numerics gates on its own.TRAINER_FP8=false, which they do not depend on; that combination was not timed.mlperf_397b_512_v7x.sh,mlperf_397b_1024_v7x.sh,mlperf_gbs512.shandmlperf_397b_256_v7x.sh, whereTRAINER_TUNED_FLAGSdefaults to false, also pass them when run withTRAINER_TUNED_FLAGS=trueat their default 32x2x4 mesh. That combination was not timed.b16ee522f). If [Qwen3.5-397B MLPerf] Fused sharded logps head, start-padded packing, and 128c/256c trainer recipe defaults #2884 moves, this needs a rebase.Logs
gs://parambole-397b-replay-nap/win-tpu/sched/final/<job>/<arm>.log, md5 prefixes below.parambole-w21g2a/o275x-g2-m, controlo275x-g2-ctl, both scanso275x-g2-bothparambole-w21gr/o275x-gr-mparambole-w21g3/o275x-g3-dh, controlo275x-g3-ctlparambole-w21gm/nx-gm-mparambole-w21ec/o275x-ec-e6k, replicateo275x-ec-e6kr, controlo275x-ec-ctl(booking r-b6424)parambole-w21e6k/o275x-e6k1,parambole-w21p3/o275x-p3-e6kparambole-w21ds/o275x-d1-ctl,parambole-w21cm/o275x-c1-ctlparambole-w21e6k/nx-e6k1parambole-v19a/o275x-v19a-fq8parambole-v19a/o275x-v19a-ctlparambole-v19b/o275x-v19b-fq8rparambole-v19b/nx-v19b-fq8parambole-bpr-v18/o275x-pr-v18,o275x-pr-r-v18,nxpr-v18gs://parambole-397b-replay-nap/win-tpu/sched/v16/parambole-rb4-v16/o275x-rbd-v16,o275x-rbd-r-v16,nxrbd-v16parambole-bf2-v18/o275x-prfc-f2,o275x-f2-l3-prof,o275x-f2-l3-r,nxf2