Repository navigation
Conversation
cp-as-ep shards the dense weights (attention and shared-expert projections,
the MoE router and the shared-expert gate) over fsdp x context x expert, so
every all-gather of them spans all devices and synchronizes the context ranks,
which do unequal attention work. The two new rules keep cp-as-ep's
activations, vocabulary weights and routed experts and take those weights off
context: cp-as-ep-dense-no-cp shards them over fsdp x expert,
cp-as-ep-dense-fsdp over fsdp only. Their all-gathers then run in independent
groups of 64 or 32 devices on the 256-device mesh, at the cost of replicating
the weights and their gradients over the axes they leave.
- RoutedMoE keeps ("context", "expert") as its expert-parallel axes under both
rules, as it does for cp-as-ep.
- Left to GSPMD, the gradient of a weight replicated over context is
all-reduced at full size over every device and then sliced. Under a rule set
that maps the new logical axis embed_grad, linears.route_dense_weight_grad
reduce-scatters it over every token axis instead, as under cp-as-ep, and
all-gathers the shard over the replicated axes. An optimization barrier
between the two keeps XLA:TPU from folding them back into the all-reduce.
DenseGeneral and GateLogit apply it to kernels with an embed or embed_attn
dimension.
- The GatedDeltaNet projections take their embed axis from embed_gdn when the
rules map it. Both new rules keep them on fsdp x context x expert: their
all-gathers are already overlapped, and replicating them would cost far more
HBM than the rest of the dense weights.
Rule sets that map neither embed_grad nor embed_gdn compile to the same
program as before.
The CPU test runs a tiny hybrid Qwen3.5 on 16 devices (fsdp 2, context 4,
expert 2) under cp-as-ep and both rules: exactly the dense weights move off
context, the loss is bit-identical, and the gradients agree within bfloat16
reduction order (1e-5 with float32 activations). A rule without the RoutedMoE
change and a perturbed weight both fail the check.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces two new custom mesh rules, 'cp-as-ep-dense-no-cp' and 'cp-as-ep-dense-fsdp', which optimize dense weight sharding by taking them off the context axis. It implements the necessary gradient routing logic in linears.py, integrates the rules into the MoE and Qwen3 models, and adds a unit test. Feedback was provided to propagate sys.path to PYTHONPATH in the test's subprocess environment to prevent potential import failures in CI/CD environments.
Comment on lines
+87
to
+90
| env = os.environ.copy() | ||
| env["XLA_FLAGS"] = env.get("XLA_FLAGS", "") + f" --xla_force_host_platform_device_count={_NUM_DEVICES}" | ||
| env["JAX_PLATFORMS"] = "cpu" | ||
| result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False) |
There was a problem hiding this comment.
In certain CI/CD environments (such as Bazel/Blaze or custom test runners), sys.path is dynamically populated with temporary directories containing the dependencies, but these are not exported to os.environ["PYTHONPATH"]. When subprocess.run is called with sys.executable, the child process starts with a clean sys.path (only containing standard library and site-packages), which will fail to import maxtext or other dependencies, leading to a test failure.
To ensure the child process inherits the exact same module search path as the parent process, propagate sys.path to PYTHONPATH in the copied environment.
Suggested change
| env = os.environ.copy() | |
| env["XLA_FLAGS"] = env.get("XLA_FLAGS", "") + f" --xla_force_host_platform_device_count={_NUM_DEVICES}" | |
| env["JAX_PLATFORMS"] = "cpu" | |
| result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False) | |
| env = os.environ.copy() | |
| env["PYTHONPATH"] = os.pathsep.join(sys.path) | |
| env["XLA_FLAGS"] = env.get("XLA_FLAGS", "") + f" --xla_force_host_platform_device_count={_NUM_DEVICES}" | |
| env["JAX_PLATFORMS"] = "cpu" | |
| result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False) |
Rohan-Bierneni
pushed a commit
to Rohan-Bierneni/tunix
that referenced
this pull request
Oct 16, 2026
…UNED_FLAGS Cherry-picked from google#2940 (2cfa895), companion to AI-Hypercomputer/maxtext#5783 (20c93dc4f), which AI-Hypercomputer/maxtext#5721 now carries in place of its earlier 2-line cp-as-ep.yml edit. TRAINER_TUNED_FLAGS=true now also passes custom_mesh_and_rule=cp-as-ep-dense-no-cp after the recipe's default custom_mesh_and_rule=cp-as-ep (MaxText keeps the last value of a repeated key): the dense weights (attention and shared-expert projections, the MoE router and the shared-expert gate) are sharded over fsdp x expert and replicated over context, their gradients reduce-scattered over every device and then all-gathered over context (linears.route_dense_weight_grad); the GDN projections (embed_gdn), routed experts, embeddings and activations stay as under cp-as-ep. TRAINER_TUNED_FLAGS=false renders as before. 128 TPU v7x (maxtext#5701 bench, from google#2940): 82.185 -> 81.253 s/step (3,588 -> 3,630 tok/s/chip), numerics PASS 6/6. Moves requirements/maxtext_requirements.txt to AI-Hypercomputer/maxtext#5721 at d66a82096.
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
1 participant
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.
Draft until #5758 lands; lint at the CI pins has not run yet. Follow-up to #5758, stacked on its head
parambole/rl-trainer-4600(20c877aa1), so the diff shows only this commit,20c93dc4f. Retarget it when #5758 lands.Two custom mesh rules, both
cp-as-epwith the dense weights taken off thecontextaxis. Undercp-as-ep, the attention and shared-expert projections, the MoE router and the shared-expert gate are sharded over fsdp x context x expert. On the 128-chip v7x bench mesh (256 devices, fsdp 32 x context 4 x expert 2), every all-gather of them spans all 256 devices and also synchronizes the four context ranks.cp-as-ep-dense-no-cpshards them over fsdp x expert and replicates them over context, so those all-gathers run in four independent groups of 64 devices.cp-as-ep-dense-fsdpshards them over fsdp only (32-way). The GDN projections, routed experts, embedding, LM head and activations are sharded as incp-as-ep.On the RL trainer bench (#5701; Qwen3.5-397B-A17B, 128 TPU v7x), with #5758's code and the companion tunix recipe (google/tunix#2913), changing only
custom_mesh_and_rulefromcp-as-eptocp-as-ep-dense-no-cptakes a step from 82.185 s (mean of four runs in three pods) to 81.253 s (mean of two runs in two pods, 81.261 and 81.245 s): 3,588 -> 3,630 tok/s/chip (+1.1%), with the bench's numerics checks passing in both pods. The regular micro-batch goes from 8.94-8.96 s to 8.838 s in both pods. The prediction written before the run said the opposite: 0 to +0.5 s per step, probably slower. Why it is faster is only partly established (see Results).What moves
The decoder's weights on the bench mesh. Blue: moved by
cp-as-ep-dense-no-cp. Gray: unchanged.flowchart TB classDef pr fill:#dbeafe,stroke:#1d4ed8,stroke-width:2px,color:#0b1f4d classDef same fill:#f3f4f6,stroke:#9ca3af,color:#6b7280 subgraph MOVED["moved: fsdp x context x expert (256-way) -> fsdp x expert (64-way), replicated over context"] direction LR ATT["full attention q, k, v, out<br/>15 layers, 1.57 B params"]:::pr SHX["shared expert wi_0, wi_1, wo<br/>and its gate, 60 layers, 0.75 B params"]:::pr RTR["MoE router<br/>60 layers, frozen in the RL recipe"]:::pr end subgraph KEPT["unchanged: as in cp-as-ep"] direction LR GDN["GDN in_proj_qkvz, in_proj_ba, out_proj<br/>45 layers, 5.31 B params, 256-way"]:::same EXP["routed experts<br/>experts over context x expert"]:::same VOC["embedding, LM head<br/>256-way"]:::same end GRAD["gradients of the moved weights:<br/>256-way reduce-scatter as in cp-as-ep,<br/>then a 4-way all-gather over context"]:::pr ATT --> GRAD SHX --> GRAD style MOVED fill:none,stroke:#1d4ed8 style KEPT fill:none,stroke:#9ca3afWhat changes
src/maxtext/configs/custom_mesh_and_rule/cp-as-ep-dense-no-cp.yml(new)cp-as-ep.ymlwithembed,embed_attn,q_loraandkv_loraon['fsdp', 'expert'], plus two new logical axes,embed_gdnandembed_grad, both on['fsdp', 'context', 'expert']src/maxtext/configs/custom_mesh_and_rule/cp-as-ep-dense-fsdp.yml(new)['fsdp']src/maxtext/common/common_types.pyCustomRule.CP_AS_EP_DENSE_NO_CPandCustomRule.CP_AS_EP_DENSE_FSDPsrc/maxtext/layers/moe.pyRoutedMoEkeeps("context", "expert")as its expert-parallel axes under both rules, as undercp-as-ep(lines 1179-1185);GateLogitroutes the router kernel's gradient (line 891)src/maxtext/layers/linears.pyroute_dense_weight_grad(lines 148-206), called byDenseGeneralafter its two-stage all-gather (line 510)src/maxtext/models/qwen3.pyembed_gdnwhen the rules map it, otherwise fromembed_attnas before (line 700)docs/guides/optimization/custom_mesh_and_rule.mdtests/unit/cp_as_ep_dense_no_cp_test.py(new)expertalone. On CPU the loss is then off by 28.5 (a positive control, see Testing). The bench program still compiles, with no error, but its ring-of-experts collectives become 2-way instead of 8-way.route_dense_weight_grad). A weight's gradient is a partial sum over every mesh axis the tokens are sharded on: fsdp, context and expert. Once the weight is replicated over context, GSPMD reduces that partial sum with a full-size all-reduce over all 256 devices and keeps a slice of the result. In the compiled bench program, each backward scan iteration then has full-size all-reduces of the query gradient (bf16[4096,32,512], 128 MiB), the attention output gradient (bf16[8192,4096], 64 MiB) and one 104 MiB tuple of the 12 shared-expert gradients and the key and value gradients, wherecp-as-ephas reduce-scatters: twice the bytes, and still a 256-way synchronization. When the rules mapembed_grad, the backward of every kernel with anembedorembed_attndimension first constrains the gradient to the kernel's axes with that dimension onembed_grad. That is one reduce-scatter over all 256 devices, as undercp-as-ep. It then constrains the result to the kernel's own sharding: a 4-way all-gather of the shard over context. The forward is unchanged. Under a rule set that does not mapembed_grad, the function returns the kernel unchanged. It raises if the rules mapembed_gradand no mesh is set withjax.set_mesh.cp-as-ephas, plus 16 4-way all-gathers over context.embed_gdn). They hold 5.31 B parameters, 2.3x the moved attention and shared-expert weights. See Not in this PR.embed_gradnorembed_gdnare unaffected. At the bench cell withcp-as-ep, this commit compiles to the same program as Qwen3.5-397B RL trainer on 128 v7x: MoE and GDN speed options as config keys #5758, apart from source locations inside 7 Mosaic kernel bodies (see Testing).Results
RL trainer bench (#5701): Qwen3.5-397B-A17B on 128 TPU v7x, mesh fsdp 32 x context 4 x expert 2, 9 micro-batches per optimizer step, 294,912 tokens per chip per step. Seconds per step: mean of steps 1-4 (step 0 compiles). Regular micro-batch: median of fwd_bwd dispatch entries 3-9 over steps 1-4. Every arm runs MaxText at #5758's head
20c877aa1and the bench image's tunix with google/tunix#2913's changed files atbee63532f, with the same 90 MaxText keys apart fromcustom_mesh_and_rule. The arms of this PR also have this commit's six source files installed over that MaxText tree.o275x-prf-l3-prof(bprf, profiled)cp-as-epo275x-prf-l3-r(bprf)cp-as-epo275x-prfc-l3(bl3)cp-as-epo275x-prf-roll4c-prof(bslow, rows rolled by 4, 4 chips profiled)cp-as-epcp-as-epo275x-l2-l3-prof(bl2, profiled)cp-as-ep-dense-no-cpo275x-l2-l3-profreplicate (bl2r, profiled)cp-as-ep-dense-no-cpcp-as-ep-dense-no-cpo275x-prf-l3-prof), and -0.850 and -0.866 s against the fastest. Regular micro-batch: -0.104 to -0.121 s. The fourcp-as-epruns span 0.158 s per step across three pods and the two runs of this PR 0.016 s, so the difference is about six times the larger spread. None of this PR's 56 regular micro-batches (8.829-8.844 s) overlaps any of the 112 of thecp-as-epruns (8.932-8.962 s) (ARITH). Every comparison is across pods; there is no same-pod control.bl2) and 81.262, 81.231, 81.247, 81.241 s (bl2r).Why it is faster: partly established. What the compiled bench program shows (MEASURED, per scan iteration of 4 layers): the moved weights' all-gathers become 64-way, 20 in the forward and 26 in the backward. That leaves 9 and 18 256-way all-gathers, all of them the GDN projections'. The backward keeps its 25 reduce-scatters over 256 devices and adds 16 4-way all-gathers over context. The prediction expected the added replication and the extra all-gathers to cost more than taking the context ranks out of the gather groups saves.
What the TPU profiles show (MEASURED;
bl2's profile againsto275x-prf-l3-prof's, the same two devices of one chip, across pods):cp-as-epeach runs in one group of 256. The GDN projections' all-gathers stay 256-way and the routed-expert collectives are unchanged.INFERRED: the chip that sets the pace reaches these synchronization points sooner; the profiled chip only shows where it waits less. A profile of the pacing chip under this rule would separate the causes; none was taken.
Memory
Compiled for the bench cell (AOT on CPU for 128 TPU v7x chips, the bench mesh and compiler flags, including the bench's
--xla_tpu_max_hbm_size_mib=84992, which is 83.0 GiB). Per device:cp-as-ep(#5758)cp-as-ep-dense-no-cpcp-as-ep-dense-fsdpcp-as-ep-dense-no-cpand +0.64 GiB forcp-as-ep-dense-fsdp. The rule that replicates more adds less, so the temporaries track the schedule, not the replication (INFERRED).How to turn it on
The mesh axes (data, stage, fsdp, context, expert) and every other logical axis rule are
cp-as-ep's.custom_mesh_and_rule=cp-as-ep-dense-fsdpis the 32-way variant. It is compiled and CPU-tested but has not run on TPU, and its value is unknown: the same reasoning that predictedcp-as-ep-dense-no-cpwould be slower also predicted it would be slower.Testing
tests/unit/cp_as_ep_dense_no_cp_test.py, new, CPU): a tiny hybrid Qwen3.5 (GatedDeltaNet kernel path, gated full attention, routed and shared experts, ring of experts, sequence-sharded GDN context parallelism) on 16 CPU devices, fsdp 2 x context 4 x expert 2, undercp-as-epand both rules, from the same parameter values. It runs in a child process. On the commit's own sources it passes all 9 checks, rc 0:cp-as-epsharding.cp-as-ep-dense-no-cp), and all-gathers over 2 devices from 24 to 74 (cp-as-ep-dense-fsdp).cp-as-epon this commit vs Qwen3.5-397B RL trainer on 128 v7x: MoE and GDN speed options as config keys #5758 (20c877aa1): the optimized fwd_bwd HLO is identical (55,912 lines) once op metadata, stack-frame tables and Mosaic kernel bodies are stripped. 7 of the 110 kernel bodies differ by a few bytes at equal length, in what reads as their source-location table (INFERRED: line numbers that moved). Memory is identical (104.46 / 78.28 GiB). The comparator gives the same verdict for Qwen3.5-397B RL trainer on 128 v7x: MoE and GDN speed options as config keys #5758 against itself (0 differing bodies) and a different one againstcp-as-ep-dense-no-cp.cp-as-ephas, plus 16 4-way all-gathers, each behind a small collective-permute (0.02-0.5 MiB).cp-as-ep-dense-fsdpon TPU.TPU validation status
cp-as-ep-dense-no-cpvscp-as-ep, across podsparambole-bl2r-v14/o275x-l2-l3-profparambole-bl2-v14/nxl2andparambole-bl2r-v14/nxl2, againstnxrefandnxref2, withparambole-bl3-v14/nxprf-l3(#5758's code, same protocol) as the same-cell referencebl2's profile (MEASURED)cp-as-ep-dense-fsdpon TPUNumerics A/B (MEASURED, read with the bench's readout). Gates-off runs: log-prob gates off, random advantages, 3 steps. The six checks compare the first step with two reference runs,
nxrefandnxref2: every arm keeps its tokens; loss, relative difference (tolerance 1e-5); mean log-prob, absolute difference (5e-5); the last micro-batch's log-prob, absolute difference (5e-5); entropy, relative difference (1e-4); gradient norm, relative difference (2e-3).nxl2(bl2)nxl2(bl2r)nxprf-l3(#5758's code)bl2) and 0.0557773 / 0.0558073 / 0.0558132 (bl2r), against 0.0557847 / 0.0558033 / 0.0558120 for Qwen3.5-397B RL trainer on 128 v7x: MoE and GDN speed options as config keys #5758's code: at most 1.33e-4 apart (relative).Not in this PR
embed_gradto['fsdp', 'expert', 'context']instead of['fsdp', 'context', 'expert']should leave each 64-way shard's four pieces contiguous after the 256-way reduce-scatter, so the 4-way all-gather would need no permute first (INFERRED; not compiled, not run).q_loraandkv_lora(multi-head latent attention) move with the other dense axes, so the rule's dense set matchescp-as-ep's, but no MLA model was compiled or run.How it was measured: arms, logs, overlay, profiles
r-371ee. Mesh fsdp 32 x context 4 x expert 2; 9 micro-batches per optimizer step; 294,912 tokens per chip per step;ragged_buffer_factor2.75. Recipe: Qwen3.5-397B RL trainer on 128 v7x: MoE and GDN speed options as config keys #5758's keys as Qwen3.5-397B RL trainer on 128 v7x: pad-first sequence packing, fused sharded log-prob head VJP, and the tuned trainer recipe google/tunix#2913'sTRAINER_TUNED_FLAGS, the fused log-prob head (TUNIX_LOGPS_SHARDED_HEAD=fused), pad first (applied by the bench driver), the staged step boundary, and none of the overlay builds' environment switches.o275x-prf-l3-profdiffer only incustom_mesh_and_rule. Their config dumps (932 keys, paths excluded) differ only in that key and in the logical axis rules it selects:embed,embed_attn,q_lora,kv_lorafrom('fsdp', 'context', 'expert')to('fsdp', 'expert'), withembed_gdnandembed_gradadded. The replicate (bl2r) ran the same job manifest with only the names changed, with the same 90 keys.gs://parambole-397b-replay-nap/win-tpu/team-pr-arms/overlay/parambole-lev-l2-8ed9e36874f514521acf64f598fb6546.tgz(md58ed9e36874f514521acf64f598fb6546). Each file in it is byte-identical to its blob at20c93dc4f:common_types.pyd6085955,cp-as-ep-dense-no-cp.yml3f965a8a,cp-as-ep-dense-fsdp.ymlc27af2af,linears.py1952c83a,moe.py13e7196c,qwen3.py600a57d8. All four arm logs of this PR (two timing, two numerics) printcustom_mesh_and_rule: CustomRule.CP_AS_EP_DENSE_NO_CP, an enum value that only this commit defines, and none has a Traceback or RESOURCE_EXHAUSTED.gs://parambole-397b-replay-nap/win-tpu/sched/v14/, with md5:parambole-bl2-v14/o275x-l2-l3-prof.log(916f2fac3ee4c91fe670cce6d2a75e04)parambole-bprf-v14/o275x-prf-l3-prof.log(6c6039ba8d3f470b7070fca0d22be626),parambole-bprf-v14/o275x-prf-l3-r.log(5689ffadccbcfe1428c2eb444bb5b9ef)parambole-bl3-v14/o275x-prfc-l3.log(f87a786767e202d862c5881932bd87dd)parambole-bslow-v14/o275x-prf-roll4c-prof.log(f88f61517c0a1d8629bfb9e71ec93741)parambole-bl3-v14/nxprf-l3.log(ed5d4cfa58dbb17495e6312be07623a8); references../r-06ca7/parambole-bnnp-06ca7/nxref.log(530911504a42b98292fd0c76222cf872) andnxref2.log(acff70f5358e3702cd5ece191fedf2c4)parambole-bl2r-v14/o275x-l2-l3-prof.log(ac16b85d39a0ab19e8f77455651fc228)parambole-bl2-v14/nxl2.log(8851f887348b2db7a3c37c439cbcf40d) andparambole-bl2r-v14/nxl2.log(37e4297b91759bb7feebbf0b15c0b36b), read with the bench'snum_readout.py(md5ba65e0c0, self-test 11 of 11).../parambole-bl2-v14/o275x-l2-l3-prof/parambole-bl2-v14-o275x-l2-l3-prof/tensorboard/plugins/profile/2026_10_15_21_38_03/; its control under.../parambole-bprf-v14/o275x-prf-l3-prof/parambole-bprf-v14-o275x-prf-l3-prof/tensorboard/plugins/profile/2026_10_15_21_24_47/. The replicate's profile was captured and not analysed. Dissection of the first two, CPU only, oncmcs-checkpoint:~/p4600/ag/v14res/l2prof-heavy/(tablesout/l2an.md, collective groupsout/census2.txt, SparseCore time per callout/sccmp.txt).manifests-ncp/v14/levers/parambole-bl2-v14.yaml(md5ef03b68b2863c050f74d711b85c2c5b3; the prediction is header line 9) andparambole-bl2r-v14.yaml(md571a8327524a0984820273d936fdf74ac, a renamed copy; header lines 1-2 state the -0.92 s per step it set out to reproduce).CPU and compile evidence
On the build host
cmcs-checkpoint, bench image (local idsha256:dbdcf5f8dcba6e71d7edbba31d1378057ce22ca4771735923a67d12862c662c5),--network none:~/p4600/ag/l2/cpu/num5.log(md5ced4b54240b4dbad56081c3937bd394d):CP_AS_EP_DENSE_NO_CP_CHECKS_PASSED, rc 0. The script's md5 (7f4ad15e) and the tree's six source files equal the committed blobs.~/p4600/ag/l2/aot/logs/l2g-A.log(cp-as-ep-dense-no-cp, md571f775a7372b5e569c55147892ffa474),l2g-B.log(cp-as-ep-dense-fsdp,df96d92d7d65ab54bf8d9f6dd421276a),l2g-cpaep.log(cp-as-epon this commit,7a6889d001bd57f5ea35891b4165a04a), each rc 0. Memory-usage reports and optimized HLO under~/p4600/ag/l2/aot/dumps/l2g-*/; Qwen3.5-397B RL trainer on 128 v7x: MoE and GDN speed options as config keys #5758's compile of the same cell under~/p4600/ag/base/aot/dumps/prcell-base/.~/p4600/ag/l2/ana/l2g-{A,B}/census_agg.mdagainst~/p4600/ag/base/ana/base/census_agg.md. Compared with the rest gate~/p4600/ag/l2/tools/restgate.py(md5a6873957) and the HLO comparatortools/bodydiff.py.~/p4600/ag/l2/ana/l2f-A/(routing, no barrier) andl2x-Ak/(rules only) both have the full-size all-reduces.Predictions written before measuring, and how they did