From 780e3e45cf766c4c1c4d043b7f2e215778fe478c Mon Sep 17 00:00:00 2001 From: chengnuojin Date: Wed, 2 Sep 2026 16:00:17 +0000 Subject: [PATCH 1/9] Emit the LM head weight gradient in kernel order under explicit sharding Under `shard_mode: explicit`, JAX's `lower_with_sharding_in_types` annotates every sharding-in-types primitive's output with a `custom_call_target="Sharding"` op. One of those lands between the LM head weight-gradient `dot` and the `transpose` that follows it, and XLA's algebraic simplifier will not fold `transpose(dot(A, B))` into `dot(B, A)` across it. Census on the dp4 module: auto has 17 adjacent (foldable) transpose-of-dot pairs, explicit has 0. The transpose survives to codegen, which costs twice -- the gradient is left sharded on its minor-most dimension, below the 512-byte extent `tpu-all-reduce-scatter-fusion` requires, so the weight-gradient all-reduce is never rewritten into a reduce-scatter; and the resulting layout propagates through the Adam update, forcing six relayout copies on the entry parameters. Remove the need for the fold rather than the barrier. `DenseGeneral` gains a `weight_grad_in_kernel_order` option backed by a `custom_vjp` whose backward pass contracts `dk` straight into the kernel's stored axis order, so no transpose is emitted at all. The stored kernel, its sharding and its initialization are untouched, so this is checkpoint-compatible; the gradient differs from the default rule only by floating-point reassociation. It self-disables under `shard_mode: auto`, where XLA folds the transpose on its own, and under quantization, where it cannot own the dot. The untied LM head opts in via `lm_head_weight_grad_in_kernel_order` (default false) on both the linen and nnx decoder paths. Measured on v5p, 4 chips, xprof `train_step_ns` median on `/device:TPU:0`, 56 runs, arms run back-to-back per config so within-row deltas carry no session drift: config auto us explicit +kernel-order B-A% D-A% rep spread deepseek 9081.0 9883.5 9216.8 +8.84 +1.50 0.56% llama2 4027.4 4223.8 3995.3 +4.88 -0.80 0.79% mistral 3905.2 4084.6 3853.0 +4.59 -1.34 0.97% mixtral 10104.5 10407.2 10129.8 +3.00 +0.25 0.12% qwen3 8261.8 8608.6 8214.8 +4.20 -0.57 0.51% w2048 14972.1 15215.1 14886.3 +1.62 -0.57 0.04% prod 196039.4 196558.5 196101.9 +0.26 +0.03 0.03% Read the two halves differently. The five emb-512 rows go from a +4.59% median penalty to -0.57%, but emb 512 / fsdp 4 puts 256 bytes per shard on the minor-most dimension, below XLA's 512-byte AR->RS gate -- so those rows are a benchmark artifact of the shrunk proxies and no production job has that geometry. The ship criterion is the two emb-2048 rows, where the bar is no regression rather than a win: prod +0.03% against a 0.03% rep spread, w2048 -0.57%. Acceptance: on llama2/shrink/fsdp4 the explicit dump gains `%all-reduce-scatter.2`, byte-for-byte identical to auto's `%all-reduce-scatter.5` (`bf16[512,32000] -> bf16[128,32000]`, layout `{1,0:T(8,128)(2,1)}`, dim-0 dynamic-slice), and `f32[emb/n, 32000]` relayout copies in `after_optimizations` drop from 6 to 0. Numerical equivalence over 20 steps against auto, max absolute loss deviation: llama2 explicit 2.62e-04, with the flag 2.52e-04; mixtral 1.41e-04 and 9.63e-05. The flag sits at explicit's own noise level, as reassociation alone predicts. `logits_via_embedding` is rejected rather than silently ignored, since the flag only applies to the untied head. An earlier revision of this branch shipped `lm_head_kernel_transposed`, which stores the kernel as `[vocab, embed]` to reach the same layout. It is dropped here: it is 1.85-4.02% slower than the custom_vjp on four of seven configs because it recovers only the relayout copies and not the reduce-scatter, it has no checkpoint-conversion path, and at the flipped shape `jax.random` draws a different matrix from the same seed, which perturbs the loss ~35x more than explicit's own noise (llama2 9.42e-03 vs 2.52e-04). Analysis, including the measured negative result for eliding redundant reshard barriers, is in docs/guides/optimization/shard_mode_performance.md. --- docs/guides/optimization.md | 8 + .../optimization/shard_mode_performance.md | 1280 +++++++++++++++++ src/maxtext/configs/base.yml | 8 + src/maxtext/configs/types.py | 20 + src/maxtext/layers/decoders.py | 1 + src/maxtext/layers/linears.py | 133 +- src/maxtext/layers/nnx_decoders.py | 1 + tests/unit/linears_test.py | 147 +- 8 files changed, 1596 insertions(+), 2 deletions(-) create mode 100644 docs/guides/optimization/shard_mode_performance.md diff --git a/docs/guides/optimization.md b/docs/guides/optimization.md index bdf9b1e4e6..1bd7aa3b1f 100644 --- a/docs/guides/optimization.md +++ b/docs/guides/optimization.md @@ -49,6 +49,13 @@ Optimize with Pallas kernels for fine-grained control. Guide to setting up benchmarks, performing performance tuning, and analyzing metrics. ``` + +```{grid-item-card} 🔀 Shard Mode: explicit vs auto +:link: optimization/shard_mode_performance +:link-type: doc + +Measured HLO and xprof comparison of `shard_mode: explicit` against `shard_mode: auto` on the onboarded models. +``` ```` ```{toctree} @@ -61,4 +68,5 @@ optimization/sharding.md optimization/custom_mesh_and_rule.md optimization/pallas_kernels_performance.md optimization/benchmark_and_performance.md +optimization/shard_mode_performance.md ``` diff --git a/docs/guides/optimization/shard_mode_performance.md b/docs/guides/optimization/shard_mode_performance.md new file mode 100644 index 0000000000..12a3359725 --- /dev/null +++ b/docs/guides/optimization/shard_mode_performance.md @@ -0,0 +1,1280 @@ + + +# `shard_mode: explicit` vs `shard_mode: auto` — a measured performance analysis + +**Status:** investigation report. Measured 2026-08-31 on a TPU v5p slice (4 JAX devices / 8 TensorCores, +95.7 GiB HBM per device), JAX 0.11.1, MaxText branch `chengnuojin-trainer-fix` @ `671cec44d`. +Revised 2026-09-02 with the shipped fix (`lm_head_weight_grad_in_kernel_order`) and two negative +results; see §8 item 1 and §8 item 2. Where a claim was superseded the original is kept and marked, +because two of the corrections reverse a conclusion that looked solid at the time. + +`shard_mode` (`src/maxtext/configs/base.yml:545`) selects how MaxText expresses sharding: + +- **`auto`** — `jax.lax.with_sharding_constraint`. A *hint*. GSPMD/Shardy is free to propagate it + backwards into producers, refine it, sink it, or drop it. +- **`explicit`** — `jax.sharding.reshard` on an all-`AxisType.Explicit` mesh (JAX sharding-in-types). + A *typed conversion*. The sharding becomes part of the value's type and the transfer is emitted at + exactly that program point. + +The entire switch is one branch, `maybe_shard_with_name` at `src/maxtext/utils/sharding.py:125-129`: + +```python +if shard_mode == ShardMode.EXPLICIT: + return reshard(inputs, named_sharding) +else: + return jax.lax.with_sharding_constraint(inputs, named_sharding) +``` + +plus mesh construction at `src/maxtext/utils/maxtext_utils.py:2426-2429`, which stamps *every* mesh +axis `Explicit` or *every* axis `Auto` — there is no partial adoption. + +______________________________________________________________________ + +## 1. Executive summary + +01. **The root cause of the regression is an optimization barrier, not sharding policy.** Under + explicit, Shardy materializes a `custom-call(custom_call_target="Sharding")` on every pinned + intermediate. On llama2 that is **65 such calls under auto vs 1,276 under explicit** (qwen3: + 65 vs 1,494), and the pre-optimization jaxpr is otherwise near-identical (42 dots and 49 + transposes in both; every non-`Sharding` opcode delta ≤ +33). The custom-call is a semantic + no-op, but it sits *between* the unembedding weight-gradient `dot` and its `transpose`, and XLA's + algebraic simplifier will not fold `transpose(dot(A,B)) → dot(B,A)` across it. Census on the + dp4 module: **auto has 17 adjacent (foldable) transpose-of-dot pairs, explicit has 0.** + +02. **Two concrete costs follow from that one barrier.** + + - **A lost reduce-scatter.** Under auto the vocab weight-gradient `all-reduce` feeds a + `dynamic-slice`, and XLA rewrites the pair into a real reduce-scatter. Under explicit the + gradient is transposed, the FSDP shard dim becomes minor-most, the pattern no longer matches, + and a plain full all-reduce is emitted — **2× the wire bytes for that tensor** (1.5·D vs + 0.75·D on a 4-way ring; 49.15 MB vs 24.58 MB per device on llama2). + - **Six relayout copies per step.** `entry_computation_layout` is identical in both modes and + the optimizer state is donated, so the `{0,1}`-pinned Adam chain must be converted in and out: + three copies for `logits_dense.kernel` / `mu` / `nu` on the way in, three on the way out. + 196.6 MB/step on llama2, up to 1.57 GB/step at emb 4096. **Added copy bytes separate the sign + of the effect perfectly**: the only three configs of 25 that add ~zero (gemma, gemma2, tp4) are + exactly the three that get faster; all 22 that add ≥ 66 MB regress. + +03. **Collective *bytes* are not identical — the byte counter was wrong.** A naive counter + reads the `all-reduce` *inside* auto's `kind=kCustom, calls=%all-reduce-scatter.N` fusion at its + full pre-slice shape, so it reports "unchanged". `after_codegen` shows the truth: + `reduce-scatter` lines are **llama2 12 → 4, mistral 12 → 4, mixtral 12 → 6, qwen3 10 → 4**. + The same fusion is booked as FUSION rather than COLLECTIVE by the xplane classifier, so a large + part of the apparent "collective time doubled" is reclassification. Everywhere this document + previously said "identical bytes", read "identical *static* byte counts, which is an artifact". + +04. **The penalty does not simply amortize with scale.** Its *instruction count* is fixed (always the + same six copies); its *cost* is not, and the ladder reads **non-monotone**: w512 +4.80% → w2048 + +0.96% → w4096 **+1.49%**. The largest absolute penalty in the whole 28-run sweep (**+1111 µs**) + is at the *largest* width. Regressing step delta on measured copy time over 12 configs gives + `Δ = −8.7 + 1.37·copy_µs`, **R² = 0.449** — the copies explain about half the variance. + + > *Correction (2026-09-02).* "Non-monotone" conflates two regimes. w512 is the only rung below + > XLA's 512-byte minor-most AR→RS gate (§4.2), so it pays mechanism A **and** mechanism B; w2048 + > and w4096 pay only B. Within the above-gate regime the ladder is monotone increasing in width + > (+0.96% → +1.49%), which is what the relayout-copy model predicts. The practical consequence is + > sharper than the original claim: **the emb-512 numbers throughout this document overstate what + > a production-geometry job sees**, and any fix validated only at emb 512 is validated on an + > artifact. See §8 item 1. + +05. **The 16-layer near-zero is cancellation, not amortization.** At d16 the six copies still cost a + measured **345.2 µs/step**; they are masked by an unrelated **−242 µs** win inside the layer + scan (and a −390 µs `concatenate` swing). Pooling all four d16 runs per mode gives **+70.7 µs = + +0.088%, t = 4.5** — a small but statistically significant penalty, not noise. + +06. **Tensor parallelism is the only win, and for a different reason entirely.** llama2-tp4's + −0.61% headline is inflated ~3× by module-launch lead time; the true device-side effect is + **−0.21%** (−9.3 µs, reproduced independently on TPU:1). The mechanism is not the vocab tensor + at all — it is the RMSNorm **scale** reshard at `normalizations.py:34-51`, which lets explicit + write the optimizer update straight into the donated parameter buffers and saves auto's 14 + buffer-aliasing copies. + +07. **One headline number in the matrix is a harness artifact, not a model property.** gemma3's + +4.49% happens because the shrink sets `base_num_decoder_layers=4`, which is *less than* + `len(GEMMA3_ATTENTION_PATTERN) == 6`, so `scan_length = 4 // 6 = 0` (`decoders.py:1341`) and the + whole decoder is **unrolled** into `layers_remainder`. gemma3 is the only model in the sweep + compiled with zero `while` loops, and unrolling is what makes explicit lose. Re-running the + same config at more layers: **L4 +4.54% → L6 +1.87% → L12 −0.58%.** Nothing about local/global + attention, the 262144 vocab, or multimodal code is involved. + +08. **`reduced=` / `unreduced=` PartitionSpec tags are a real explicit-only capability but account + for 0.0% of every number in this document.** `moe.py` contains zero occurrences; the tags live + only in `deepseek_batchsplit.py` and `gradient_accumulation.py`, behind a gate + (`gradient_accumulation.py:77`) that requires `data > 1`. Every measured run had `data=1` and + `gradient_accumulation_steps=1`. Nothing in the MoE path participates in any measured delta. + +09. **Explicit forfeits capabilities**: `check_vma` (which `base.yml:704` calls a performance win) is + hard-rejected, `fused_qkv=True` is a trace-time error, context parallelism on the + `dot_product`/GPU-flash paths is a trace-time error, and `P.UNCONSTRAINED` is illegal. + +10. **There is a fix, it is measured, and it ships here off by default.** The barrier itself cannot be + removed — it comes from JAX's `lower_with_sharding_in_types`, which annotates every + sharding-in-types primitive on an all-Explicit mesh, not from any MaxText pin (§4.5; the earlier + "stop pinning the logits output" recommendation was tested and produces a 0-line HLO diff, and + so does eliding the redundant barriers — §8 item 2). What can be removed is the *need* for the + fold: write the LM head's weight gradient by hand so it contracts straight into the kernel's + stored axis order. `lm_head_weight_grad_in_kernel_order: true` does that. It recovers the + reduce-scatter of item 02 — the emitted `%all-reduce-scatter` is byte-for-byte auto's — and + drives the six relayout copies to zero, without touching the stored kernel, the checkpoint, or + initialization. Median across seven configs: explicit alone **+4.59%**, with the flag + **−0.57%**. At production geometry, where there was nearly nothing to recover, it is neutral + (**+0.03%** and **−0.57%**). See §8 item 1. + + > *Superseded (2026-09-02).* This item previously recommended `lm_head_kernel_transposed`, which + > stores the kernel as `[vocab, embed]`. That flag reaches the same layout but is slower on four + > of seven configs, has no checkpoint-conversion path, and — because `jax.random` draws a + > different matrix from the same seed at the flipped shape — perturbs the loss ~35× more than + > explicit's own noise. It has been removed rather than shipped alongside a strictly better flag. + +**Practical guidance:** at realistic depth/sequence/batch the penalty falls to a few tenths of a +percent and capability should decide the choice. But it does **not** reliably vanish with *width*, +and under pure data parallelism it is +10%. On an untied model where you have committed to explicit, +set `lm_head_weight_grad_in_kernel_order: true` (§8) rather than assuming scale will hide it; it is +checkpoint-compatible and inert under `auto`, so it costs nothing to leave on. + +______________________________________________________________________ + +## 2. Scope: which models are actually onboarded + +`src/maxtext/configs/types.py:4142-4168` gates explicit mode to a whitelist: + +```python +supported_decoders = { + "simple", + "simple_mlp", + "llama2", + "deepseek", + "mistral", + "mixtral", + "qwen3", + "qwen3_moe", + "qwen3_custom_moe", + "gemma", + "gemma2", + "gemma3", +} +``` + +plus `use_multimodal` is rejected. Every real model in that set was measured. Exactly one shipped +config defaults to explicit: `src/maxtext/configs/models/deepseek3-671b-batchsplit.yml:61`. + +______________________________________________________________________ + +## 3. Method + +### Harness + +Purpose-built for this study and not checked in; §9 gives the exact command line. Four pieces: + +| piece | what it does | +| ----------------------- | --------------------------------------------------------------------------------------------------------------------------------------------- | +| run driver | runs one (model, mode) with `--xla_dump_to` + `profiler=xplane`, serialized under `flock` | +| xplane decoder | dependency-free XSpace protobuf reader; the installed `tensorboard_plugin_profile` cannot decode xplane on this host | +| collective byte counter | counts collectives by kind, bytes from dtype × shape — **see the caveat below, this one under-reports** | +| HLO differ | signature-bucketed instruction diff (opcode + layout-stripped shape + `dimensions`), since instruction names are unstable across compilations | + +### Timing source + +MaxText does not log per-step device time, so step time is taken from the xprof `XLA Modules` +line: the median duration of the `jit_train_step` module events on `/device:TPU:0`. With +`skip_first_n_steps_for_profiler=3, profiler_steps=3` there are exactly 3 such events per run. +Device-op time and the collective / fusion / data-movement split come from the `XLA Ops` and +`Async XLA Ops` lines, summed over all 8 TensorCores and all 3 profiled steps — so device-op totals +are larger than wall step time and should be read as *ratios*, not as wall clock. + +### Noise floor + +Same config (llama2 decoder, emb 2048, mlp 8192, 16 layers, seq 1024, pdbs 1, fsdp 4), 3 independent +runs per mode: + +| | run 1 | run 2 | run 3 | mean | stdev | +| -------- | ------: | ------: | ------: | ---------: | ---------------: | +| auto | 80366.1 | 80394.8 | 80419.1 | 80393.3 µs | 26.5 µs (0.033%) | +| explicit | 80483.6 | 80477.7 | 80466.0 | 80475.8 µs | 9.0 µs (0.011%) | + +Largest within-mode spread: 53 µs = **0.066%**. Treat |Δ| below ~0.15% as noise on a *single* pair; +with 4 runs per mode pooled, a +0.09% effect is resolvable (§5.3). + +**Cross-validation against a second, independent metric.** On the 21 device-bound configs where the +xprof module span and MaxText's own host-side step time can both be read, the two agree with +**r = 0.96**, and **17 of 21 are slower under explicit on both**. deepseek measures +8.8% device / ++7.1% host. The direction of every headline in this document survives the metric change; only the +sub-1% magnitudes are metric-sensitive (see caveat 3 below). + +### Measurement caveats — read these before trusting any number here + +1. **A static collective byte total is not a measure of communication volume.** The byte counter + used here had two independent defects: + + - It sums the *body* of auto's `kind=kCustom, calls=%all-reduce-scatter.N` fusion, so it reads + the pre-slice `all-reduce` shape and assigns a reduce-scatter and a full all-reduce of the same + operand the same bytes — even though they differ by exactly 2× on the wire. + - It is **trip-count-unweighted**: every collective is counted once regardless of its enclosing + `while`. `scale/d16` (16 layers) and `scale/w2048` (4 layers) report *byte-for-byte identical* + totals (628,146,216 B auto) while true dynamic volume is **4.20 GB vs 1.34 GB**. + + Use `grep -c reduce-scatter *after_codegen*` instead, and weight by trip count. Once weighted, + aggregate per-step collective volume really is nearly mode-invariant across all 25 configs + (|Δ| ≤ 0.70%, max −0.700% at d16) — but that is *aggregate* parity, not per-class parity: + collective-permute volume **drops 50%** under explicit in the five scale-ladder configs + (2 → 1 `collective-permute-start(bf16[48,8192])`, `source_target_pairs={{0,1},{1,2},{2,3}}`), + while gemma2 gains a real extra in-loop `all-gather bf16[1,8,128,512] dimensions={3}` worth + +4 MiB/step. This also means **collective instruction *count* is an anti-signal**: explicit's + lower count reflects auto's wins being fused away. + +2. **The xplane classifier books `fusion.N` as FUSION even when `calls=%all-reduce-scatter.N`.** + Part of every "collective % of device time" jump in §5.1 is this reclassification, not new + communication. Quantified in §4.4. + +3. **`train_step_ns_median` includes a variable module-launch lead gap.** Measured at 30–88 µs + across the 3 profiled steps on TPU:0 — larger than some of the deltas being reported. For + sub-1% results, cross-check against the op span (first op start → last op end) or the TPU:1 + module span. This matters for tp4 (§5.2) and nothing else in this document. + +4. **Shrinking a model can change its compilation structure, not just its size.** gemma3 at + 4 layers takes a completely different code path from gemma3 at 6+ layers (§4.7). **Always check + `while`-loop count parity between the two dumps before comparing them** — every other model in + this sweep compiled to `while = 2` or `6/7` in both modes; gemma3 compiled to `while = 0`. A + shrink that disables the scan is not a proxy for the real model. + +5. **The persistent compilation cache silently suppresses HLO dumps.** MaxText defaults + `jax_cache_dir: "~/jax_cache"` (`base.yml:534`). A cache hit skips compilation, so + `--xla_dump_to` produces nothing. The harness sets `jax_cache_dir=` to force compilation. Timing + is unaffected (the executable is identical), but a sweep that trusts `status: ok` alone will + silently produce empty dumps. + +6. **MaxText's own `dump_hlo` config crashes on a local `base_output_directory`.** + `src/maxtext/utils/gcs_utils.py:89-116` `upload_dump()` calls `parse_gcs_bucket_and_prefix()` + unconditionally, which raises `IndexError: string index out of range` on a non-`gs://` path — and + it fires *after* training but *before* the xplane profile is flushed, so you lose the profile. + The harness drives the dump through `XLA_FLAGS` instead. **This is a real robustness bug.** + +### Measured HBM bandwidth + +The relayout copies in §4.3 are bandwidth-bound and were used to calibrate the platform: +**2.29–2.38 TB/s** effective across the four ≥ 65 MB cases, i.e. **83–86% of the v5p 2765 GB/s +spec**. The small (16 MB) llama2 copies read back at an implied 3.64 TB/s — above spec — which means +they are partly cached or partly overlapped and their profiled durations *understate* their isolated +cost. Only the ≥ 65 MB cases give a trustworthy reading. + +______________________________________________________________________ + +## 4. The dominant mechanism + +Chain, in causal order: **Sharding custom-call barrier → blocked `transpose(dot)` fold → transposed +vocab gradient → (a) lost reduce-scatter and (b) six relayout copies.** + +### 4.1 Root cause: the `Sharding` custom-call is an optimization barrier + +MaxText builds the unembedding as a separate `dense_general` with +`kernel_axes=("embed_vocab", "vocab")` at `src/maxtext/layers/decoders.py:831-844`, so the parameter +is `[emb, V]`. `base.yml:561` maps `embed_vocab → fsdp` and `:560` maps `vocab → [tensor, ...]`; at +fsdp=4 / tensor=1 the kernel is sharded 4 ways on its **contracting** (embed) axis and replicated on +vocab. JAX's `_dot_general_transpose_rhs` computes `dW` as a dot that naturally produces `[V, emb]` +followed by an explicit transpose back to `[emb, V]`. **Both modes emit exactly that pair** in +`before_optimizations`: + +``` +# auto — llama2, lines 2273-2274: dot and transpose are ADJACENT +%dot_general.80 = bf16[32000,512]{1,0} dot(%convert_element_type.147, %transpose.85), + lhs_contracting_dims={0,1}, rhs_contracting_dims={0,1} +%transpose.86 = bf16[512,32000]{0,1} transpose(%dot_general.80), dimensions={1,0} + +# explicit — qwen3, lines 3446-3449: a Sharding custom-call is INTERPOSED +%dot_general.190 = bf16[151936,512]{1,0} dot(...) +%dot_general.191 = bf16[151936,512]{1,0} custom-call(%dot_general.190), + custom_call_target="Sharding", sharding={devices=[1,4]<=[4]}, + xla.sdy.sharding=<[{}, {"fsdp"}]> +%transpose.216 = bf16[512,151936]{0,1} transpose(%dot_general.191), dimensions={1,0} +%transpose.217 = bf16[512,151936]{1,0} custom-call(%transpose.216), + custom_call_target="Sharding", sharding={devices=[4,1]<=[4]} +``` + +llama2's explicit dump is byte-for-byte the same pattern at V=32000. + +The custom-call is semantically a no-op. Its only effect is to stop the algebraic simplifier. The +result is visible in `after_optimizations` as an **unfolded operand order**: + +``` +qwen3/auto:3721 ROOT %convolution.78 = bf16[512,151936]{1,0:T(8,128)(2,1)} convolution(%fusion.270, %fusion.298), dim_labels=fb_io->bf +qwen3/explicit:3735 ROOT %convolution.78 = bf16[151936,512]{1,0:T(8,128)(2,1)} convolution(%fusion.336, %fusion.285), dim_labels=fb_io->bf + ^^^^^^^^^^^^^^ operands swapped, transpose not folded back +``` + +Same for llama2 (`bf16[512,32000]` vs `bf16[32000,512]`), mistral, mixtral, deepseek +(`bf16[512,129280]` vs `bf16[129280,512]`) and qwen3_moe. + +Scale of the barrier, module-wide: + +| | Sharding custom-calls | non-`Sharding` instructions | adjacent foldable transpose-of-dot pairs | +| ----------------- | --------------------: | --------------------------: | ---------------------------------------: | +| auto (llama2) | 65 | 1,836 | **17** (+3 already behind a custom-call) | +| explicit (llama2) | 1,276 | 1,950 | **0** (all 17 behind a custom-call) | +| auto (qwen3) | 65 | — | — | +| explicit (qwen3) | 1,494 | — | — | + +Every other opcode delta between the two pre-optimization modules is ≤ +33 (`add` +33, +`multiply` +28, `constant` +22, `broadcast` +22). The two programs are the same program plus 1,200 +barriers. + +**Honesty about causality:** this has *not* been proven by rebuilding XLA with the barrier removed. +The evidence is structural and statistical — the modules differ almost only by the custom-calls, +17/17 foldable pairs are adjacent under auto and 0/17 under explicit, and exactly those dots come +out folded under auto and unfolded under explicit. An alternative reading is that XLA's layout +assignment simply made a worse entry-layout choice under explicit (it chose `f32[512,32000]{1,0}` +entry while pinning the internal chain to `{0,1}`; under tp4 it chose `{0,1}` entry and paid +nothing). **Both readings point at the same fix.** + +### 4.2 Consequence A — the lost reduce-scatter + +Under auto, XLA:TPU pattern-matches the FSDP weight-gradient `all-reduce` feeding a +`dynamic-slice(partition-id*128, 0)` and rewrites it into a genuine reduce-scatter, packaged as a +custom fusion: + +``` +llama2/auto: %fusion.12 = bf16[128,32000]{1,0} fusion(%fusion.226), kind=kCustom, calls=%all-reduce-scatter.5 + body: %all-reduce.96 = bf16[512,32000] all-reduce(%input.5), replica_groups={{0,1,2,3}}, + frontend_attributes={from-cross-replica-sharding="true"} + ROOT %dynamic-slice.111 = bf16[128,32000] dynamic-slice(%all-reduce.96, %multiply.222, %constant.718) +``` + +Under explicit the gradient is `[V, emb]`, so the shard dim (512 → 128) is minor-most, the rewrite +cannot match, and a plain (combiner-tupled) all-reduce is emitted: + +``` +llama2/explicit: %all-reduce.118 = (bf16[512], bf16[32000,512]) all-reduce(%copy-done.146, %fusion.242), channel_id=58 +deepseek/explicit: %all-reduce.56 = bf16[129280,512] all-reduce(%fusion.440), channel_id=136 +``` + +`after_codegen` reduce-scatter line counts (2 lines per fusion): + +| model | auto | explicit | +| ------- | ---: | -------: | +| llama2 | 12 | 4 | +| mistral | 12 | 4 | +| mixtral | 12 | 6 | +| qwen3 | 10 | 4 | + +On a 4-way ring a reduce-scatter moves `(N−1)/N · D` per device and an all-reduce moves +`2(N−1)/N · D`. For llama2's 32.77 MB vocab gradient that is **24.58 MB vs 49.15 MB** — the wire +traffic for that one tensor genuinely doubles. + +Three important limits on this half of the mechanism: + +- **It is threshold-dependent, not universal.** At emb ≥ 2048 **both** modes form + `%all-reduce-scatter.7`. qwen3 (V=151936, 77.79 MB/shard) has a bare all-reduce either way — + `%all-reduce.25 bf16[512,151936]` 1700.5 µs/step auto vs `bf16[151936,512]` 1695.0 µs/step + explicit, a 0.3% wash. So **the collective half applies to llama2 / mistral / mixtral / deepseek + at emb 512 only** — 9 of the 25 measured configs (emb 512 × fsdp ≥ 4 × untied vocab); the relayout + half (§4.3) is universal. + + Two successive explanations of *why* it bites there have been tested and discarded, so be careful + with this paragraph's history. It is **not** "the FSDP axis became minor so reduce-scatter is + impossible": the emb-2048 explicit dump contains + `%all-reduce-scatter.7 (input.7: bf16[32000,2048]) → bf16[32000,512]` with + `from-cross-replica-sharding="true"`, so XLA does reduce-scatter along dim 1 of the transposed + shape. Nor is it the **all-reduce combiner** (`xla_tpu_ars_combiner_threshold_in_bytes`), which is + what an earlier revision of this section claimed. The actual gate is in + `tpu-all-reduce-scatter-fusion`: it refuses an AR→DS rewrite when the **per-shard extent of the + minor-most dimension is below ~512 bytes** — one 128-lane × 32-bit vreg row — independent of + dtype. Under `[V, emb]` at emb 512 the per-shard minor extent is 128 elements of bf16 = 256 bytes, + under the bar; at emb 2048 it is 512 × 2 = 1024 bytes, over it. That is the whole emb-512-only + story. + +- **There is a flag, and it is a diagnostic rather than a fix.** + `--xla_tpu_relayout_group_size_threshold_for_reduce_scatter` (default `INT64_MAX`, i.e. "never + insert a relayout in order to enable a reduce-scatter"). Setting it to `1` recovers **every** lost + reduce-scatter in the real MaxText module, which makes it a clean per-config isolation of + mechanism A. Its sign is model-dependent — deepseek3-tiny 9885.2 → 9710.4 µs (−1.77%, 21.9% of the + penalty), llama2 −0.5% — so it is an experiment, never a default. + +- **It is recoverable, and that is not obvious from the paragraph above.** "The gradient is `[V, emb]`, + so the shard dim is minor-most and the rewrite cannot match" describes the *symptom*, and an + earlier revision treated it as terminal for explicit mode. It is not: the orientation is a + consequence of the unfoldable transpose, so emitting the gradient dot in kernel order removes it at + the source. With `lm_head_weight_grad_in_kernel_order: true` on llama2/shrink/fsdp4 the explicit + dump gains a third reduce-scatter fusion, `%all-reduce-scatter.2`, which is byte-for-byte auto's + `%all-reduce-scatter.5` — `bf16[512,32000] → bf16[128,32000]`, layout `{1,0:T(8,128)(2,1)}`, + dim-0 dynamic-slice. `lm_head_kernel_transposed` does **not** recover it; it only fixes the + relayout copies of §4.3. That difference is the whole reason the shipped flag is the `custom_vjp` + one (§8 item 1). + +- **The reduce-scatter is worth less than the bucket in §4.4 suggests.** It halves wire bytes but + achieves only 91.7–99.3 GB/s versus 130.1–138.8 GB/s for the all-reduce, so the win is 0.66–0.76× + time, not the ideal 0.50×. The causal control — `--xla_tpu_enable_all_reduce_scatter_fusion=false` + applied to **auto**, which removes the reduce-scatter while changing nothing else — costs auto only + **+47.2 µs** on llama2, ≤28% of that config's +168 µs penalty. §4.4's xprof bucket puts the same + quantity at +126 µs, so **that bucket over-attributes mechanism A by ~2.7×** on llama2; the + difference is time that overlaps with other work rather than adding to the critical path. + +### 4.3 Consequence B — six relayout copies per step (universal) + +`entry_computation_layout` is **identical** in both modes (llama2: 6× `f32[128,32000]{1,0:T(8,128)}` + +- 6× `f32[32000,128]{1,0:T(8,128)}`) and the optimizer state is donated / input-output-aliased. But + the `{0,1}` layout from the unfolded transpose propagates through the whole Adam update fusion. So + explicit must convert `{1,0}→{0,1}` on the way in and `{0,1}→{1,0}` on the way out, for exactly + three tensors — the `logits_dense` kernel and its Adam `mu` and `nu`: + +``` +%copy.264 = f32[128,32000]{0,1:T(8,128)} copy(%param.100) op_name="state['optimizer']['opt_state'][0]['mu']['decoder']['logits_dense']['kernel'].value" +%copy.265 = f32[128,32000]{0,1:T(8,128)} copy(%param.101) op_name="...['nu']['decoder']['logits_dense']['kernel'].value" +%copy.267 = f32[128,32000]{0,1:T(8,128)} copy(%param.55) op_name="state['model']['decoder']['logits_dense']['kernel'].value" +%copy.268 = f32[128,32000]{1,0:T(8,128)} copy(%multiply_reduce_fusion.4#1) +%copy.269 = f32[128,32000]{1,0:T(8,128)} copy(%multiply_reduce_fusion.4#3) +%copy.270 = f32[128,32000]{1,0:T(8,128)} copy(%multiply_reduce_fusion.4#2) +``` + +Auto has **zero** of these in qwen3 / deepseek / qwen3_moe, and exactly **one** in llama2 / mistral / +mixtral. Layout census on llama2, same logical tensor: + +| | `f32[128,32000]{1,0}` | `f32[128,32000]{0,1}` | +| -------- | --------------------: | --------------------: | +| auto | 51 | **0** | +| explicit | 18 | **37** | + +These are genuine data movement, not metadata. Both layouts are exactly tile-divisible under +`T(8,128)` (for `{1,0}` on `[128,V]`: 128/8 = 16, 32000/128 = 250; for `{0,1}`: 32000/8 = 4000, +128/128 = 1), so there is no padding component — it is a pure transposing gather at HBM speed. Cost +is **linear in `(d_model / n_fsdp_shards) × vocab_size`** and independent of depth, sequence length, +batch size and MoE-ness. Each copy touches the buffer twice (read + write): + +| config | bytes/step | predicted @ 2.37 TB/s | measured | +| --------------------------------------------- | ---------: | --------------------: | -------------: | +| llama2 / mistral / mixtral (emb 512, V 32000) | 196.6 MB | 83 µs | 60.2 µs\* | +| deepseek (emb 512, V 129280) | 794.3 MB | 335 µs | 336.2 µs | +| qwen3 (emb 512, V 151936) | 933.5 MB | 394 µs | 393.7 µs | +| w2048 / d16 / dp4 (emb 2048, V 32000) | 786.4 MB | 332 µs | 345.1–351.5 µs | +| w4096 (emb 4096, V 32000) | 1.573 GB | 664 µs | 659.2 µs | + +\* the 16 MB llama2 buffers are small enough to partially overlap; see the bandwidth caveat in §3. + +**Added entry-copy bytes separate the sign of the effect perfectly across all 25 configs.** This is +the strongest single generalization in the study: the only three configs that add ~zero copy bytes +(gemma +2.1 MB, gemma2 +2.1 MB, tp4 −0.1 MB) are **exactly** the three non-regressions, and all 22 +configs that add ≥ 66 MB regress. No exceptions in either direction. + +The critical observation is that the copies do **not** slow the arithmetic — there is an internal +control in the same module. On dp4, the `token_embedder` Adam fusion (same 65.5 MB tensor, same +3 outputs) costs 177,290 ns vs 177,393 ns (0.06% apart), and the `logits_dense` Adam fusion itself +costs 165,019 ns vs 164,015 ns (0.6% apart). **The `{0,1}` layout costs nothing arithmetically; the +copies are purely additive work.** + +### 4.4 xprof attribution, with the reclassification separated out + +llama2, device-op time summed over 8 cores × 3 steps: + +| class | auto | explicit | Δ | Δ per step | +| ------------- | --------------------: | --------------------: | ---------: | ----------: | +| collective | 2,761,490 ns (13.76%) | 4,358,489 ns (21.20%) | +1,596,999 | +532 µs | +| data-movement | 307,767 ns (1.53%) | 531,344 ns (2.58%) | +223,577 | +74.5 µs | +| fusion | 6,321,712 ns (31.50%) | 5,067,125 ns (24.65%) | −1,254,587 | −418 µs | +| other | 10,678,565 ns | 10,602,167 ns | −76,397 | −25 µs | +| **total** | 20,069,534 ns | 20,559,125 ns | +489,591 | **+163 µs** | + +Measured wall delta is +173 µs/step (4040.3 → 4213.6 µs); the device-op sum predicts +163 µs. + +**But the collective row above is misleading.** Auto's reduce-scatter fusions are classified as +FUSION, which is simultaneously why "collective" rises and "fusion" falls. Re-doing the split on +leaf events and separating the two real mechanisms gives an exact 7-bucket decomposition +(each row sums to the measured window delta to the nanosecond): + +| model | Δ step | vocab RS→AR | 6 relayout copies | all other comm | other data-mvmt | non-comm fusion | matmul/other | idle | [RS→AR]+[copies] | +| ---------- | -------: | ----------: | ----------------: | -------------: | --------------: | --------------: | -----------: | ------: | ---------------: | +| llama2 | +173,242 | +126,270 | +59,776 | −64,005 | +28,102 | +32,478 | +386 | −9,765 | **107.4%** | +| mistral | +213,052 | +127,686 | +73,665 | −39,648 | +25,960 | +9,410 | −1,629 | +17,608 | **94.5%** | +| qwen3 | +334,399 | −12,873 | +394,852 | −67,218 | +31,761 | −26,873 | +6,435 | +8,313 | **114.2%** | +| deepseek | +802,952 | +347,521 | +335,215 | +12,135 | +30,230 | +61,710 | −23,160 | +39,301 | **85.0%** | +| mixtral | +300,282 | +128,674 | +61,129 | — | — | — | — | — | **63.2%** | +| qwen3_moe | +440,442 | — | — | — | — | — | — | — | **85.4%** | +| qwen3_cmoe | +389,645 | — | — | — | — | — | — | — | **102.0%** | + +(ns, per profiled step window on `/device:TPU:0`.) + +Two mechanisms account for **85–114%** of the regression in 6 of the 7 affected shrunk models +(mixtral 63%). + +> **Caveat added after a causal control was run.** The `vocab RS→AR` column above is an +> *attribution*, not a measurement: it sums xprof durations for the collective that changed shape. +> Disabling the rewrite directly on **auto** (`--xla_tpu_enable_all_reduce_scatter_fusion=false`, +> which removes the reduce-scatter and touches nothing else) costs llama2 only **+47.2 µs**, against +> this table's +126,270 ns for the same quantity — the bucket over-states mechanism A by **~2.7×**, +> because much of that collective time overlaps other work instead of extending the critical path. +> The relayout-copy column does not have this problem: §4.3's independent bytes-÷-bandwidth +> prediction lands within a few percent of the measured value on 4 of 5 configs. Read the +> `[RS→AR]+[copies]` percentages as an upper bound, and the copies as the dominant term. + +Of the *apparent* collective-time increase, **75–100% is reclassification** in every +model; the residual real communication delta is +62,265 (llama2), +88,039 (mistral), +359,656 +(deepseek), +100,890 (mixtral) — and **negative** in 6 of 10 models (−80,090 qwen3, −142,763 gemma2, +−125,924 gemma3, −54,346 gemma, −53,401 qwen3_moe, −23,283 qwen3_cmoe). + +Four alternative explanations were tested and **all four refuted**: + +| hypothesis | verdict | evidence | +| ----------------------------------------------- | ---------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| explicit overlaps less / reshards serialize | **refuted** | occupancy (leaf union / window) 96.30→96.68, 97.57→97.27, 98.96→98.91, 97.35→97.16; idle delta is −5.6% to +4.9% of the step delta. The core is not idling more, it is executing more work. | +| tupled all-reduces hide latency worse | **refuted, opposite sign** | llama2 in-loop per iteration: auto = 4-operand AR (2.099 MB) 33,462 ns + five AR-scatter fusions (5.243 MB) 87,798 ns = 121,260 ns for 7.081 MB (58.4 GB/s). explicit = 7-operand AR (5.245 MB) 75,461 ns + two AR-scatter fusions (2.097 MB) 34,840 ns = 110,301 ns for 9.440 MB (85.6 GB/s). **Explicit moves +33% bytes in −9% time.** Same sign for mistral (−69,745 ns) and qwen3 (−66,514 ns). Bigger tuples *help*. | +| the transposed collective is slower on the wire | **refuted** | qwen3's `all-reduce.28` is the identical collective in both modes, differing only by transposition: `bf16[512,151936]` 1,702,101 ns vs `bf16[151936,512]` 1,689,229 ns = **−0.76%**, inside noise. All-gather totals: qwen3 +0.06%, deepseek +0.06%, llama2 −2.1%, mixtral −0.5%. **Transposing a collective is free on the wire**; the cost is the relayout copies it induces around it. | +| MoE / `unreduced` machinery | **refuted — not exercised at all** | see §6.1 | + +### 4.5 Where the barrier comes from — *not* from MaxText's pin + +**This section originally blamed MaxText's `out_sharding` pin. That was measured and is wrong.** +The pin is real, but removing it changes nothing: gating it on `shard_mode` the way +`norm_out_sharding` is gated produces a **0-line diff** in the explicit `after_optimizations` dump, +and all six relayout copies survive. Of the 1,276 `Sharding` custom-calls in the explicit llama2 +module, **0** originate from the `out_sharding=` argument. Keep reading for what does cause them; +the description of the pin below is retained because it is what a reader will find in the code. + +`src/maxtext/layers/decoders.py:805-810` computes an `out_sharding` for the logits +**unconditionally** (unlike `norm_out_sharding` two lines above, which *is* gated on shard_mode): + +```python +if model_mode in (MODEL_MODE_PREFILL, MODEL_MODE_AUTOREGRESSIVE): + out_sharding = create_sharding(self.mesh, (None, None, "activation_vocab")) +else: + out_sharding = create_sharding( + self.mesh, + ("activation_embed_and_logits_batch", "activation_length", "activation_vocab"), + ) +``` + +and passes it into `linears.dense_general` at `decoders.py:843`. +`src/maxtext/layers/linears.py:309-311` is where it becomes explicit-only: + +```python +# out_sharding should be None for auto mesh axis +if self.shard_mode != ShardMode.EXPLICIT: + out_sharding = None +``` + +The real source is one level down, in JAX. On an all-`AxisType.Explicit` mesh, +`lower_with_sharding_in_types` (`jax/_src/interpreters/mlir.py:3170-3190`) annotates the output of +**every** sharding-in-types primitive — ~60 of them — with a `Sharding` custom-call. It early-returns +only for an empty, all-manual or all-auto mesh; there is no "the result sharding already equals the +propagated one, skip it" short-circuit. So the barrier between the gradient `dot_general` and its +`transpose` is emitted whether or not MaxText asks for anything. (Correspondingly, +`_dot_general_lower` at `jax/_src/lax/lax.py:6266-6292` accepts `out_sharding` and never reads it — +only `_dot_general_sharding_rule` uses it, to override the output aval.) + +Two independent controls pin this down: + +- One redundant `jax.lax.with_sharding_constraint` inserted between the dot and the transpose on an + **auto** mesh reproduces the explicit artifact byte for byte. +- `auto` plus `--xla_disable_hlo_passes=algsimp` reproduces explicit's structure exactly. + +So the causal chain is: *any* barrier between the two ops → `algsimp` cannot fold +`transpose(dot(A,B)) → dot(B,A)` → the weight gradient materializes as `[V, emb]`. The fold window is +gone for good by the time the partitioner runs (`algsimp` sits in `simplification-1` at pass 0008-0009; +`xla-partitioner` / `shardy-xla` at 0012-0016), so nothing later can recover it. + +That is why the fix in §8 attacks the **orientation** rather than the barrier: the barrier is +structural to sharding-in-types, but a kernel already stored as `[V, emb]` needs no fold in the +first place. + +This is still an instance of the general pattern recorded as finding [55] in the accompanying code +review — *pinning an intermediate blocks XLA from propagating the constraint backwards into the +producer* — but here the pin is JAX's, not MaxText's, and it manifests through the **algebraic +simplifier and layout assignment** rather than through sharding propagation. + +### 4.6 The natural control: tied embeddings + +gemma and gemma2 are **faster** under explicit (−1.39%, −0.66%). Both set +`logits_via_embedding: true`, so `apply_output_head` takes the `attend_on_embedding` branch +(`decoders.py:813-829`) and **never reaches `logits_dense` at all**. There is no separate vocab +weight-gradient dot, no reduce-scatter to lose, and no `f32[emb/4, V]` optimizer chain to relayout. + +This is the cleanest evidence for the mechanism: **remove the untied vocab projection and the sign +of the effect flips.** Across the 7 non-tied matrix models the mean regression is +4.5%. + +**The all-gather dimension flip that replaces it is free.** All three gemmas show + +``` +auto: %all-gather.30 = bf16[262144,512]{1,0:T(8,128)(2,1)} all-gather(%convert_element_type.862), dimensions={1} +explicit: %all-gather.29 = bf16[512,262144]{0,1:T(8,128)(2,1)} all-gather(%convert_bitcast_fusion), dimensions={0} +``` + +HLO layout is minor-to-major, so `{1,0}` on `[262144,512]` and `{0,1}` on `[512,262144]` both put +512 minor. Same physical `(262144, 512)` grid, same `T(8,128)(2,1)` tiling, and the gather axis is +the physical-minor 512 axis (exactly one 128-lane tile per device) either way. It is a pure logical +relabel, and it measures that way — per-step durations over 3 profiled steps: + +| model | auto | explicit | Δ | +| ------ | --------------------------: | --------------------------: | --------------------: | +| gemma3 | 1225.1 / 1223.3 / 1224.8 µs | 1223.7 / 1225.1 / 1224.9 µs | **+0.2 µs (+0.016%)** | +| gemma | 1206.6 / 1206.1 / 1206.3 µs | 1204.7 / 1207.6 / 1204.9 µs | −0.6 µs (−0.05%) | +| gemma2 | 1205.7 µs | 1204.0 µs | −1.7 µs (−0.14%) | + +**But the flip does cost something one instruction upstream, and this is a universal tied-embedding +tax.** `src/maxtext/layers/embeddings.py:235-240` does +`jnp.asarray(embedding_table, jnp.bfloat16).T` inside `attend_on_embedding`. Under auto XLA folds +the `.T` into the dot's dimension numbers (`dim_labels=bf_oi->bf`), so **one** convert serves both +consumers — `%convert_element_type.862 = bf16[262144,128]` feeds both the input `%gather_fusion` and +the tied-output `%all-gather.30`. Under explicit the transposed value is a distinct typed aval, so +XLA emits `dim_labels=bf_io->bf` and needs a second, differently-shaped value: + +``` +%convert_element_type.396 = bf16[262144,128]{1,0} convert(%param.193) [81.8 µs] -> %gather_fusion +%convert_bitcast_fusion = bf16[128,262144]{0,1} fusion(%param.193) [82.3 µs] -> %all-gather.29 + fused body: convert(f32[262144,128]) -> ROOT bitcast(bf16[128,262144]) +``` + +The transpose really is a bitcast (free); the **duplicated f32→bf16 conversion of the 128 MB shard +is not**: +82.3 µs (gemma3), +80.6 µs (gemma), +81.2 µs (gemma2) — about 0.75% of a step, in +duplicated HBM traffic, paid by every tied-embedding model including the two that win overall. + +What actually carries gemma/gemma2's win is one memory-bound op: the tied-embedding logits matmul, +**−116.6 µs (gemma)** and **−93.4 µs (gemma2)**, i.e. 69% and 78% of their entire net gain. Full +device-op ledger: fusion −314.9 / −433.9 µs against all-reduce + copy penalties of +162.7 / +316.5, +netting −169.1 / −119.7 µs of device-op delta versus measured step deltas of −150.7 / −92.9 µs. + +**Unexplained, and flagged rather than asserted:** that same logits matmul is **+47.1 µs in gemma3** +(auto max 772.7 µs < explicit min 788.4 µs, so it is real and not overlap). The HLO transformation +is byte-for-byte the same shape in all three gemmas — rhs `[V,E]{1,0} bf_oi->bf` becomes +`[E,V]{0,1} bf_io->bf`, same physical operand, same `bf16[1024,V]{1,0}` output, same fused epilogue. +The most likely cause is HBM contention (gemma3/explicit has 132 MB/step of extra layout-copy +traffic competing with a 750 µs memory-bound op), but that cannot be proven from a 3-step profile. +It is 10% of gemma3's delta and 69–78% of gemma/gemma2's gain, so it is not a small caveat. + +### 4.7 gemma3's +4.49% is an artifact of the shrink, and it reverses at real depth + +gemma3 has `logits_via_embedding: true`, identical `bf16[262144,512]` convolutions in both modes, +and **zero vocab relayout copies**. Its communication goes *down* by 125,924 ns. The unembedding +mechanism explains **0% of gemma3**. What it has instead is a structural change nothing else in the +sweep has: **+1243 HLO instructions** (11,209 → 12,452). For comparison, every other model's +instruction delta is between −64 and +71. + +**Root cause: the shrunk model never gets scanned.** +`src/maxtext/models/gemma3.py:39-46` defines `GEMMA3_ATTENTION_PATTERN` with length 6, and +`src/maxtext/layers/decoders.py:1341` computes `scan_length = num_layers // 6`. At +`base_num_decoder_layers=4` that is **0**, the main scan is skipped entirely, and all four layers +fall into the unscanned `layers_remainder` block (`decoders.py:1386-1391`). Every parameter in the +gemma3 dumps is named `state['model']['decoder']['layers_remainder']['layers_N'][...]` — 262 such +references in auto, 298 in explicit, and zero scanned-layer references. **gemma3 is the only model +in the matrix compiled with `while = 0`;** gemma and gemma2 both scan normally. + +Unrolling costs two things that only bite under explicit: + +1. **Remat stops being CSE'd.** `remat_policy` defaults to `'full'` (`base.yml:374`). When the block + is unrolled, the recomputed forward and the original forward live in the *same* HLO computation, + so XLA can CSE them — and under auto it does: only **61 of 11,209** instructions still carry + `rematted_computation` metadata. Under explicit the interposed Sharding custom-calls prevent it: + **1,013 of 12,452**. That is **952 of the +1243 extra instructions (76.6%)**. Directly + observable: `splash_mha_fwd_segmented_residuals` appears 4× in auto (= 4 layers) but **7×** in + explicit (4 forward + 3 recomputed), 91.7 → 158.3 µs; `convolution` goes 34 → 54 with 16 of the + 20 extra tagged `rematted_computation`. Device time of remat-tagged top-level ops: **auto 0.0 µs + (n=0), explicit 212.9 µs (n=75)**. +2. **Every parameter gets its own entry/exit layout copy.** Explicit picks `{0,1}` for 1,202 tensors + (auto: 475; auto prefers `{1,0}` 1,724× vs explicit 1,250×). Module entry/exit buffers are + `{1,0}` and donated, so each mismatch is a real copy — and the counts match the unroll exactly: + 48× `f32[128,1024]` (4 layers × {wi_0, wi_1} × {param, mu, nu} × {in, out}), 36× `f32[128,8,128]` + (q/k/v), 24× `f32[8,128,128]`, 24× `f32[1024,128]` (wo). ~66 MB extra copied per step + (132 MB HBM traffic) = the observed **+132.6 µs**. Inside a `while` body the loop-carried buffers + have one fixed layout, so this tax does not exist for gemma/gemma2 — or for gemma3 at ≥ 6 layers. + +**The decisive experiment.** Same harness, same flags, only `base_num_decoder_layers` changed: + +| layers | structure | auto | explicit | Δ | remat-tagged (auto/expl) | Δ instructions | +| -----: | -------------------------- | ---------: | ---------: | ---------------------: | -----------------------: | -------------: | +| 4 | `scan_length=0`, `while=0` | 10659.8 µs | 11144.3 µs | **+484.4 µs (+4.54%)** | 61 / 1059 | +1243 | +| 6 | `scan_length=1` | 12827.8 µs | 13067.8 µs | +240.1 µs (+1.87%) | 1943 / 1952 | −71 | +| 12 | `scan_length=2`, `while=2` | 18093.2 µs | 17987.6 µs | **−105.6 µs (−0.58%)** | 2519 / 2552 | −18 | + +L4 reproduces the original +4.49% to within 0.05 pp. The remat asymmetry that drives item 1 +disappears at exactly the point the scan appears, and so does the instruction delta. Opcode ledger +across the ladder: + +``` + fusion delta +115.7 (L4) -> -193.3 (L6) -> -596.7 (L12) + custom-call delta +70.4 (L4) -> +2.7 (L6) -> -16.6 (L12) + copy delta +132.6 (L4) -> +146.2 (L6) -> +59.8 (L12) + all-reduce delta +112.4 (L4) -> +218.6 (L6) -> +378.7 (L12) <- universal tax, GROWS with depth +``` + +Attribution of the +486.4 µs at L4 (per-opcode ledger sums to +442.5 µs = 91.0%; residual ~44 µs is +inter-op scheduling gaps): un-CSE'd remat **279.5 µs (57%)**, entry/exit layout copies +**132.6 µs (27%)**, gradient all-reduce repack **112.4 µs (23%)**, duplicated embedding convert +**82.3 µs (17%)**, logits matmul **47.1 µs (10%)**, offset by a **−192 µs** win on weight-gradient +dot relabels. Items 1 and 2 (412 µs, 85% of the delta) exist **only** because `scan_length == 0`. + +Three things the +4.49% is *not*, all checked directly: + +- **Not multimodal.** `use_multimodal` defaults false (`base.yml:1250`). Neither dump contains a + single vision tensor — no 1152 or 4304 hidden dims, no 14×14 patch convolution, zero ops with + image/vision/patch metadata. +- **Not local/global alternation.** With 4 layers, `get_attention_type(0..3)` returns + `GEMMA3_ATTENTION_PATTERN[0..3]` = `LOCAL_SLIDING` for all four — the shrunk model has **zero** + global layers, so the alternation never happens. Both dumps carry the same two `s32[1,512,512]` + mask constants, and the per-layer FSDP weight all-gathers (`all-gather.60-.83`, `bf16[512,8,128]`, + 6 per layer) are 24 in auto vs 25 in explicit and *cheaper* under explicit (245.3 → 223.7 µs). It + is the pattern **length** (6 > 4), not the alternation, that matters. +- **Not the 262144 vocab.** gemma and gemma2 have the same tied-embedding structure, the same 256 MB + all-gather and the same 268 MB gradient all-reduce — and the vocab path is precisely where they + *win*. + +**Practical consequence: do not benchmark gemma3 below 6 layers, and always compare `while`-loop +counts before comparing dumps.** The one durable gemma3 finding is the last ledger row: the +gradient all-reduce repack is a real tax that **grows with depth** (+112 → +219 → +379 µs), even +though the total flips negative because the fusion win grows faster. + +______________________________________________________________________ + +## 5. Results + +### 5.1 Onboarded-model matrix + +Shrunk to fit a v5p-8 while preserving the structural features that drive sharding (MoE, GQA, MLA, +local/global attention): emb 512, mlp 1024, 4 layers, 8 heads, head_dim 128, seq 1024, pdbs 1, +fsdp 4. These are *deliberately small* — see §5.3 before drawing conclusions from the magnitudes. + +| model | tied emb | auto µs | explicit µs | Δ step | mechanism coverage | +| ---------------------- | :------: | ------: | ----------: | ---------: | ---------------------------------: | +| `deepseek3-tiny` | no | 9088 | 9890 | **+8.84%** | 85% | +| `mistral-7b` | no | 3890 | 4103 | **+5.48%** | 95% | +| `gemma3-4b` ⚠ | **yes** | 10831 | 11317 | **+4.49%** | 0% — **shrink artifact, see §4.7** | +| `llama2-7b` | no | 4040 | 4214 | **+4.29%** | 107% | +| `qwen3-8b` | no | 8264 | 8598 | **+4.05%** | 114% | +| `qwen3-30b-a3b` | no | 14484 | 14924 | **+3.04%** | 85% | +| `mixtral-8x7b` | no | 10105 | 10405 | **+2.97%** | 63% | +| `qwen3-custom-30b-a3b` | no | 14488 | 14878 | **+2.69%** | 102% | +| `gemma2-2b` | **yes** | 14160 | 14067 | **−0.66%** | n/a (win) | +| `gemma-2b` | **yes** | 10839 | 10688 | **−1.39%** | n/a (win) | + +"Mechanism coverage" = fraction of the step delta attributed to [lost reduce-scatter] + \[six relayout +copies\] by the per-op decomposition in §4.4. + +⚠ **gemma3's row does not belong in this table.** At 4 layers gemma3 compiles with `while = 0` +(§4.7) — a different code path from the model at any realistic depth. Re-run at 12 layers it is +**−0.58%**. Read the matrix as 9 models plus one methodology finding. + +> **Removed from this table:** the "Δ collective bytes" and "collective count" columns that appeared +> in the first revision of this document. Both are artifacts — see §3 caveats 1 and 2. Explicit's +> *lower* collective count reflects auto's reduce-scatters being counted as all-reduces inside their +> fusions, and the "≈ 0% byte delta" is the byte counter reading pre-slice shapes. + +### 5.2 Parallelism strategy (llama2 / mixtral, same shrunk model, 4 chips cut different ways) + +| config | auto µs | explicit µs | Δ step | note | +| -------------------------- | ------: | ----------: | -----------------------: | ------------------------------------ | +| `ici_data_parallelism=4` | 4442 | 4892 | **+10.12%** | worst case | +| `ici_fsdp_parallelism=4` | 4059 | 4220 | +3.96% | the default | +| `fsdp=2 × tensor=2` | 4335 | 4463 | +2.95% | same six copies, as `f32[256,16000]` | +| `mixtral fsdp=4` | 10112 | 10405 | +2.89% | 72% attributed | +| `mixtral expert=4` | 10732 | 10987 | +2.38% | 98% attributed | +| `fsdp=4, seq 4096` | 12058 | 12134 | +0.63% | | +| `ici_tensor_parallelism=4` | 4563 | 4535 | **−0.61%** (true −0.21%) | only win; different mechanism | + +**Pure DP is the worst case and it is almost pure copy overhead.** With no FSDP the vocab parameter +is *fully replicated*, so the transposing copies operate on `f32[512,32000]` = 65.5 MB each. The six +copies measure **350,463 ns of the +449,471 ns delta = 78.0%**. Widening to the whole optimizer +tail: **+380,640 ns = 84.7%** of the delta occurs *after* the gradient all-reduce finishes, and +92.1% of that tail delta is the six copies. The remaining 15.3% is pre-tail and is *not* relayout: ++35,861 ns is a single **4-byte `s32[]` all-reduce** (`all-reduce.9`) whose duration is pure +inter-core wait, and +27,378 ns is the splash-attention forward Pallas kernel (identical shapes, +identical bytes, most plausibly scheduling skew). + +Note the counter-intuitive detail: explicit's *non-vocab* copies are **cheaper** — per-layer +`copy.283-291` = 73,501 ns vs auto's `copy.199-203` + `copy_dynamic-update-slice_fusion.6-8` = +107,610 ns. The instruction-count delta ("copy 27 → 35") overstates the story; only 6 of the 8 matter +(441.5 MB vs 44.1 MB of copy output bytes — a 10× byte increase behind a 1.3× count increase). + +**Pure TP is the only win, and it has nothing to do with the vocab tensor.** Two corrections to the +obvious reading: + +- **Magnitude.** The −0.61% headline is inflated ~3× by the module-launch lead gap (§3 caveat 3). + Measured on the op span, or on TPU:1's module span, the delta is **−9.2 to −9.3 µs = −0.21%**. + TPU:1 medians: auto 4,475,492 ns vs explicit 4,466,266 ns. Still a genuine win, still the only one, + but not 0.6%. + +- **Mechanism.** It is the RMSNorm **scale** parameter and its AdamW state. Under TP=4 the scale is + a scanned parameter (`param_scan_axis=1`) whose *entry* layout is replicated in both modes + (`f32[512,4]{1,0} parameter(23), sharding={replicated}` — byte-identical). But the activation fed + to the norm is `tensor`-sharded on its last axis, so the VJP w.r.t. `scale` produces a + `tensor`-sharded `[128]` cotangent that must become replicated. Somebody has to all-gather it; the + modes differ only in **where**. + + *Explicit* gathers inside the scan. `src/maxtext/layers/normalizations.py:34-51` + `_align_scale_with_normalized_axis`, gated at `:116-117` on `ShardMode.EXPLICIT`, fires because + `scale_spec[-1] is None != activation_spec[-1] == 'tensor'`: + + ```python + return jax.sharding.reshard(scale, jax.sharding.PartitionSpec(activation_spec[-1])) + ``` + + visible in the pre-opt dump as + `%reshard.37 = bf16[512]{0} custom-call(%add.53), custom_call_target="Sharding", sharding={devices=[4]<=[4]}, metadata={op_name="reshard"}` + (plus `%reshard.25`, `.96`, `.119` in the remat body and `.147` for `decoder_norm`). The gather is + pinned at that program point — inside the backward scan, in bf16. + + *Auto* has **zero** resharded scales (65 `Sharding` custom-calls, none with `op_name="reshard"`). + GSPMD keeps the cotangent `tensor`-sharded and propagates that through the *entire* elementwise + optimizer chain, doing AdamW on `f32[128,4]` shards — then has to honour the mandated replicated + entry/exit layout at the module boundary, so it sinks **nine f32 all-gathers into the epilogue**, + back-to-back, with nothing to overlap them with. And because a value produced by a collective + cannot be written in place into a donated input buffer, copy-insertion adds **5 prologue + 9 + epilogue copies** that explicit does not need. + + Phase decomposition of the median step (split at the scan `while` boundaries): + + ``` + pre-while -5,633 ns (auto's 5 donated-param copies measure 5,540 ns -> 98% attributed) + while +41,647 ns (explicit's 8 in-loop bf16 all-gathers = 34,325 ns; + bitcast_convert_fusion.7/.8 +4,173; dynamic_update_slice.194/.195 10,688) + post-while -45,300 ns (auto's 9 epilogue all-gathers 37,654 ns + 9 epilogue copies 10,268 ns) + ------------------------ + net -9,286 ns + ``` + + Three honest caveats. (i) Explicit's in-scan gathers are **not** well overlapped — 34,325 of the + +41,647 ns while-loop regression is the gathers themselves. Explicit wins by moving a serial cost + from 100%-exposed to ~80%-exposed, which is a weaker story than "explicit enables overlap". + (ii) Explicit does strictly **more** arithmetic — it runs AdamW on replicated `f32[512,4]` instead + of `f32[128,4]`, 4× redundant FLOPs — and only wins because at 2 KB these fusions are + tile-quantized anyway (2,238 ns replicated vs 4,109 ns sharded). **At a real embedding width that + sign could flip.** (iii) "Explicit moves fewer bytes" (−52,224 B of 63.1 MB = 0.08%) is true but + causally irrelevant: at 1–8 KB messages the ICI cost is pure per-message latency, ~4.2 µs either + way. + + So GSPMD's choice here is a mild, real, ~0.2% mistake — defensible in isolation (sharded Adam = + 4× fewer FLOPs, sink-the-collective-to-the-boundary is standard) — not evidence that GSPMD is + broadly worse under TP. + +### 5.3 Scale — the ladder is non-monotone + +llama2 decoder, fsdp 4, varying one dimension at a time: + +| config | auto µs | explicit µs | Δ step | measured copy cost | copy / Δ | +| ------------------------------- | ------: | ----------: | --------------------: | -----------------: | -------: | +| emb 512, 4 layers, seq 1k | 4008 | 4200 | **+4.80%** (+192 µs) | 60.2 µs | 32% | +| emb 2048, 4 layers, seq 1k | 22399 | 22614 | +0.96% (+216 µs) | 351.5 µs | 163% | +| emb 4096, 4 layers, seq 1k | 72210 | 73283 | **+1.49%** (+1073 µs) | 659.2 µs | 61% | +| emb 2048, 4 layers, **seq 4k** | 45492 | 46286 | +1.74% (+794 µs) | 351 µs | 44% | +| emb 2048, 4 layers, **seq 8k** | 86953 | 87229 | +0.32% (+276 µs) | 340 µs | 123% | +| emb 2048, 4 layers, **pdbs 4** | 41233 | 41599 | +0.89% (+366 µs) | 351 µs | 96% | +| emb 2048, **16 layers**, seq 1k | 80398 | 80434 | **+0.04%** (+36 µs) | 345.2 µs | 972% | +| mixtral emb 2048, 4 layers | 96076 | 96620 | +0.57% (+544 µs) | 364 µs | 67% | + +Read that last column carefully. **The copies are always there and always cost roughly what +bandwidth says they should — but how much of that reaches wall clock varies from 32% to 972%.** + +Three conclusions, replacing the "it amortizes" claim in the first revision of this document: + +1. **The width ladder reads non-monotone — because it spans two regimes.** +4.80% → +0.96% → + **+1.49%** along width; +0.96% → +1.74% → +0.32% along sequence. The largest *absolute* penalty + in the entire 28-run sweep (**+1111 µs** against a 2-run auto mean of 72,171 µs) is at the + **largest** width. A fixed-cost model calibrated at w512 (+192 µs) predicts d16 = +0.24% and + w4096 = +0.27%; measured are +0.088% and +1.54%, i.e. **wrong by 2.7× low and 5.7× high**. + + *Correction (2026-09-02):* a single-mechanism model was never going to fit this, and the reason is + structural rather than statistical. **w512 is the only width below XLA's 512-byte minor-most AR→RS + gate** (§4.2: bf16 needs `emb/n_fsdp ≥ 256`, and 512/4 = 128). It therefore pays mechanism A *and* + mechanism B; w2048 and w4096 pay only B. Split that way the ladder is monotone within each regime + — the above-gate rungs rise +0.96% → +1.49% with width, exactly as the relayout-copy model says + they should, and the below-gate rung is a different program. **Read every emb-512 number in this + document as an upper bound that production geometry does not reach.** + +2. **d16's near-zero is cancellation, not absence.** The copies still cost 345.2 ± 0.6 µs (4 runs). + They are offset by an unrelated **−242 µs** win inside the layer scan and a **−390.5 µs** + `concatenate` swing (auto 778.1 µs → explicit 387.6 µs). Had the loop been neutral, d16 would + read +313 µs = **+0.39%**, ~4× the headline. Pooling all 4 runs per mode (rep1-3 + scale/d16, + stdev ~22 µs each) gives **+70.7 µs = +0.088%, t = 4.5** — significant, not noise. + +3. **A second, variable-sign mechanism lives inside the scan.** The out-of-loop budget is fully + explained (copies are 110% of the d16 out-of-loop delta), but the in-loop residual is **−242 µs + at d16, +477 µs at s4k, +409 µs at w4096** — no consistent sign. Its signature is an operand + orientation flip in the attention/MLP backward: auto materializes 2× `bf16[528,8192]` + 1× + `bf16[8192,512]` where explicit materializes 1× + 2× (the 528 = 512 embed shard + 16-row halo, + which drags in `concatenate bf16[576,8192]` and `collective-permute bf16[48,8192]`). + +Regressing step delta on measured copy time over 12 configs: **`Δ = −8.7 + 1.37·copy_µs`, +R² = 0.449**, per-config residual −274 to +467 µs (stdev 240 µs). Mean copy cost 330.5 µs vs mean +delta 443.0 µs — so on *average* the copies explain essentially the whole penalty, but they predict +any individual config only about half the time. + +**What does amortize, in relative terms** — the copy component alone: 1.51% → 1.57% → 0.91% along +width; 1.57% → 0.43% along depth (L4→L16); 1.57% → 0.77% → 0.39% along seq (1k→4k→8k); 1.57% → 0.85% +along batch. So it amortizes over depth, sequence and batch, and only weakly over width. + +Extrapolating to a real llama2-7B (emb 4096, 32 layers, seq 4096, 4-way FSDP): the copy cost stays +around 659 µs on a step that scales to roughly 2.5 s, i.e. **~0.03%**. That extrapolation is one +order of magnitude beyond anything measured here and should be treated as an estimate, not a result. + +______________________________________________________________________ + +## 6. What explicit sharding buys you + +These come from a systematic read of the sharding core, the model layers, the MoE path, and the +test suite (125 candidate findings, 20 upheld under adversarial verification). + +### 6.1 `reduced=` / `unreduced=` PartitionSpec tags — a real capability, unexercised here + +**First, the disclaimer, because it is easy to misread this section:** these tags account for +**0.0% of every measured number in this document.** `src/maxtext/layers/moe.py` contains **zero** +occurrences of `reduced` or `unreduced` (its only two `shard_mode` branches are an `output_sharding` +at `:391-395` and a router-weight scatter `out_sharding` at `:2717-2721`). The tags exist in exactly +two places — `src/maxtext/models/deepseek_batchsplit.py:231-254, 323-378, 436-452` (model +`deepseek3-671b-batchsplit`, never run here) and `src/maxtext/utils/gradient_accumulation.py:198-209` +— and both are gated at `gradient_accumulation.py:77`: + +```python +data_parallel_active = ( + config.shard_mode == ShardMode.EXPLICIT and param_mesh.shape.get("data", 1) > 1 +) +``` + +Every matrix / parallelism / scale run used `gradient_accumulation_steps: 1` (`base.yml:986`) and a +pure fsdp4 or ep4 mesh (`data = 1`), so this was **never True**. There are no fewer and no smaller +gradient all-reduces under explicit anywhere in the data. + +**What the capability actually is.** JAX's `_check_unreduced` (`shard_map.py:457-465`) rejects +`reduced=`/`unreduced=` tags unless *every* mesh axis is `AxisType.Explicit`, which makes them +categorically unavailable under auto. `_unmentioned2` computes +`name_set = _spec_to_vma(spec) | spec.unreduced`, and `to_ct_spec()` swaps reduced ↔ unreduced +(`partition_spec.py:238-242`). A `reduced={'data'}` in-spec removes `data` from `shard_map`'s +defensive `psum`, leaving the cotangent partial — which lets the all-reduce be **hoisted out of a +gradient-accumulation scan**. N per-microbatch all-reduces of `E × M × H` floats collapse to one. + +`deepseek3-671b-batchsplit.yml:61` is `shard_mode: explicit` for exactly this reason: +`deepseek_batchsplit.gather_weights` (`deepseek_batchsplit.py:125-259`) runs +`jax.lax.pcast(w, axis_name='data', to='reduced')` and `jax.lax.all_gather(..., to='reduced')` inside +a `shard_map` whose `out_specs` carry `reduced={'data','fsdp','expert'}`, with the backward +counterpart at `:263-384` carrying `unreduced={...}`. **This config cannot trace under auto.** It is +not a preference; it is a compile-or-die requirement, and it is the fullest expression of what +explicit mode is for. + +Caveat: on this branch the tag is largely dead code outside `deepseek_batchsplit`. Several helpers +still index/iterate `PartitionSpec` directly (`sharding.py:193`, `:242`, `:446-449`), which raises +once a spec carries a tag — `P.__getitem__`/`__iter__` are disabled for tagged specs +(`jax/_src/partition_spec.py:164-176`). Only `_truncate_pspec` (`sharding.py:398-406`) was made +tag-safe. This is why the optimization has been added and reverted twice. + +### 6.2 Larger, merged collectives (the count reduction is *not* the win) + +Explicit tuples gradient all-reduces: 11 → 7 on llama2, 11 → 8 on gemma/qwen3/mistral, 16 → 10 on +gemma2, 25 → 13 on gemma3. + +**Do not read the lower count as a win.** As §3 caveat 1 explains, most of the "missing" all-reduces +are the ones auto had fused into reduce-scatters; the counter is an anti-signal. Similarly deepseek +is the one model where explicit has *more* collectives (51 → 52, all-reduce 11 → 12) and that extra +one is a **benefit**: `%all-reduce-scatter.4` producing `bf16[128,8,128]`, and the pair +auto `all-reduce.143` 27.2 µs → explicit `all-reduce.165` 18.3 µs is **8.9 µs/step faster**. + +Where tupling *is* a genuine win, it is measurable: llama2's in-loop tuple moves **+33.3% wire bytes +in −9.0% time** (85.6 GB/s vs 58.4 GB/s), −43,835 ns over 4 iterations; mistral −69,745 ns; qwen3 +−66,514 ns. Today that win is swamped by §4. + +**But bigger tuples are not universally better, and the counter-example scales with depth.** On +gemma3 the repack is a cost: auto has 25 all-reduces totalling 297,815,400 B (one 33-element tuple +`all-reduce.195` = 8,408,064 B at 105.8 µs, plus 20 standalone 1 MB ARs); explicit has 13 totalling +297,815,516 B (one 45-element tuple `all-reduce.212` = 20,990,976 B at 249.7 µs, plus 8 standalone). +Total bytes differ by 116 B (0.00004%), but the tuple grew **2.50× in bytes and 2.36× in time +(+143.9 µs)** — more traffic pulled into a single blocking collective with nothing left to overlap +it. And unlike everything else in §4.7, this tax **grows with depth**: +112.4 µs at L4, +218.6 at +L6, +378.7 at L12. Caveat: 20 of auto's 25 and 8 of explicit's 13 all-reduces never appear in the +profile's `XLA Ops` line at all; the omission is symmetric, but the ledger only accounts for the +5 timed collectives in each run. + +### 6.3 Errors instead of silent pessimization + +A mismatched ZeRO-1 layout is a hard type error under explicit; under auto GSPMD "papers over it with +an extra collective" (`src/maxtext/configs/types.py:4548-4560`, the repo's own words). Explicit turns +a class of silent performance bugs into compile failures. + +______________________________________________________________________ + +## 7. What explicit sharding costs you + +### 7.1 Capabilities that are hard-blocked + +| capability | status under explicit | evidence | +| ------------------------------------------------ | --------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------- | +| `check_vma=True` | **rejected by MaxText** (`types.py:3089-3090`) although JAX 0.11.1 accepts it on an Explicit mesh | `base.yml:704` calls it "recommended for improved performance" | +| `fused_qkv=True` | **trace-time `ShardingTypeError`** — `attentions.py:1252` is the one projection that forwards no `out_sharding` | no validator catches it; you get a raw JAX error | +| context parallelism on `dot_product` / GPU flash | **trace-time error** — `qk_product` puts the same `context` axis on both `t` and `s` (`attention_op.py:2749`) | TPU splash escapes via `shard_map` manual axes | +| `P.UNCONSTRAINED` | **illegal** when no axis is Auto; the pipeline substitutes `None`, which means *replicate*, the opposite | `pipeline.py` placeholder swap | +| partial adoption | **impossible** — mesh axis types are all-or-nothing (`maxtext_utils.py:2426-2429`) | | + +`check_vma` is the most consequential: it gates the varying-manual-axes fast path in megablox +(`ops.py:296`, `ops.py:665`, `backend.py:528`, `backend.py:789`) on exactly the MoE models that are +on the explicit whitelist. Note the validator is also incomplete — `pipeline.py:385`, `:392`, `:1186` +hardcode `check_vma=True` regardless of config, and nothing forbids +`ici_pipeline_parallelism > 1` with explicit. + +### 7.2 Silent losses + +- **`with_sharding_constraint` becomes a zero-HLO no-op assert** under an all-Explicit mesh. Any + sharding hint not routed through `maybe_shard_with_name` simply vanishes. Confirmed live sites: + `maxtext_utils_nnx.py:175` (init-time), `attention_op.py:2567-2579` and `:2619-2693` + (decode/prefill), `kvcache.py:532/541/602/604/846/847`, `qwen3.py:632`, `diloco_sharding.py:78`. + Training is unaffected (those branches are gated on non-train model modes), but decode and prefill + under explicit are running without their sharding hints. +- **Pinned intermediates can multiply matmul FLOPs.** Reproduced standalone: `x = P('fsdp',None,None) [8,256,512]`, `w` replicated `[512,2048]`, `einsum('bse,ef->bsf')`, result constrained to + `P('fsdp','tensor',None)`. Auto emitted `f32[4,64,512] fusion(... dynamic-slice ...)` — it sliced + the *operand* before the matmul. Explicit emitted `f32[1024,2048] ynn_fusion(...)` then sliced + after: **4× the per-device matmul FLOPs** on that op. This is the same class of defect as §4.1. +- **`mistral`/`qwen3` pin the MLP intermediate with `activation_norm_length`** while + `llama2`/`gemma`/`gemma2`/`gemma3` use `activation_length`. Under flax's first-come mesh-axis + consumption these resolve differently — `activation_norm_length` ends up **replicating** the length + axis. Under context parallelism that is an extra all-gather per `wi` projection per layer (2/layer + for gated MLPs) plus the backward reduce-scatter. At `context=1` they coincide, which is why the + matrix above does not show it. + +### 7.3 No parity guarantee exists + +Both golden suites run **auto only**. `tests/unit/sharding_compare_test.py:124-135` never sets +`shard_mode`; neither does the case table at `tests/utils/sharding_dump.py:39-82` nor the generator +at `tests/utils/run_sharding_dump.py:66-88`. The HLO regression test +`tests/integration/hlo_diff_test.py:105-136` parametrizes `compile_topology` and layer count only. + +**Consequence:** every divergence in this document is invisible to CI. A change that doubled +collectives under explicit but left the loss curve intact would pass the entire test tree. The repo +already has the machinery (`collective_lines()` in `tests/utils/hlo_test_utils.py`) — it is simply +not pointed at explicit mode. + +### 7.4 A dead validation gate + +`validate_shard_mode(shard_mode, decoder_block, quantization)` rejects explicit + quantization and +whitelists only `{simple, simple_mlp, llama2}`. Its only caller is +`pyconfig_deprecated.validate_keys`, reachable only from the legacy `_HyperParameters.__init__`. The +live entry point (`pyconfig.initialize` → `_initialize_pydantic` → `types.MaxTextConfig`) never runs +it, and `types.py` says nothing about quantization + `shard_mode`. So **`shard_mode=explicit` + +`quantization=int8` now passes validation and compiles**, even though the AQT/Qwix paths were never +onboarded. + +______________________________________________________________________ + +## 8. Recommendations + +### For users choosing a mode today + +- **Tied-embedding models (`logits_via_embedding: true`):** explicit is already as fast or faster — + including gemma3, once it is deep enough to scan (−0.58% at 12 layers). Nothing to wait for. +- **gemma3 specifically:** never benchmark it below 6 layers. `scan_length = num_layers // 6`, so a + 4-layer proxy unrolls the whole decoder and measures a different program (§4.7). +- **Untied dense models at production depth/sequence/batch:** the penalty is a few tenths of a + percent. Choose on capability. If you need `check_vma`, `fused_qkv`, or context parallelism on + non-TPU attention, you must use `auto`. +- **Wide models (emb ≥ 4096):** do *not* assume scale hides the cost — it is the configuration with + the largest absolute penalty measured. Profile before committing. +- **MoE with gradient accumulation and `data > 1`:** explicit is the only mode that can express + deferred expert-weight all-reduce. This is the case explicit exists for — though note that path + is unexercised by anything measured here, so its benefit is a code-level claim, not a measurement. +- **Pure data parallelism:** prefer `auto`, or set `lm_head_weight_grad_in_kernel_order: true` if you + are committed to explicit — that closes the whole +9.3% gap (§8 item 1). +10% is not noise. +- **Any untied model on explicit:** `lm_head_weight_grad_in_kernel_order: true` is checkpoint-safe + and inert under `auto`, so the only reason not to set it is that you have measured it neutral for + your geometry and prefer fewer knobs. +- **Pure tensor parallelism:** explicit is mildly better (−0.2%), for the reason in §5.2. The + advantage may not survive at a real embedding width. +- **Small models / debugging runs:** explicit costs a few percent; usually irrelevant, but do not + benchmark a 4-layer proxy and extrapolate. + +### For MaxText, ranked by value + +01. **Emit the untied LM head's weight gradient in kernel order — `lm_head_weight_grad_in_kernel_order` + (§4.1, §4.2, §4.3, §4.5).** + + > Two earlier revisions of this item have been withdrawn. The first recommended "stop pinning the + > logits output under explicit": **measured false**, gating the `out_sharding` produces a 0-line + > HLO diff, because the barriers are JAX's and not MaxText's (§4.5). The second recommended + > `lm_head_kernel_transposed`, storing the kernel as `[vocab, embed]`: that one works, but it is + > dominated on every axis by what follows, so it has been removed rather than shipped beside it. + + Since the barrier cannot be removed, remove the *need* for the fold. Autodiff's default transpose + rule builds the weight gradient as `transpose(dot(g, inputs))` and relies on `algsimp` folding the + transpose back into the dot; under explicit that fold never happens. A hand-written `custom_vjp` + contracts `dk` straight into the kernel's own axis order instead, so there is no transpose to + fold. `_dot_general_kernel_ordered_grad` in `linears.py` implements it, `DenseGeneral`'s + `weight_grad_in_kernel_order` gates it, and `lm_head_weight_grad_in_kernel_order: false` in + `base.yml` wires it to the untied head. It self-disables under `auto` (where the fold does happen) + and under quantization (where it does not own the dot). + + **Why this and not the transposed kernel.** They reach the same layout, but only the `custom_vjp` + recovers the *collective*. On llama2/shrink/fsdp4 the kernel-order dump gains + `%all-reduce-scatter.2`, byte-for-byte identical to auto's `%all-reduce-scatter.5` + (`bf16[512,32000] → bf16[128,32000]`, layout `{1,0:T(8,128)(2,1)}`, dim-0 dynamic-slice); the + transposed kernel does not. It matches the transposed kernel's copy reduction as well + (35 → 29 data-movement ops). And it leaves the stored array untouched, so there is no checkpoint + conversion and no change to initialization. + + **Measured**, v5p, 4 chips, xprof `train_step_ns` median on `/device:TPU:0`, 56 runs, all four + arms back-to-back per config so within-row deltas carry no session drift + (`A` auto, `B` explicit, `C` explicit + transposed kernel, `D` explicit + kernel-order gradient): + + | config | A µs | B µs | C µs | D µs | B−A | C−A | D−A | D−C | rep spread | + | -------- | -------: | -------: | -------: | -------: | ----: | ----: | ----: | ----: | ---------: | + | deepseek | 9081.0 | 9883.5 | 9603.2 | 9216.8 | +8.84 | +5.75 | +1.50 | −4.02 | 0.56% | + | llama2 | 4027.4 | 4223.8 | 4136.2 | 3995.3 | +4.88 | +2.70 | −0.80 | −3.41 | 0.79% | + | mistral | 3905.2 | 4084.6 | 4014.4 | 3853.0 | +4.59 | +2.80 | −1.34 | −4.02 | 0.97% | + | mixtral | 10104.5 | 10407.2 | 10320.4 | 10129.8 | +3.00 | +2.14 | +0.25 | −1.85 | 0.12% | + | qwen3 | 8261.8 | 8608.6 | 8188.9 | 8214.8 | +4.20 | −0.88 | −0.57 | +0.32 | 0.51% | + | *w2048* | 14972.1 | 15215.1 | 14872.8 | 14886.3 | +1.62 | −0.66 | −0.57 | +0.09 | 0.04% | + | *prod* | 196039.4 | 196558.5 | 196145.3 | 196101.9 | +0.26 | +0.05 | +0.03 | −0.02 | 0.03% | + + The first five rows are the shrunk `emb 512` proxies; the two italic rows are `emb 2048` + (`w2048` = 4 layers / seq 1k, `prod` = 16 layers / seq 4k / 16 query heads). + + **Read the two halves of that table differently — this is the part that is easy to get wrong.** + + - The five emb-512 rows go from a **+4.59% median penalty to −0.57%**, D beating auto outright on + three of five. That is real, but it is *the artifact regime*: emb 512 / fsdp 4 puts 256 bytes + per shard on the minor-most dimension, below XLA's 512-byte AR→RS gate (§4.2), so mechanism A + is active only there. No production job has that geometry. **A fix validated only on these rows + is validated on a benchmark artifact.** + - The two above-gate rows are the actual ship criterion, and the bar there is *no regression*, not + a win: **`prod` +0.03%** against a 0.03% rep spread, **`w2048` −0.57%**. There was almost + nothing left to recover, and the flag does not cost anything to have recovered it. That is what + licenses turning it on; the emb-512 column is what makes it worth having the switch at all. + - **D vs C across all seven: median −1.85%, worst +0.32%.** D wins by 1.85–4.02% on the four + configs where C left the most on the table, and ties C within noise on the two where C already + won. There is no configuration measured here where C is the better choice. + + **Numerical equivalence** — 20 training steps against the `auto` baseline, same seed, max absolute + loss deviation (relative in parentheses): + + | config | B (explicit) | C (transposed) | D (kernel-order) | + | ------- | -----------------: | ---------------------: | -----------------: | + | gemma2 | 6.41e-04 (4.9e-05) | — | — | + | llama2 | 2.62e-04 (2.4e-05) | **9.42e-03 (8.7e-04)** | 2.52e-04 (2.3e-05) | + | mixtral | 1.41e-04 (1.3e-05) | **7.93e-03 (7.3e-04)** | 9.63e-05 (8.9e-06) | + + D sits at explicit's own noise level — on mixtral slightly *below* it — because reordering a + contraction only reassociates a floating-point sum. C is ~35× worse for a reason that is not + rounding at all: at the flipped kernel shape `jax.random` draws a different matrix from the same + seed, so C trains a differently-initialized model. That is survivable for a fresh run and fatal + for resuming one. + + **Acceptance tests**, in order of how much they tell you: + + 1. `%all-reduce-scatter` with post-slice shape `bf16[emb/n_fsdp, vocab]` present in the explicit + `after_codegen` dump — the mechanism, and the one thing only this flag delivers. + 2. Zero `copy | f32[*,vocab]` instructions in the explicit `after_optimizations` dump. + 3. Step time at **production geometry** within noise of `auto`. Do not gate on an emb-512 proxy in + either direction. + + Note that raw `reduce-scatter` *counts* are a poor acceptance test on their own — the causal + control in §4.2 values that mechanism at ≤28% of the penalty. + +02. **Do not bother eliding redundant `reshard` barriers — tested, worth exactly zero.** The obvious + companion fix to item 1 is to stop emitting a barrier when the reshard is a no-op (the value's + aval already carries the requested sharding), on the theory that fewer barriers means more folds. + It was implemented and measured, and it does not work: + + | config | barriers off → on | optimized ops | differing ops | + | ------------- | ----------------: | ------------: | ------------- | + | llama2 shrink | 1254 → 1201 | 2889 → 2889 | **none** | + | prod | 1258 → 1205 | 2942 → 2942 | **none** | + | deepseek | 3537 → 3402 | 9642 → 9642 | **none** | + + The predicate is correct — verified on CPU, where it collapses N redundant reshards to the + irreducible 2 — but it only reaches ~4% of the barriers in a real module, and the optimized + executable is **op-identical** in all three cases. XLA already discards no-op `Sharding` + custom-calls; the ones that matter are the ones on real reshards, which by definition cannot be + elided. Recorded here so the next person does not spend the week: **the barrier count is not the + lever, the gradient orientation is.** + +03. **Add `shard_mode` to the golden suites (§7.3).** Until CI compares modes, every regression here + can silently return. `tests/integration/hlo_diff_test.py` is the cheapest place to start; assert + on `after_codegen` reduce-scatter counts, not on static collective byte totals. + +04. **Any HLO byte-counter used to compare modes must descend into + `kind=kCustom, calls=%all-reduce-scatter.*` fusions** and report the post-slice shape. Counting + the `all-reduce` inside the fusion at its pre-slice shape reports auto's wins as if they were + losses, which is how the "identical bytes" claim survived a full sweep here. + +05. **Re-examine the `check_vma` ban (§7.1).** JAX 0.11.1 accepts `check_vma=True` on an all-Explicit + mesh; the restriction appears to be MaxText-imposed. Lifting it would give the MoE models on the + explicit whitelist back a documented optimization. + +06. **Validate `fused_qkv` + explicit and context-parallel + explicit** in `types.py` so users get a + MaxText message instead of a raw JAX trace error. + +07. **Fix `gcs_utils.upload_dump()`** to no-op (or write locally) for non-`gs://` paths — today it + crashes after training and destroys the profile. + +08. **Make explicit mode not defeat remat CSE on unrolled blocks (§4.7).** With `remat_policy: full` + and no scan, auto CSEs the recomputed forward against the original (61 remat-tagged instructions + of 11,209); explicit cannot (1,013 of 12,452), costing 279.5 µs. This bites any config where a + block is unrolled — gemma3 below 6 layers today, but also `scan_layers: false` runs generally. + +09. **Deduplicate the tied-embedding bf16 cast (§4.6).** `embeddings.py:235-240`'s `.T` makes explicit + materialize the 128 MB embedding shard in bf16 **twice** (+80–82 µs on every gemma). Hoisting the + cast above the transpose, or reusing one converted value for both consumers, is a pure win and is + independent of everything else in this document. + +10. **Investigate the tuple-repack all-reduce tax (§6.2)** — the one measured explicit cost that grows + with depth (+112 → +219 → +379 µs from 4 to 12 layers on gemma3). + +11. **Delete or rewire `validate_shard_mode`** (§7.4) — as dead code it actively misinforms. + +______________________________________________________________________ + +## 9. Reproducibility + +The measurement harness was a set of throwaway scripts and is not checked in; it is short enough to +restate. Every run was: + +```bash +python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \ + model_name= shard_mode= \ + steps=12 per_device_batch_size=1 max_target_length=1024 \ + dataset_type=synthetic enable_checkpointing=False \ + skip_first_n_steps_for_profiler=3 profiler_steps=3 profiler=xplane profile_cleanly=True \ + jax_cache_dir= \ + override_model_config=True base_emb_dim=512 base_mlp_dim=1024 base_num_decoder_layers=4 \ + base_num_query_heads=8 base_num_kv_heads=8 head_dim=128 +``` + +with `XLA_FLAGS="--xla_dump_to= --xla_dump_hlo_module_re=.*train_step.*"` for the HLO. Two notes +that cost time to discover: + +- Drive the dump through `XLA_FLAGS`, **not** MaxText's `dump_hlo` config — that hook unconditionally + uploads to GCS and raises on a local `base_output_directory`, aborting the run before the xplane + profile is flushed (§8 item 6). +- Clear `jax_cache_dir`. A persistent-cache hit skips compilation and therefore suppresses the HLO + dump, while producing an identical executable and step time. + +Timings are `train_step` durations read off `/device:TPU:0` in the `.xplane.pb`, median of the 3 +profiled steps; HLO facts are from `*after_optimizations.txt` and `*after_codegen*`. The TPU is a +serial resource, so runs were serialized with `flock`. 12 steps, profile steps 4-6, synthetic data, +checkpointing off, throughout. + +### Limitations, stated plainly + +- **The barrier→simplifier causal link is inferred, not proven.** No XLA rebuild with the + `Sharding` custom-call suppressed was performed. See §4.1 for the alternative layout-assignment + reading; both imply the same MaxText-side fix, but the XLA-side attribution should be confirmed + before filing an XLA bug. +- **4 chips, single host.** No cross-host ICI/DCN behaviour is exercised. The collective costs here + are intra-slice. Conclusions about *relayout* cost should hold; conclusions about *collective + latency* may not scale to multi-slice. In particular the reduce-scatter loss (§4.2) is + combiner-threshold-dependent and its 2× wire-byte penalty is a 4-way-ring figure. +- **Shrunk models, and one of them shrank into a different program.** Real parameter counts do not + fit a v5p-8. §5.1's magnitudes are proxies, not production numbers; §5.3 exists to bound the + extrapolation and its answer is "partly, and non-monotonically". gemma3 (§4.7) is the cautionary + case: its shrink disabled the scan, and its headline number reverses sign at real depth. The other + nine models compiled to the same `while` structure in both modes, which is the check that caught it. +- **The unembedding mechanism is the well-evidenced one.** The in-scan variable-sign residual + (§5.3 item 3), mixtral's 28% unattributed remainder, gemma3's logits-matmul divergence (§4.6, + ±100 µs with no mechanistic explanation), and the context-parallelism MLP spec divergence (§7.2) + are identified but not fully attributed. +- **The gemma3 layer ladder is 1 run per point**, not 3 — enough to establish the sign reversal + (+484 µs → −106 µs dwarfs the noise floor) but not to quantify L6/L12 precisely. +- **3 profiled steps per run** for most configs (4 runs per mode only for the d16 point). Adequate + given the 0.07% noise floor, but not a large sample — and the module-span lead gap (§3 caveat 3) + puts a floor of a few tenths of a percent on any single-pair comparison. diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 269339e056..06ee18f976 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -200,6 +200,14 @@ logits_via_embedding: false normalize_embedding_logits: true # whether to normalize pre-softmax logits if logits_via_embedding is true logits_dot_in_fp32: false # whether to use fp32 in logits_dense or shared_embedding dot product for stability cast_logits_to_fp32: true # whether to cast the logits to fp32. the higher precision is generally beneficial, but it can vary slightly. +# Compute the untied LM head's weight gradient with a hand-written backward pass that contracts straight into the +# kernel's stored orientation, instead of letting autodiff emit transpose(dot(...)) and hoping XLA folds it away. +# Under shard_mode: explicit that fold never happens, because the Sharding custom-call on every dot output is a +# barrier algsimp will not cross, so the gradient is left sharded on its minor-most dimension and loses the +# all-reduce -> reduce-scatter fusion (see docs/guides/optimization/shard_mode_performance.md). The stored kernel +# and its initialization are untouched, so no checkpoint conversion is needed; gradients differ from the default +# rule only by floating-point reassociation. No effect under shard_mode: auto, where XLA folds the transpose itself. +lm_head_weight_grad_in_kernel_order: false float32_qk_product: false # in dot_product attention, whether to cast to fp32 the inputs to qk product float32_logits: false # in dot_product attention, whether to cast to fp32 the inputs to softmax mla_qk_head_chunk_size: 0 # Limits HBM footprint by sequentially evaluating the QK matrix in the Indexer across the unsharded local heads dimension natively. diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 394e57f514..e220bec38a 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -604,6 +604,15 @@ class LogitsAndLoss(BaseModel): ) logits_dot_in_fp32: bool = Field(False, description="Use fp32 for the logits dot product for stability.") cast_logits_to_fp32: bool = Field(True, description="Whether to cast the final logits to fp32.") + lm_head_weight_grad_in_kernel_order: bool = Field( + False, + description=( + "Compute the untied LM head's weight gradient directly in the kernel's stored axis order via a " + "custom_vjp, so autodiff emits no transpose after the gradient dot. Recovers the weight-gradient " + "reduce-scatter that shard_mode=explicit otherwise loses, leaving the stored kernel and its " + "initialization untouched. No effect under shard_mode=auto." + ), + ) final_logits_soft_cap: None | NonNegativeFloat = Field( None, description="Soft-cap value for the final logits. None or 0.0 means no cap.", @@ -3269,6 +3278,17 @@ def validate_shard_embed_moe_on_fsdp(self) -> "MaxTextConfig": ) return self + @model_validator(mode="after") + def validate_lm_head_weight_grad_in_kernel_order(self) -> "MaxTextConfig": + """Reject lm_head_weight_grad_in_kernel_order where it cannot be honored.""" + if not self.lm_head_weight_grad_in_kernel_order: + return self + if self.logits_via_embedding: + raise ValueError( + "lm_head_weight_grad_in_kernel_order only applies to the untied LM head, but logits_via_embedding is True." + ) + return self + @model_validator(mode="after") def set_derived_and_validate_values(self) -> "MaxTextConfig": """ diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index 18a9906c15..d09becc49b 100644 --- a/src/maxtext/layers/decoders.py +++ b/src/maxtext/layers/decoders.py @@ -838,6 +838,7 @@ def apply_output_head(self, shared_embedding: nn.Module | nnx.Module, y, determi name="logits_dense", matmul_precision=self.config.matmul_precision, parameter_memory_host_offload=cfg.parameter_memory_host_offload, + weight_grad_in_kernel_order=cfg.lm_head_weight_grad_in_kernel_order, )( y, out_sharding=out_sharding, diff --git a/src/maxtext/layers/linears.py b/src/maxtext/layers/linears.py index e088f7ec50..92beaeeee6 100644 --- a/src/maxtext/layers/linears.py +++ b/src/maxtext/layers/linears.py @@ -85,6 +85,104 @@ def _compute_dot_general(inputs, kernel, kernel_axes, axis, contract_ind, matmul return dot_general(inputs, kernel, ((axis, contract_ind), ((), ())), precision=matmul_precision) +def _aval_sharding(value): + """The sharding carried by `value`'s type, or None when it carries none. + + Under `shard_mode: explicit` every aval knows its own sharding, so this is how + a hand-written backward pass pins its outputs to the same layout the default + transpose rule would have produced. + """ + try: + sharding = jax.typeof(value).sharding + except Exception: # pylint: disable=broad-except + return None + spec = getattr(sharding, "spec", None) + if spec is None or all(p is None for p in spec): + return None + return sharding + + +def _restore_axis_order(value, produced_axes): + """Permute `value` so that its axis j is the operand's original axis j. + + `produced_axes[i]` names the original axis that dot_general placed at position + i. For every layout this path is enabled for the mapping is already the + identity, so the transpose is elided rather than merely folded later. + """ + produced_axes = tuple(produced_axes) + if produced_axes == tuple(range(len(produced_axes))): + return value + return lax.transpose(value, tuple(produced_axes.index(j) for j in range(len(produced_axes)))) + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(2, 3, 4)) +def _dot_general_kernel_ordered_grad(inputs, kernel, dimension_numbers, precision, out_sharding): + """`lax.dot_general` whose weight gradient is emitted in stored-kernel layout. + + JAX's default transpose rule builds the kernel gradient as + `transpose(dot(g, inputs))`. XLA normally folds that transpose back into the + dot, but under `shard_mode: explicit` every dot output carries a `Sharding` + custom-call and algebraic simplification will not fold across it. The + transpose then survives to codegen, which leaves the gradient sharded on its + minor-most dimension -- the one position where XLA declines to fuse the + weight-gradient all-reduce into a reduce-scatter. See + docs/guides/optimization/shard_mode_performance.md for the measurements. + + Writing the backward pass by hand avoids the transpose instead of relying on + XLA to undo it: `dk` is contracted straight into the kernel's own axis order. + That reassociates the sum, so the gradients match the default rule to float + rounding (~1e-7 relative) rather than bit for bit. + """ + return lax.dot_general(inputs, kernel, dimension_numbers, precision=precision, out_sharding=out_sharding) + + +def _dot_general_kernel_ordered_grad_fwd(inputs, kernel, dimension_numbers, precision, out_sharding): + out = lax.dot_general(inputs, kernel, dimension_numbers, precision=precision, out_sharding=out_sharding) + return out, (inputs, kernel) + + +def _dot_general_kernel_ordered_grad_bwd(dimension_numbers, precision, out_sharding, res, g): + """Backward pass for `_dot_general_kernel_ordered_grad`.""" + del out_sharding # Constrains the forward output only. + inputs, kernel = res + (lhs_contract, rhs_contract), (lhs_batch, rhs_batch) = dimension_numbers + if lhs_batch or rhs_batch: + raise ValueError("_dot_general_kernel_ordered_grad does not support batch dimensions.") + lhs_free = tuple(i for i in range(inputs.ndim) if i not in lhs_contract) + rhs_free = tuple(i for i in range(kernel.ndim) if i not in rhs_contract) + n_lhs_free = len(lhs_free) + + # dk contracts the free input axes against g's leading axes, which correspond + # to them one for one. The result carries the kernel's contracting axes (in + # ascending operand order) followed by its free axes -- the kernel's own order + # whenever it is stored in_features-major. + dk = lax.dot_general( + inputs, + g, + ((lhs_free, tuple(range(n_lhs_free))), ((), ())), + precision=precision, + out_sharding=_aval_sharding(kernel), + ) + dk_axes = tuple(rhs_contract[lhs_contract.index(a)] for a in sorted(lhs_contract)) + rhs_free + dk = _restore_axis_order(dk, dk_axes) + + # dx contracts g's trailing axes against the kernel's free axes, likewise + # paired one for one. + dx = lax.dot_general( + g, + kernel, + ((tuple(range(n_lhs_free, g.ndim)), rhs_free), ((), ())), + precision=precision, + out_sharding=_aval_sharding(inputs), + ) + dx_axes = lhs_free + tuple(lhs_contract[rhs_contract.index(a)] for a in sorted(rhs_contract)) + dx = _restore_axis_order(dx, dx_axes) + return dx, dk + + +_dot_general_kernel_ordered_grad.defvjp(_dot_general_kernel_ordered_grad_fwd, _dot_general_kernel_ordered_grad_bwd) + + def _compute_dot_general_nnx( inputs, kernel, @@ -94,6 +192,7 @@ def _compute_dot_general_nnx( quant_dot_general: nnx_wrappers.ToNNX | None, initializing: bool, out_sharding: NamedSharding | None = None, + weight_grad_in_kernel_order: bool = False, ): """Computes a dot_general operation that may be quantized.""" dot_general = lax.dot_general @@ -107,6 +206,11 @@ def _compute_dot_general_nnx( out_ndim = (inputs.ndim - len(axis)) + (kernel.ndim - len(contract_ind)) out_sharding = truncate_out_sharding(out_sharding, out_ndim) + if weight_grad_in_kernel_order: + return _dot_general_kernel_ordered_grad( + inputs, kernel, ((axis, contract_ind), ((), ())), matmul_precision, out_sharding + ) + return dot_general( inputs, kernel, ((axis, contract_ind), ((), ())), precision=matmul_precision, out_sharding=out_sharding ) @@ -132,6 +236,7 @@ def __init__( mesh: Mesh | None = None, use_two_stage_all_gather: bool = False, debug_sharding: bool = False, + weight_grad_in_kernel_order: bool = False, *, # Following arguments are keyword-only rngs: nnx.Rngs = None, ): @@ -158,6 +263,16 @@ def __init__( transpose XLA emits for a single combined 2-axis all-gather. debug_sharding: when True, log the logical/physical sharding of the two-stage all-gather constraints to the sharding dump files. + weight_grad_in_kernel_order: compute the weight gradient with a hand-written + backward pass that contracts straight into the kernel's stored axis order, + rather than letting autodiff emit `transpose(dot(...))` and relying on XLA + to fold the transpose away. See + docs/guides/optimization/shard_mode_performance.md section 4.3 for why + that fold does not happen under `shard_mode: explicit`. The stored kernel + and its initialization are untouched, so no checkpoint conversion is + needed; the gradients differ from the default rule only by floating-point + reassociation. Only takes effect under `shard_mode: explicit`; ignored + when the layer is quantized. rngs: RNG state for initialization in nnx. """ self.in_features_shape = canonicalize_tuple(in_features_shape) @@ -175,6 +290,7 @@ def __init__( self.mesh = mesh self.use_two_stage_all_gather = use_two_stage_all_gather self.debug_sharding = debug_sharding + self.weight_grad_in_kernel_order = weight_grad_in_kernel_order # Parameter initialization kernel_shape = self.in_features_shape + self.out_features_shape @@ -311,15 +427,26 @@ def __call__( out_sharding = None contract_ind = tuple(range(0, len(self.axis))) + quant_dot_general = self.quant_dot_general if slice_bounds is None else None + # The hand-written weight gradient only pays off where XLA refuses to fold + # the transpose away, and it has to own the whole dot to do it -- so it is + # off under auto sharding and under any quantized path. + weight_grad_in_kernel_order = ( + self.weight_grad_in_kernel_order + and self.shard_mode == ShardMode.EXPLICIT + and self.quant is None + and quant_dot_general is None + ) output = _compute_dot_general_nnx( inputs, kernel, norm_axis, contract_ind, self.matmul_precision, - self.quant_dot_general if slice_bounds is None else None, + quant_dot_general, _initializing, out_sharding, + weight_grad_in_kernel_order, ) if self.bias is not None: @@ -346,6 +473,7 @@ def dense_general( shard_mode: ShardMode = ShardMode.AUTO, matmul_precision: str = "default", parameter_memory_host_offload: bool = False, + weight_grad_in_kernel_order: bool = False, name: None | str = None, ): """Creates a DenseGeneral Linen module using nnx.bridge.to_linen. @@ -365,6 +493,8 @@ def dense_general( shard_mode: indicating the shard mode matmul_precision: Precision for matrix multiplication. parameter_memory_host_offload: Determines whether to offload params to host + weight_grad_in_kernel_order: compute the weight gradient directly in the + kernel's stored axis order; see DenseGeneral for details. name: name passed to the ToLinen Module """ if not (inputs_shape is not None) ^ (in_features_shape is not None): @@ -389,6 +519,7 @@ def dense_general( shard_mode=shard_mode, matmul_precision=matmul_precision, parameter_memory_host_offload=parameter_memory_host_offload, + weight_grad_in_kernel_order=weight_grad_in_kernel_order, name=name, metadata_fn=variable_to_logically_partitioned, abstract_init=False, diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index e0347b8f1e..91206d4ab9 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -471,6 +471,7 @@ def __init__( shard_mode=config.shard_mode, matmul_precision=self.config.matmul_precision, parameter_memory_host_offload=config.parameter_memory_host_offload, + weight_grad_in_kernel_order=config.lm_head_weight_grad_in_kernel_order, rngs=rngs, ) diff --git a/tests/unit/linears_test.py b/tests/unit/linears_test.py index a20d29a0d9..976458a40d 100644 --- a/tests/unit/linears_test.py +++ b/tests/unit/linears_test.py @@ -14,9 +14,10 @@ """Tests for linears.py.""" +import functools import sys import unittest -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch from flax import nnx from flax.linen import partitioning as nn_partitioning import jax @@ -24,6 +25,7 @@ import numpy as np import pytest +from maxtext.common.common_types import ShardMode from maxtext.layers import linears from maxtext.configs import pyconfig from maxtext.utils import maxtext_utils @@ -180,6 +182,149 @@ def test_slice_bounds_invalid(self): with self.assertRaisesRegex(ValueError, "slice_bounds .* must be valid and within"): layer(inputs, slice_bounds=(0, 100)) + def _make_kernel_ordered_pair(self, in_features_shape, out_features_shape, axis=-1, use_bias=False): + """Builds a default layer and a weight_grad_in_kernel_order twin sharing weights. + + Both store the kernel identically and share the RNG stream, which is exactly + the property that makes the flag checkpoint-compatible. + """ + kwargs = { + "in_features_shape": in_features_shape, + "out_features_shape": out_features_shape, + "axis": axis, + "use_bias": use_bias, + "shard_mode": ShardMode.EXPLICIT, + } + default = linears.DenseGeneral(**kwargs, rngs=nnx.Rngs(params=0)) + ordered = linears.DenseGeneral(**kwargs, weight_grad_in_kernel_order=True, rngs=nnx.Rngs(params=0)) + return default, ordered + + def test_weight_grad_in_kernel_order_leaves_kernel_untouched(self): + default, ordered = self._make_kernel_ordered_pair(4, 8) + # No checkpoint impact is the point of this flag, so the stored array must be + # bit-identical, not merely the same shape. + self.assertEqual(default.kernel[...].shape, ordered.kernel[...].shape) + self.assertEqual(default.kernel_axes, ordered.kernel_axes) + np.testing.assert_array_equal(default.kernel[...], ordered.kernel[...]) + + def test_weight_grad_in_kernel_order_forward_matches(self): + cases = ((4, 8, -1), (4, (2, 8), -1), ((2, 3), 8, (1, 2))) + for in_features_shape, out_features_shape, axis in cases: + with self.subTest(in_features_shape=in_features_shape, out_features_shape=out_features_shape): + default, ordered = self._make_kernel_ordered_pair(in_features_shape, out_features_shape, axis=axis) + shape = (2, 3) if axis == -1 else (5,) + shape += (in_features_shape,) if isinstance(in_features_shape, int) else tuple(in_features_shape) + inputs = jax.random.normal(jax.random.PRNGKey(0), shape) + np.testing.assert_array_equal(default(inputs), ordered(inputs)) + + def test_weight_grad_in_kernel_order_gradients_match(self): + """Reordering the contraction must not change the gradient beyond float noise.""" + cases = ((4, 8, -1), (4, (2, 8), -1), ((2, 3), 8, (1, 2))) + for in_features_shape, out_features_shape, axis in cases: + with self.subTest(in_features_shape=in_features_shape, out_features_shape=out_features_shape): + default, ordered = self._make_kernel_ordered_pair(in_features_shape, out_features_shape, axis=axis) + shape = (2, 3) if axis == -1 else (5,) + shape += (in_features_shape,) if isinstance(in_features_shape, int) else tuple(in_features_shape) + inputs = jax.random.normal(jax.random.PRNGKey(0), shape) + + def loss(layer, x): + return jnp.sum(jnp.sin(layer(x))) + + grad_default = nnx.grad(loss)(default, inputs).kernel[...] + grad_ordered = nnx.grad(loss)(ordered, inputs).kernel[...] + self.assertEqual(grad_ordered.shape, ordered.kernel[...].shape) + np.testing.assert_allclose(grad_default, grad_ordered, rtol=1e-5, atol=1e-6) + + grad_x_default = jax.grad(functools.partial(loss, default))(inputs) + grad_x_ordered = jax.grad(functools.partial(loss, ordered))(inputs) + np.testing.assert_allclose(grad_x_default, grad_x_ordered, rtol=1e-5, atol=1e-6) + + def test_weight_grad_in_kernel_order_lm_head_shape(self): + """The LM head's 2-D kernel is the only shape the config flag enables. + + Reordering the contraction reassociates the sum, so the gradient agrees to + float rounding rather than exactly. The stored kernel -- and therefore the + initialization -- is untouched, so the difference stays at rounding level + instead of changing the training trajectory. + """ + for dtype in (jnp.float32, jnp.bfloat16): + with self.subTest(dtype=jnp.dtype(dtype).name): + default, ordered = self._make_kernel_ordered_pair(8, 32) + inputs = jax.random.normal(jax.random.PRNGKey(0), (2, 5, 8), dtype) + weights = jax.random.normal(jax.random.PRNGKey(1), (2, 5, 32), dtype) + + def loss(layer, x, weights=weights): + return jnp.sum((layer(x) * weights).astype(jnp.float32)) + + grad_default = nnx.grad(loss)(default, inputs).kernel[...] + grad_ordered = nnx.grad(loss)(ordered, inputs).kernel[...] + scale = float(jnp.abs(grad_default).max()) + self.assertLess(float(jnp.abs(grad_default - grad_ordered).max()), 1e-5 * scale) + + def test_weight_grad_in_kernel_order_composes_with_remat(self): + # The LM head sits outside the scanned stack, but training still wraps it in + # remat, and a custom_vjp that breaks under checkpointing is unusable. + default, ordered = self._make_kernel_ordered_pair(4, 8) + inputs = jax.random.normal(jax.random.PRNGKey(0), (2, 3, 4)) + + def grads(layer): + return nnx.grad(nnx.remat(lambda l, x: jnp.sum(jnp.sin(l(x)))))(layer, inputs).kernel[...] + + np.testing.assert_allclose(grads(default), grads(ordered), rtol=1e-5, atol=1e-6) + + def test_weight_grad_in_kernel_order_ignored_under_auto(self): + # Under auto sharding XLA folds the transpose itself, so the custom_vjp buys + # nothing and must not take over the dot. + layer = linears.DenseGeneral( + in_features_shape=4, + out_features_shape=8, + shard_mode=ShardMode.AUTO, + weight_grad_in_kernel_order=True, + rngs=nnx.Rngs(params=0), + ) + inputs = jax.random.normal(jax.random.PRNGKey(0), (2, 3, 4)) + reference = linears.DenseGeneral( + in_features_shape=4, out_features_shape=8, shard_mode=ShardMode.AUTO, rngs=nnx.Rngs(params=0) + ) + np.testing.assert_array_equal(layer(inputs), reference(inputs)) + + def test_weight_grad_in_kernel_order_yields_to_quantization(self): + # The custom_vjp has to own the whole dot, so it must step aside for a + # quantized one rather than wrap it. + mock_quant = MagicMock() + mock_quant.quant_mode = None + mock_instance = mock_quant.dot_general_cls.return_value.return_value + mock_instance.init_with_output.return_value = (MagicMock(), {}) + mock_instance.apply.return_value = (MagicMock(), {}) + + layer = linears.DenseGeneral( + in_features_shape=4, + out_features_shape=8, + shard_mode=ShardMode.EXPLICIT, + quant=mock_quant, + weight_grad_in_kernel_order=True, + rngs=self.rngs, + ) + inputs = jax.random.normal(jax.random.PRNGKey(0), (2, 3, 4)) + with patch.object(linears, "_dot_general_kernel_ordered_grad") as spy: + layer(inputs) + spy.assert_not_called() + + def test_weight_grad_in_kernel_order_emits_no_transpose(self): + """The mechanism itself: no transpose on the weight-gradient path.""" + default, ordered = self._make_kernel_ordered_pair(4, 8) + inputs = jax.random.normal(jax.random.PRNGKey(0), (2, 3, 4)) + + def count_transposes(layer): + def loss(l, x): + return jnp.sum(jnp.sin(l(x))) + + text = nnx.jit(nnx.grad(loss)).lower(layer, inputs).as_text() + return text.count("stablehlo.transpose") + + self.assertEqual(count_transposes(ordered), 0) + self.assertGreater(count_transposes(default), 0) + def _run_dense_test(self, axis, in_feat_shape, expected_shape): batch_size = 2 seq_len = 3 From acc0a815176520170caa435bd587049a2bf5f57e Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Thu, 3 Sep 2026 11:55:00 +0000 Subject: [PATCH 2/9] Contract over the embedding table's feature axis in attend_on_embedding `attend_on_embedding` materialized `table.T` to compute the tied output head's logits. Under `shard_mode: explicit` the transposed table is a distinct typed value, so XLA cannot fold it into the dot's dimension numbers: the table shard is cast to bf16 once for the input lookup in `Embed.__call__` and a second time here, and the tied head's weight gradient comes out flipped and needs a transpose to put back. Expressing the transpose as dimension numbers via einsum keeps both consumers on one cast and leaves the gradient in the table's own axis order. Under `shard_mode: auto` XLA already folded the `.T` away, so this is a no-op there. It reassociates the accumulation, so logits agree to float rounding rather than bit for bit. --- src/maxtext/layers/embeddings.py | 13 +++++++++-- tests/unit/embeddings_test.py | 37 ++++++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 2 deletions(-) diff --git a/src/maxtext/layers/embeddings.py b/src/maxtext/layers/embeddings.py index a744ea041a..b0f8d23f6b 100644 --- a/src/maxtext/layers/embeddings.py +++ b/src/maxtext/layers/embeddings.py @@ -232,9 +232,18 @@ def attend_on_embedding( if out_sharding is not None: out_sharding = truncate_out_sharding(out_sharding, query.ndim) embedding_table = _maybe_move_embedding_to_device(embedding_table, config) - return jnp.dot( + # Contract over the table's feature axis instead of materializing `table.T`. + # Under `shard_mode: explicit` the transposed table is a distinct typed value, + # so XLA cannot fold it into the dot's dimension numbers: the (large) table + # shard is cast to bf16 once for the input lookup in `Embed.__call__` and a + # second time here, and the tied head's weight gradient comes out flipped. + # Expressing the transpose as dimension numbers keeps both consumers on one + # cast and leaves the gradient in the table's own axis order. Under `auto` + # this is a no-op -- XLA already folded the `.T` away. + return jnp.einsum( + "...e,ve->...v", query, - jnp.asarray(embedding_table, jnp.bfloat16).T, + jnp.asarray(embedding_table, jnp.bfloat16), preferred_element_type=attend_dtype, out_sharding=out_sharding, ) diff --git a/tests/unit/embeddings_test.py b/tests/unit/embeddings_test.py index 88ef09d09c..ac54155e37 100644 --- a/tests/unit/embeddings_test.py +++ b/tests/unit/embeddings_test.py @@ -84,6 +84,43 @@ def test_attend(self): self.assertEqual(outputs.shape, (batch_size, seq_len, num_embeddings)) + def test_attend_on_embedding_matches_transposed_dot(self): + """`attend_on_embedding` contracts over the table's feature axis directly. + + Expressing the transpose as dimension numbers instead of materializing + `table.T` is what lets the input lookup and the tied output head share one + bf16 cast under `shard_mode: explicit`. It reassociates the accumulation, so + the logits agree to float rounding rather than bit for bit. + """ + table = jax.random.normal(jax.random.PRNGKey(0), (32, 8)) + for query_shape in ((8,), (3, 8), (2, 3, 8)): + with self.subTest(query_shape=query_shape): + query = jax.random.normal(jax.random.PRNGKey(1), query_shape) + expected = jnp.dot(query, jnp.asarray(table, jnp.bfloat16).T, preferred_element_type=jnp.float32) + got = embeddings.attend_on_embedding(query, table, jnp.float32, self.cfg) + self.assertEqual(got.shape, expected.shape) + np.testing.assert_allclose(got, expected, rtol=1e-5, atol=1e-6) + + def test_attend_on_embedding_weight_gradient_in_table_order(self): + """The tied head's table gradient comes out in the table's own axis order. + + With `table.T` the gradient is built transposed and handed back through a + second transpose, which under explicit sharding survives to codegen. The + gradient itself must be unchanged. + """ + table = jax.random.normal(jax.random.PRNGKey(0), (32, 8)) + query = jax.random.normal(jax.random.PRNGKey(1), (2, 3, 8)) + + def loss(t, attend): + return jnp.sum(jnp.sin(attend(query, t))) + + got = jax.grad(loss)(table, lambda q, t: embeddings.attend_on_embedding(q, t, jnp.float32, self.cfg)) + expected = jax.grad(loss)( + table, lambda q, t: jnp.dot(q, jnp.asarray(t, jnp.bfloat16).T, preferred_element_type=jnp.float32) + ) + self.assertEqual(got.shape, table.shape) + np.testing.assert_allclose(got, expected, rtol=1e-5, atol=1e-6) + class RotaryEmbeddingTest(unittest.TestCase): """Tests for RotaryEmbedding.""" From efb45afda5f9b19a50d13334e68cbc2be4af4d88 Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Thu, 3 Sep 2026 11:55:17 +0000 Subject: [PATCH 3/9] Build only the YaRN frequency rows a step actually reads `YarnRotaryEmbedding.freqs_cis` built the whole `[max_position_embeddings, half_dim]` table and then indexed `max_target_length` rows out of it. The table is traced, not a constant XLA can hoist, so it was rebuilt inside every layer of every step: on deepseek2-16b that is `f32[163840, 32]` per layer per step, of which 0.6% is read. Row `p` is `exp(1j * p * corrected_freqs)`, so the rows can be produced straight from `position`. `freqs_cis_at` does that; `freqs_cis` stays as the reference definition the new test checks against, and the two agree bit for bit. Measured on deepseek at emb 2048 / mlp 8192 / L8 this is -1.4% of step time, and the same -1.4% under `shard_mode: auto` and `explicit` alike. `_cis` writes `exp(1j * theta)` as `cos(theta) + 1j*sin(theta)`. It is bit-identical, but it avoids XLA's overflow-safe complex `exponential` expansion, which evaluates `exp(real(1j*theta)) == exp(0)` over the whole tensor -- twice under explicit sharding, where the `Sharding` custom-call on the broadcast `1j` blocks the reassociation that would let CSE merge them. --- src/maxtext/layers/embeddings.py | 44 ++++++++++++++++++++++++-------- tests/unit/embeddings_test.py | 22 ++++++++++++++++ 2 files changed, 56 insertions(+), 10 deletions(-) diff --git a/src/maxtext/layers/embeddings.py b/src/maxtext/layers/embeddings.py index b0f8d23f6b..c86a1ea77c 100644 --- a/src/maxtext/layers/embeddings.py +++ b/src/maxtext/layers/embeddings.py @@ -35,6 +35,20 @@ _MAX_WAVELENGTH = 10_000 +def _cis(theta: Array) -> Array: + """`exp(1j * theta)` for real `theta`, written as `cos(theta) + 1j*sin(theta)`. + + Bit-identical to `jnp.exp(1j * theta)`, but it does not go through XLA's + overflow-safe complex `exponential` expansion, which evaluates + `exp(real(1j*theta)) == exp(0)` over the whole tensor. Under + `shard_mode: explicit` it evaluates it *twice*: the `Sharding` custom-call on + the broadcast `1j` keeps the simplifier from commuting the constant to the + right, so the reassociation that lets CSE merge the two `exponential`s never + fires. + """ + return jax.lax.complex(jnp.cos(theta), jnp.sin(theta)) + + def _maybe_move_embedding_to_device(embedding_table: Array, config: Config) -> Array: """Moves embedding table to device if parameter offloading is enabled.""" if config.parameter_memory_host_offload: @@ -880,8 +894,8 @@ def __init__( raise ValueError("Embedding dim for rotary position embedding must be a multiple of 2.") @property - def freqs_cis(self): - """Frequencies for rotary embedding.""" + def corrected_freqs(self): + """The per-dimension rotary frequencies, shape [half_dim].""" half_dim = self.embedding_dims // 2 # Compute base frequencies for each (even-indexed) dimension. # (Note: We use jnp.arange with float32 for precision.) @@ -897,15 +911,26 @@ def freqs_cis(self): ) smooth = 1 - self._linear_ramp_factor(low, high, half_dim) # The corrected frequency is a weighted mix of the scaled and base values. - freqs = freqs / self.rope_factor * (1 - smooth) + freqs * smooth + return freqs / self.rope_factor * (1 - smooth) + freqs * smooth + @property + def freqs_cis(self): + """Frequencies for every position, shape [max_position_embeddings, half_dim].""" # Precompute frequencies for all positions by taking the outer product. t = jnp.arange(self.max_position_embeddings, dtype=jnp.float32) # shape [max_position_embeddings] # This gives a [max_position_embeddings, half_dim] tensor with rows as time steps. - freqs = jnp.outer(t, freqs) + return _cis(jnp.outer(t, self.corrected_freqs)) - # Compute the complex “cis” values: exp(i * theta). - return jnp.exp(1j * freqs) # shape [max_position_embeddings, half_dim] + def freqs_cis_at(self, position: Array, out_sharding=None) -> Array: + """`freqs_cis[position]`, without building the whole table to index it. + + Row `p` of `freqs_cis` is `exp(1j * p * corrected_freqs)`, so the rows a step + actually reads can be produced straight from `position`. The table itself is + traced, not a constant XLA can hoist or fold, so building it costs + `max_position_embeddings` rows *inside every layer, every step* — on + deepseek2-16b that is 163,840 rows of which `max_target_length` are read. + """ + return _cis(jnp.einsum("bs,h->bsh", position.astype(jnp.float32), self.corrected_freqs, out_sharding=out_sharding)) def _find_correction_dim(self, num_rotations: float, dim: int, base: float, max_position_embeddings: int) -> float: """Compute the correction dimension for a given number of rotations.""" @@ -976,10 +1001,9 @@ def __call__(self, inputs: Array, position: None | Array = None) -> Array: else: position = position.astype(jnp.int32) - # Lookup the precomputed frequencies using the position indices. - # self.freqs_cis has shape [max_position_embeddings, half_dim] so we use jnp.take along axis 0. - # After indexing, shape becomes [B, S, half_dim]; we then add an axis for the heads. - freqs = self.freqs_cis.at[position].get(out_sharding=self.freqs_sharding) # shape: [B, S, half_dim] + # Build the frequencies for these positions directly, rather than indexing + # them out of the full [max_position_embeddings, half_dim] table. + freqs = self.freqs_cis_at(position, out_sharding=self.freqs_sharding) # shape: [B, S, half_dim] freqs = freqs[:, :, jnp.newaxis, :] # shape: [B, S, 1, half_dim] if self.interleave and self.pairwise: diff --git a/tests/unit/embeddings_test.py b/tests/unit/embeddings_test.py index ac54155e37..45664d982d 100644 --- a/tests/unit/embeddings_test.py +++ b/tests/unit/embeddings_test.py @@ -332,6 +332,28 @@ def test_explicit_shard_mode_call(self): default_outputs = default_layer(inputs, position=position) np.testing.assert_allclose(outputs, default_outputs, atol=1e-5) + def test_freqs_cis_at_matches_indexing_the_table(self): + """`freqs_cis_at` is what `__call__` uses; the full table is what it stands in for. + + Row `p` of `freqs_cis` is `exp(1j * p * corrected_freqs)`, so the two agree + bit for bit — `arange(max_position_embeddings)[p]` is exactly `float(p)` for + every representable position. + """ + layer = embeddings.YarnRotaryEmbedding( + embedding_dims=64, + mesh=self.mesh, + max_position_embeddings=163840, + original_max_position_embeddings=4096, + rngs=self.rngs, + ) + table = layer.freqs_cis + self.assertEqual(table.shape, (163840, 32)) + # Sample the bottom, the middle and the very top of the table. + for lo in (0, 4096, 163840 - 128): + with self.subTest(lo=lo): + position = jnp.arange(lo, lo + 128, dtype=jnp.int32)[jnp.newaxis, :] + np.testing.assert_array_equal(layer.freqs_cis_at(position), table.at[position].get()) + def test_pairwise_explicit_shard_mode_call(self): layer = embeddings.YarnRotaryEmbedding( embedding_dims=4, From 2b29dda339f845916f2ec71547a7b14c530eee98 Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Thu, 3 Sep 2026 11:55:32 +0000 Subject: [PATCH 4/9] Add dense_weight_grad_in_kernel_order for the in-loop projections The in-loop counterpart of `lm_head_weight_grad_in_kernel_order`: wire the same `custom_vjp` to the three `MlpBlock` projections, the four attention projections and the nine MLA / Indexer projections, so their weight gradients are contracted straight into each kernel's stored axis order. Unlike the LM-head flag this is a trade, not a fix, so it defaults off. Under `shard_mode: explicit` the `Sharding` custom-calls decouple the gradient's layout from the stored parameter's, leaving XLA free to pick it; the flag declines that freedom and reproduces `shard_mode: auto`'s layout op for op. Whether that is a win depends on the scanned gradient stack's depth. At emb 2048 / mlp 8192 it is worth 0.90pp on gemma3-4b and 0.16pp on deepseek3-16b, whose stacks are 3 and 1 deep, and costs 0.06%-0.60% on the six models whose stacks are 16 deep -- two points below the line against six above is not enough to fit a default on, so it ships off with the two models named in the docs. `_permuted_sharding` handles the >2-D kernels the in-loop sites have; on a 2-D kernel it degenerates to the identity, which is why the LM head never needed it. Stored kernels and their initialization are untouched, so no checkpoint conversion is needed, and the flag is inert under `shard_mode: auto` and on quantized layers. See docs/guides/optimization/shard_mode_performance.md section 4.8. --- src/maxtext/configs/base.yml | 11 ++++ src/maxtext/configs/types.py | 15 ++++++ src/maxtext/layers/attention_mla.py | 8 +++ src/maxtext/layers/attentions.py | 4 ++ src/maxtext/layers/linears.py | 35 ++++++++++--- tests/unit/linears_test.py | 80 +++++++++++++++++++++++++++++ 6 files changed, 145 insertions(+), 8 deletions(-) diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 06ee18f976..82ba134c78 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -718,6 +718,17 @@ shard_optimizer_over_data: false # when dense MLP weight matrices are sharded on both fsdp and fsdp-transpose axes, use two separate all-gather calls dense_fsdp_use_two_stage_all_gather: false +# The in-loop counterpart of lm_head_weight_grad_in_kernel_order: compute the attention q/k/v/out and MLP wi/wo +# weight gradients with a hand-written backward pass that contracts straight into each kernel's stored orientation. +# Unlike the LM head flag this is a trade, not a fix, so it defaults off. Under shard_mode: explicit XLA is free to +# pick the gradient's layout -- the Sharding custom-calls decouple it from the stored parameter's -- and this flag +# declines that freedom, reproducing shard_mode: auto's layout op for op. Measured, that is worth roughly -0.3% on +# gemma3 and -0.2% on deepseek, whose scanned gradient stacks are 3 and 1 deep, and costs 0.06%-0.60% on the models +# whose stacks are 16 deep (see docs/guides/optimization/shard_mode_performance.md section 4.8). Turn it on for +# gemma3 and deepseek; measure before turning it on for anything else. Stored kernels and their initialization are +# untouched, so no checkpoint conversion is needed. No effect under shard_mode: auto or on quantized layers. +dense_weight_grad_in_kernel_order: false + # Unless explicitly specified, the number of TPU slices is automatically determined. It should only be set for # disaggregated reinforcement learning workloads using multiple slices. For ahead of time compilation, # you should set compile_toplogy_num_slices, which will in turn set this value. For non-TPU environments this is set to 1. diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index e220bec38a..324defbe20 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1267,6 +1267,21 @@ class LayoutAndSharding(BaseModel): False, description="Use two separate All-Gather calls for dense MLP weights sharded on both FSDP and FSDP-transpose.", ) + dense_weight_grad_in_kernel_order: bool = Field( + False, + description=( + "Compute the per-layer attention and MLP projection weight gradients directly in each kernel's stored " + "axis order via a custom_vjp, so autodiff emits no transpose after the gradient dot. This is the " + "in-loop counterpart of lm_head_weight_grad_in_kernel_order, but unlike that flag it is a trade rather " + "than a fix, so it defaults off. Under shard_mode=explicit the Sharding custom-calls decouple the " + "gradient's layout from the stored parameter's, leaving XLA free to pick it; this flag declines that " + "freedom and reproduces shard_mode=auto's layout op for op. Measured, that is worth roughly -0.3% on " + "gemma3 and -0.2% on deepseek, whose scanned gradient stacks are 3 and 1 deep, and costs 0.06%-0.60% on " + "the models whose stacks are 16 deep (docs/guides/optimization/shard_mode_performance.md section 4.8). " + "Measure before turning it on for anything else. Stored kernels and their initialization are untouched. " + "No effect under shard_mode=auto or on quantized layers." + ), + ) internal_compile: bool = Field( False, description="Use internal_compile to bypass open-source topology mappings.", diff --git a/src/maxtext/layers/attention_mla.py b/src/maxtext/layers/attention_mla.py index 13dd34a1d4..38eff43890 100644 --- a/src/maxtext/layers/attention_mla.py +++ b/src/maxtext/layers/attention_mla.py @@ -138,6 +138,7 @@ def __init__( quant=self.quant, matmul_precision=self.config.matmul_precision, shard_mode=self.config.shard_mode, + weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order, rngs=self.rngs, ) @@ -153,6 +154,7 @@ def __init__( quant=self.quant, matmul_precision=self.config.matmul_precision, shard_mode=self.config.shard_mode, + weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order, rngs=self.rngs, ) @@ -172,6 +174,7 @@ def __init__( quant=None, matmul_precision=self.config.matmul_precision, shard_mode=self.config.shard_mode, + weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order, rngs=self.rngs, ) @@ -808,6 +811,7 @@ def _init_projections(self, inputs_q_shape: Tuple, inputs_kv_shape: Tuple) -> No quant=self.quant, matmul_precision=self.config.matmul_precision, shard_mode=self.config.shard_mode, + weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order, rngs=self.rngs, ) else: @@ -823,6 +827,7 @@ def _init_projections(self, inputs_q_shape: Tuple, inputs_kv_shape: Tuple) -> No quant=self.quant, matmul_precision=self.config.matmul_precision, shard_mode=self.config.shard_mode, + weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order, rngs=self.rngs, ) self.q_norm = RMSNorm( @@ -844,6 +849,7 @@ def _init_projections(self, inputs_q_shape: Tuple, inputs_kv_shape: Tuple) -> No quant=self.quant, matmul_precision=self.config.matmul_precision, shard_mode=self.config.shard_mode, + weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order, rngs=self.rngs, ) @@ -859,6 +865,7 @@ def _init_projections(self, inputs_q_shape: Tuple, inputs_kv_shape: Tuple) -> No quant=self.quant, matmul_precision=self.config.matmul_precision, shard_mode=self.config.shard_mode, + weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order, rngs=self.rngs, ) self.kv_norm = RMSNorm( @@ -883,6 +890,7 @@ def _init_projections(self, inputs_q_shape: Tuple, inputs_kv_shape: Tuple) -> No quant=self.quant, matmul_precision=self.config.matmul_precision, shard_mode=self.config.shard_mode, + weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order, rngs=self.rngs, ) diff --git a/src/maxtext/layers/attentions.py b/src/maxtext/layers/attentions.py index c61b945b0e..37623ab7b8 100644 --- a/src/maxtext/layers/attentions.py +++ b/src/maxtext/layers/attentions.py @@ -697,6 +697,7 @@ def query_init(*args): matmul_precision=self.config.matmul_precision, use_bias=self.use_bias_in_projections, shard_mode=self.config.shard_mode, + weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order, rngs=self.rngs, ) @@ -733,6 +734,7 @@ def init_kv_w(self, inputs_kv_shape: Tuple) -> nnx.Module: weight_dtype=self.weight_dtype, quant=self.quant, shard_mode=self.config.shard_mode, + weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order, matmul_precision=self.config.matmul_precision, use_bias=self.use_bias_in_projections, rngs=self.rngs, @@ -774,6 +776,7 @@ def init_qkv_w(self, inputs_shape: Tuple) -> nnx.Module: weight_dtype=self.weight_dtype, quant=self.quant, shard_mode=self.config.shard_mode, + weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order, matmul_precision=self.config.matmul_precision, use_bias=self.use_bias_in_projections, rngs=self.rngs, @@ -829,6 +832,7 @@ def init_out_w(self, output_dim: int) -> nnx.Module: weight_dtype=self.weight_dtype, quant=self.quant, shard_mode=self.config.shard_mode, + weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order, matmul_precision=self.config.matmul_precision, use_bias=False if self.is_qwen2 else self.use_bias_in_projections, rngs=self.rngs, diff --git a/src/maxtext/layers/linears.py b/src/maxtext/layers/linears.py index 92beaeeee6..d78880afaf 100644 --- a/src/maxtext/layers/linears.py +++ b/src/maxtext/layers/linears.py @@ -96,10 +96,26 @@ def _aval_sharding(value): sharding = jax.typeof(value).sharding except Exception: # pylint: disable=broad-except return None - spec = getattr(sharding, "spec", None) - if spec is None or all(p is None for p in spec): - return None - return sharding + # A fully replicated aval still needs its `out_sharding` spelled out: the + # gradient dots below contract over sharded token axes, which JAX refuses to + # resolve on its own. + return sharding if getattr(sharding, "spec", None) is not None else None + + +def _permuted_sharding(value, produced_axes): + """`value`'s sharding with its axes reordered into `produced_axes` order. + + The gradient dots emit the operand's axes in `produced_axes` order and only + then permute them back, so the `out_sharding` handed to the dot has to be + permuted the same way. For every layout this path is enabled for + `produced_axes` is the identity and this is just `_aval_sharding`. + """ + sharding = _aval_sharding(value) + produced_axes = tuple(produced_axes) + if sharding is None or produced_axes == tuple(range(len(produced_axes))): + return sharding + spec = sharding.spec + return sharding.update(spec=jax.sharding.PartitionSpec(*(spec[a] for a in produced_axes))) def _restore_axis_order(value, produced_axes): @@ -156,26 +172,26 @@ def _dot_general_kernel_ordered_grad_bwd(dimension_numbers, precision, out_shard # to them one for one. The result carries the kernel's contracting axes (in # ascending operand order) followed by its free axes -- the kernel's own order # whenever it is stored in_features-major. + dk_axes = tuple(rhs_contract[lhs_contract.index(a)] for a in sorted(lhs_contract)) + rhs_free dk = lax.dot_general( inputs, g, ((lhs_free, tuple(range(n_lhs_free))), ((), ())), precision=precision, - out_sharding=_aval_sharding(kernel), + out_sharding=_permuted_sharding(kernel, dk_axes), ) - dk_axes = tuple(rhs_contract[lhs_contract.index(a)] for a in sorted(lhs_contract)) + rhs_free dk = _restore_axis_order(dk, dk_axes) # dx contracts g's trailing axes against the kernel's free axes, likewise # paired one for one. + dx_axes = lhs_free + tuple(lhs_contract[rhs_contract.index(a)] for a in sorted(rhs_contract)) dx = lax.dot_general( g, kernel, ((tuple(range(n_lhs_free, g.ndim)), rhs_free), ((), ())), precision=precision, - out_sharding=_aval_sharding(inputs), + out_sharding=_permuted_sharding(inputs, dx_axes), ) - dx_axes = lhs_free + tuple(lhs_contract[rhs_contract.index(a)] for a in sorted(rhs_contract)) dx = _restore_axis_order(dx, dx_axes) return dx, dk @@ -631,6 +647,7 @@ def __init__( quant=self.quant, use_bias=self.use_bias, shard_mode=self.config.shard_mode, + weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order, matmul_precision=self.config.matmul_precision, mesh=self.mesh, use_two_stage_all_gather=self.config.dense_fsdp_use_two_stage_all_gather, @@ -650,6 +667,7 @@ def __init__( quant=self.quant, use_bias=self.use_bias, shard_mode=self.config.shard_mode, + weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order, matmul_precision=self.config.matmul_precision, mesh=self.mesh, use_two_stage_all_gather=self.config.dense_fsdp_use_two_stage_all_gather, @@ -668,6 +686,7 @@ def __init__( quant=self.quant, use_bias=self.use_bias, shard_mode=self.config.shard_mode, + weight_grad_in_kernel_order=self.config.dense_weight_grad_in_kernel_order, matmul_precision=self.config.matmul_precision, mesh=self.mesh, use_two_stage_all_gather=self.config.dense_fsdp_use_two_stage_all_gather, diff --git a/tests/unit/linears_test.py b/tests/unit/linears_test.py index 976458a40d..c720b9a3a5 100644 --- a/tests/unit/linears_test.py +++ b/tests/unit/linears_test.py @@ -325,6 +325,86 @@ def loss(l, x): self.assertEqual(count_transposes(ordered), 0) self.assertGreater(count_transposes(default), 0) + @staticmethod + def _explicit_mesh(): + return jax.make_mesh((1, 1), ("fsdp", "tensor"), axis_types=(jax.sharding.AxisType.Explicit,) * 2) + + def test_weight_grad_in_kernel_order_with_replicated_kernel(self): + """A replicated kernel still has to spell out the gradient's sharding. + + The gradient dot contracts over the sharded token axes, which JAX refuses to + resolve on its own -- omitting `out_sharding` there raises instead of + defaulting to replicated. + """ + mesh = self._explicit_mesh() + with jax.set_mesh(mesh): + inputs = jax.sharding.reshard( + jax.random.normal(jax.random.PRNGKey(0), (2, 4, 8)), + jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec("fsdp", None, None)), + ) + kernel = jax.sharding.reshard( + jax.random.normal(jax.random.PRNGKey(1), (8, 16)), + jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(None, None)), + ) + dimension_numbers = (((2,), (0,)), ((), ())) + + def loss(x, k, dot): + return jnp.sum(jnp.sin(dot(x, k))) + + ordered = functools.partial( + linears._dot_general_kernel_ordered_grad, # pylint: disable=protected-access + dimension_numbers=dimension_numbers, + precision=None, + out_sharding=None, + ) + reference = functools.partial(jax.lax.dot_general, dimension_numbers=dimension_numbers) + grad_ordered = jax.jit(jax.grad(loss, argnums=(0, 1)), static_argnums=2)(inputs, kernel, ordered) + grad_reference = jax.jit(jax.grad(loss, argnums=(0, 1)), static_argnums=2)(inputs, kernel, reference) + for got, want in zip(grad_ordered, grad_reference): + self.assertEqual(jax.typeof(got).sharding.spec, jax.typeof(want).sharding.spec) + np.testing.assert_allclose(got, want, rtol=1e-5, atol=1e-6) + + def test_weight_grad_in_kernel_order_with_permuted_axes(self): + """Dimension numbers that do not pair operand axes in ascending order. + + No `DenseGeneral` layout builds these today, but the backward pass permutes + its dots' outputs back into operand order and the `out_sharding` it hands + those dots has to be permuted the same way or the custom_vjp returns a + cotangent whose type does not match its primal. + """ + mesh = self._explicit_mesh() + with jax.set_mesh(mesh): + # `inputs` contracts its trailing axes against the kernel's *leading* axes + # in reverse order, so both the dk and the dx dot come out permuted, and + # the permutation moves a sharded axis. + inputs = jax.sharding.reshard( + jax.random.normal(jax.random.PRNGKey(0), (2, 5, 3, 4)), + jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec("fsdp", None, "tensor", None)), + ) + kernel = jax.sharding.reshard( + jax.random.normal(jax.random.PRNGKey(1), (4, 3, 7)), + jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(None, "tensor", None)), + ) + dimension_numbers = (((2, 3), (1, 0)), ((), ())) + out_sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec("fsdp", None, None)) + + def loss(x, k, dot): + return jnp.sum(jnp.sin(dot(x, k))) + + ordered = functools.partial( + linears._dot_general_kernel_ordered_grad, # pylint: disable=protected-access + dimension_numbers=dimension_numbers, + precision=None, + out_sharding=out_sharding, + ) + reference = functools.partial(jax.lax.dot_general, dimension_numbers=dimension_numbers, out_sharding=out_sharding) + grad_ordered = jax.jit(jax.grad(loss, argnums=(0, 1)), static_argnums=2)(inputs, kernel, ordered) + grad_reference = jax.jit(jax.grad(loss, argnums=(0, 1)), static_argnums=2)(inputs, kernel, reference) + for got, want in zip(grad_ordered, grad_reference): + self.assertEqual(got.shape, want.shape) + self.assertEqual(jax.typeof(got).sharding.spec, jax.typeof(want).sharding.spec) + np.testing.assert_allclose(got, want, rtol=1e-5, atol=1e-6) + def _run_dense_test(self, axis, in_feat_shape, expected_shape): batch_size = 2 seq_len = 3 From d98257d44543f35ae7bb63067013134c7bcc6763 Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Thu, 3 Sep 2026 11:55:50 +0000 Subject: [PATCH 5/9] Turn lm_head_weight_grad_in_kernel_order on where it can act Re-measured across all eight onboarded models at emb 2048 / mlp 8192 / 16 layers -- above the geometry gate an emb-512 proxy sits below -- the flag has never lost to `shard_mode: auto`: on any model, at any depth (L12/L14/L16), at any vocab (4096 or 32000). Without it, explicit is +3.25% on qwen3-8b and +0.92% on deepseek3-16b. There is no configuration measured where leaving it off is the better choice, so off is the wrong default. It is now tri-state. `None` resolves to on for an untied head under `shard_mode: explicit` and off everywhere else, which is exactly where it can do something: under `auto` XLA folds the transpose itself, and a tied head has no kernel of its own. An explicit `true` on a tied head is now an error rather than silently ignored. With this default and nothing else in the config but `shard_mode: explicit`, the median across the eight models is -0.330% and the worst case is gemma3 at +0.803%; adding `dense_weight_grad_in_kernel_order` on the two models that want it leaves no model regressing by more than 0.041%. --- src/maxtext/configs/base.yml | 5 +++-- src/maxtext/configs/types.py | 25 +++++++++++++++++-------- tests/unit/pyconfig_test.py | 21 +++++++++++++++++++++ 3 files changed, 41 insertions(+), 10 deletions(-) diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 82ba134c78..82dd903898 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -206,8 +206,9 @@ cast_logits_to_fp32: true # whether to cast the logits to fp32. the higher preci # barrier algsimp will not cross, so the gradient is left sharded on its minor-most dimension and loses the # all-reduce -> reduce-scatter fusion (see docs/guides/optimization/shard_mode_performance.md). The stored kernel # and its initialization are untouched, so no checkpoint conversion is needed; gradients differ from the default -# rule only by floating-point reassociation. No effect under shard_mode: auto, where XLA folds the transpose itself. -lm_head_weight_grad_in_kernel_order: false +# rule only by floating-point reassociation. No effect under shard_mode: auto, where XLA folds the transpose itself, +# and none on a tied head, which has no kernel of its own -- so None means on exactly where it can do something. +lm_head_weight_grad_in_kernel_order: None # on for an untied model under shard_mode: explicit if None float32_qk_product: false # in dot_product attention, whether to cast to fp32 the inputs to qk product float32_logits: false # in dot_product attention, whether to cast to fp32 the inputs to softmax mla_qk_head_chunk_size: 0 # Limits HBM footprint by sequentially evaluating the QK matrix in the Indexer across the unsharded local heads dimension natively. diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 324defbe20..39a532df3c 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -604,13 +604,14 @@ class LogitsAndLoss(BaseModel): ) logits_dot_in_fp32: bool = Field(False, description="Use fp32 for the logits dot product for stability.") cast_logits_to_fp32: bool = Field(True, description="Whether to cast the final logits to fp32.") - lm_head_weight_grad_in_kernel_order: bool = Field( - False, + lm_head_weight_grad_in_kernel_order: bool | None = Field( + None, description=( "Compute the untied LM head's weight gradient directly in the kernel's stored axis order via a " "custom_vjp, so autodiff emits no transpose after the gradient dot. Recovers the weight-gradient " "reduce-scatter that shard_mode=explicit otherwise loses, leaving the stored kernel and its " - "initialization untouched. No effect under shard_mode=auto." + "initialization untouched. No effect under shard_mode=auto. None means on for an untied model " + "under shard_mode=explicit and off everywhere else." ), ) final_logits_soft_cap: None | NonNegativeFloat = Field( @@ -3294,11 +3295,19 @@ def validate_shard_embed_moe_on_fsdp(self) -> "MaxTextConfig": return self @model_validator(mode="after") - def validate_lm_head_weight_grad_in_kernel_order(self) -> "MaxTextConfig": - """Reject lm_head_weight_grad_in_kernel_order where it cannot be honored.""" - if not self.lm_head_weight_grad_in_kernel_order: - return self - if self.logits_via_embedding: + def resolve_lm_head_weight_grad_in_kernel_order(self) -> "MaxTextConfig": + """Resolve the tri-state flag, and reject it where it cannot be honored. + + The transpose it removes only exists under explicit sharding, and on an untied + model removing it has been a win or a wash on every configuration measured + (docs/guides/optimization/shard_mode_performance.md section 5) -- on qwen3-8b it + is the difference between +3.25% and -0.41% against `auto`. So the default is + "on wherever it can do anything", and writing the flag out is only needed to + reproduce a measurement. + """ + if self.lm_head_weight_grad_in_kernel_order is None: + self.lm_head_weight_grad_in_kernel_order = self.shard_mode == ShardMode.EXPLICIT and not self.logits_via_embedding + elif self.lm_head_weight_grad_in_kernel_order and self.logits_via_embedding: raise ValueError( "lm_head_weight_grad_in_kernel_order only applies to the untied LM head, but logits_via_embedding is True." ) diff --git a/tests/unit/pyconfig_test.py b/tests/unit/pyconfig_test.py index a5395c516b..9b42f52736 100644 --- a/tests/unit/pyconfig_test.py +++ b/tests/unit/pyconfig_test.py @@ -257,6 +257,27 @@ def test_explicit_sharding_mistral_decoder_support(self): ) self.assertEqual(config.decoder_block.value, decoder_block) + def test_lm_head_weight_grad_in_kernel_order_default(self): + """Left unset, the flag is on exactly where it can do something.""" + + def resolve(**kwargs): + return pyconfig.initialize( + [os.path.join(MAXTEXT_PKG_DIR, "train.py"), get_test_config_path()], + skip_jax_distributed_system=True, + **kwargs, + ).lm_head_weight_grad_in_kernel_order + + # The transpose it removes only exists under explicit sharding, on an untied head. + self.assertTrue(resolve(shard_mode="explicit")) + self.assertFalse(resolve(shard_mode="auto")) + self.assertFalse(resolve(shard_mode="explicit", logits_via_embedding=True)) + # Writing it out still wins over the default, in both directions. + self.assertFalse(resolve(shard_mode="explicit", lm_head_weight_grad_in_kernel_order=False)) + self.assertTrue(resolve(shard_mode="auto", lm_head_weight_grad_in_kernel_order=True)) + # But asking for it on a tied head is still a hard error, not a silent no-op. + with self.assertRaisesRegex(Exception, "only applies to the untied LM head"): + resolve(shard_mode="explicit", logits_via_embedding=True, lm_head_weight_grad_in_kernel_order=True) + def test_resolve_config_path(self): self.assertEqual(resolve_config_path("foo"), os.path.join("src", "foo")) self.assertEqual(resolve_config_path(__file__), __file__) From 6828d80a9a2fa8089acf176720f750c4b0d6c823 Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Thu, 3 Sep 2026 11:56:03 +0000 Subject: [PATCH 6/9] Revise the shard_mode performance guide for the 2026-09-03 measurements Re-measures the whole matrix at above-gate geometry with nothing written out but `shard_mode` (section 5.4, the table to quote), documents the LM-head flag's new default, adds section 4.8 on what the dense-layer flag actually trades, section 5.5 on the YaRN table, and section 6.3 on the rotated reduce-scatter -- explicit's largest and most stable win, 480 of 495 `collective-permute-start` ops deleted on every untied model. Also records what did not work: thirteen candidates that were implemented or measured and refuted, and a systematic hunt for any remaining explicit-only cost on llama2/mistral/qwen3 that found nothing above ~0.10% of step. --- .../optimization/shard_mode_performance.md | 428 ++++++++++++++++-- 1 file changed, 393 insertions(+), 35 deletions(-) diff --git a/docs/guides/optimization/shard_mode_performance.md b/docs/guides/optimization/shard_mode_performance.md index 12a3359725..7ecf65e2c6 100644 --- a/docs/guides/optimization/shard_mode_performance.md +++ b/docs/guides/optimization/shard_mode_performance.md @@ -21,8 +21,11 @@ Revised 2026-09-02 with the shipped fix (`lm_head_weight_grad_in_kernel_order`) and two negative results; see §8 item 1 and §8 item 2. Where a claim was superseded the original is kept and marked, because two of the corrections reverse a conclusion that looked solid at the time. +Revised 2026-09-03: the whole matrix re-measured at above-gate geometry (§5.4), the LM-head flag +turned **on by default** where it can act, a second flag added for the dense-layer gradients (§4.8), +and a mode-independent 1.4% win found and shipped along the way (§5.5). -`shard_mode` (`src/maxtext/configs/base.yml:545`) selects how MaxText expresses sharding: +`shard_mode` (`src/maxtext/configs/base.yml:554`) selects how MaxText expresses sharding: - **`auto`** — `jax.lax.with_sharding_constraint`. A *hint*. GSPMD/Shardy is free to propagate it backwards into producers, refine it, sink it, or drop it. @@ -134,17 +137,40 @@ ______________________________________________________________________ **−0.57%**. At production geometry, where there was nearly nothing to recover, it is neutral (**+0.03%** and **−0.57%**). See §8 item 1. + > *Revised (2026-09-03).* "Off by default" no longer holds. Re-measured across all eight models + > at above-gate geometry the flag has **never lost to `auto`** — on any model, at any depth + > (L12/L14/L16), at any vocab (4096 or 32000) — while explicit without it is +3.25% on qwen3. + > It is therefore now tri-state and resolves to **on** wherever it can act: `shard_mode: explicit` and an untied head. Writing it out is only needed to reproduce a measurement. + +11. **With that default, explicit is faster than `auto` on five of eight onboarded models and inside + the run-to-run spread on a sixth.** The one number a user needs is §5.4: at emb 2048 / mlp 8192 / + 16 layers, with nothing in the config but `shard_mode`, the median across the eight models is + **−0.330%** and the worst case is gemma3 at **+0.803%** — down from a +3.25% worst case before + this revision. Both remaining non-wins (gemma3, deepseek) are fixed by a second flag, + `dense_weight_grad_in_kernel_order`, which ships default-off because its sign is model-dependent + (§4.8): it is worth 0.90 pp on gemma3 and 0.16 pp on deepseek, and costs 0.06–0.60% on the other + six. With those two per-model settings applied, no model regresses by more than 0.041%. + +12. **The largest single win found here is not a sharding fix at all.** `YarnRotaryEmbedding` built + its whole `[max_position_embeddings, head_dim/2]` frequency table inside every layer of every + step in order to read `max_target_length` rows out of it — `f32[163840, 32]` on deepseek, 0.6% + of it used. Computing the needed rows directly from `position` is **−1.4% of step time, equally + in both modes** (§5.5). It is called out here because it was invisible to every mode-comparison + in this document until a detector was pointed at ops whose output is ≥ 8× their input: A/B + diffing two modes cannot find work that both modes do. + > *Superseded (2026-09-02).* This item previously recommended `lm_head_kernel_transposed`, which > stores the kernel as `[vocab, embed]`. That flag reaches the same layout but is slower on four > of seven configs, has no checkpoint-conversion path, and — because `jax.random` draws a > different matrix from the same seed at the flipped shape — perturbs the loss ~35× more than > explicit's own noise. It has been removed rather than shipped alongside a strictly better flag. -**Practical guidance:** at realistic depth/sequence/batch the penalty falls to a few tenths of a -percent and capability should decide the choice. But it does **not** reliably vanish with *width*, -and under pure data parallelism it is +10%. On an untied model where you have committed to explicit, -set `lm_head_weight_grad_in_kernel_order: true` (§8) rather than assuming scale will hide it; it is -checkpoint-compatible and inert under `auto`, so it costs nothing to leave on. +**Practical guidance:** as of this revision, `shard_mode: explicit` with nothing else written out is +the faster mode on most onboarded models at realistic geometry (§5.4), so choose on capability +(§6, §7) and treat performance as a tie-breaker rather than a tax. Two carve-outs remain: **gemma3 +and deepseek want `dense_weight_grad_in_kernel_order: true`** (§4.8), and **pure data parallelism** +was +10% before the LM-head fix — re-verify it for your geometry rather than assuming the fix +covers it, since the DP measurement predates the default flip. ______________________________________________________________________ @@ -732,6 +758,88 @@ counts before comparing dumps.** The one durable gemma3 finding is the last ledg gradient all-reduce repack is a real tax that **grows with depth** (+112 → +219 → +379 µs), even though the total flips negative because the fusion win grows faster. +> *Correction (2026-09-03).* The all-reduce repack tax does not survive the move to above-gate +> geometry. Re-measured at emb 2048 / mlp 8192 / L18, gemma3's `all-reduce` opcode delta is +> **+41.4 µs of 2.52 s of device-op time** across the whole profile (84 ops in both modes) — three +> orders of magnitude below the emb-512 figure as a share of step. It was a small-tensor effect: +> at emb 512 the repacked tuple is latency-bound, and merging it removes the only work that was +> available to overlap it. Read the ladder above as a property of the shrink, not of gemma3. +> What *does* survive at real width is §4.8. + +### 4.8 Consequence C — the dense-layer weight gradient's layout is un-pinned, and that cuts both ways + +§4.2 and §4.3 are about *one* dot, the untied LM head, and they are about a transpose that XLA +declines to fold. There is a second, broader effect on every other `DenseGeneral` in the model, and +it is not a missed optimization at all — it is a **layout choice that explicit hands to XLA and +`auto` does not**. + +Under `auto` the weight gradient's layout is derived from the stored parameter's, because GSPMD +propagates the constraint through. Under explicit the `Sharding` custom-calls sit between the two, +layout assignment cannot propagate across them, and XLA picks whatever layout the gradient `dot` +itself prefers — `{1,2,0}` for an `f32[1, embed_shard, mlp]` MLP kernel gradient, i.e. **embed +minor-most**, the transposed orientation. Whether that costs or pays depends entirely on what +consumes the gradient, and the sign is not the same for all models. + +The evidence is a single opcode. Device-op totals for `copy`, d16 geometry (emb 2048, mlp 8192), +3 profiled steps summed over all eight core planes, one rep: + +| model (d16) | `auto` | `explicit` | `explicit` + dense flag | +| ------------------------- | -------------------: | -----------------------: | ----------------------: | +| llama2 (scan stack 16) | 1128 ops / 27,694 µs | 1896 ops / **26,234 µs** | 1128 ops / 27,682 µs | +| mistral (16) | 1116 / 27,850 | 1884 / **26,298** | 1116 / 27,849 | +| qwen3 (16) | 1116 / 31,180 | 1884 / **29,556** | 1116 / 31,771 | +| gemma3-L18 (scan stack 3) | 792 / **2,425** | 2304 / 10,110 | 792 / 2,391 | + +Two things to read off that table. + +**The flag is a become-`auto` switch for this opcode, exactly.** Turning +`dense_weight_grad_in_kernel_order` on reproduces `auto`'s copy count *op for op* on all four models +— 1128 vs 1128 on llama2, 792 vs 792 on gemma3 — and its cost to within 1.4% (llama2 −0.04%, +mistral −0.01%, gemma3 −1.4%, qwen3 +1.9%). It is not "removing a redundant copy"; it is declining +XLA's alternative layout. + +**On the deep-scan models XLA's alternative layout is the better one.** llama2/mistral/qwen3 relayout +the MLP kernel gradient in *both* modes — the target is the optimizer's `f32[1,512,8192]{2,0,1:T(1,128)}` +— and the only difference is the source: + +``` +auto f32[1,512,8192]{2,0,1:T(1,128)} copy(f32[1,512,8192]{2,1,0:T(8,128)} %get-tuple-element) + 384 ops, 10,154.5 us (26.4 us each) +explicit f32[1,512,8192]{2,0,1:T(1,128)} copy(f32[1,512,8192]{1,2,0:T(8,128)} %convert_bitcast_fusion) + 384 ops, 7,769.8 us (20.2 us each) +``` + +Same count, same destination, **24% cheaper per copy** from the transposed source. Explicit nets +−1,460 µs on llama2's `copy` opcode even after paying for two extra attention-kernel families +(+739 µs on `f32[1,512,8,128]`, +174 µs on `f32[1,8,128,512]`). Forcing kernel order gives that back. + +**On the short-scan models `auto` does no relayout at all, so explicit's choice is pure cost.** +gemma3-L18 (`scan_length = 18 // 6 = 3`) and deepseek's length-1 dense-layer scan have no +`f32[1,512,8192]`-class copy under `auto` — its entire `copy` budget is 2,425 µs of unrelated +attention-mask and bf16-cast traffic. Explicit adds **1,512 copies costing +7,685 µs**, in four +families, all of them the same "gradient came out transposed" signature: + +``` ++4,434.6 us 432 ops f32[1,512,8192]{2,1,0} copy(f32[1,512,8192]{1,2,0} %convert_bitcast_fusion) MLP wi ++2,254.7 us 216 ops f32[1,8192,512]{2,1,0} copy(f32[1,8192,512]{1,2,0} %get-tuple-element) MLP wo ++ 809.0 us 648 ops f32[1,512,8,128]{3,2,0,1} copy(f32[1,512,8,128]{1,3,2,0} ...) q/k/v ++ 215.4 us 216 ops f32[1,8,128,512]{3,2,0,1} copy(f32[1,8,128,512]{2,3,1,0} ...) out proj +``` + +Those four families are 60% of gemma3-L18's whole explicit penalty once the aggregate `while` row +(which double-counts its own body) is excluded: +7,685 of +12,832 µs. The flag deletes all four and +the `while` row flips from +8,897 µs to −2,705 µs. + +**So the discriminator is the scanned stack depth, and it is a proxy, not a mechanism.** The six +models where the flag loses all have a 16-deep gradient stack; the two where it wins have stacks of +3 and 1. The plausible reading is that a 16-deep stack forces a relayout anyway (the optimizer wants +`T(1,128)` on a tensor that the scan writes with `T(8,128)`), so the gradient's own layout is free +to be chosen for the dot; at depth 3 or 1 no relayout is needed and any choice other than the +parameter's is wasted. But that is inference from layouts, not a proven causal chain, and two points +below the threshold against six above is not enough to fit one. **That is why the flag ships +default-off and the LM-head flag does not** — the LM-head transpose is a missed fold with one sign, +this is a layout trade with two. + ______________________________________________________________________ ## 5. Results @@ -762,6 +870,10 @@ copies\] by the per-op decomposition in §4.4. (§4.7) — a different code path from the model at any realistic depth. Re-run at 12 layers it is **−0.58%**. Read the matrix as 9 models plus one methodology finding. +> **This table is superseded by §5.4** and is kept only because §4 refers to its dumps. Every row is +> below XLA's minor-most AR→RS gate (§4.2) and none has the LM-head fix, so it measures the worst +> case of a mode that no longer exists. Quote §5.4 instead. + > **Removed from this table:** the "Δ collective bytes" and "collective count" columns that appeared > in the first revision of this document. Both are artifacts — see §3 caveats 1 and 2. Explicit's > *lower* collective count reflects auto's reduce-scatters being counted as all-reduces inside their @@ -914,6 +1026,101 @@ Extrapolating to a real llama2-7B (emb 4096, 32 layers, seq 4096, 4-way FSDP): t around 659 µs on a step that scales to roughly 2.5 s, i.e. **~0.03%**. That extrapolation is one order of magnitude beyond anything measured here and should be treated as an estimate, not a result. +### 5.4 The table to quote: all eight models, above the gate, with nothing written out but `shard_mode` + +Everything above this point measures a mode that no longer exists — emb 512 (below XLA's AR→RS +gate, §4.2/§5.3), and with `lm_head_weight_grad_in_kernel_order` off. This is the same sweep re-run +at **emb 2048 / mlp 8192 / 16 layers** (gemma3 at 18, a multiple of its 6-long attention pattern; +deepseek at 8), seq 1024, pdbs 1, fsdp 4, on one build, 54 runs, 3 reps per arm, medians. The +`explicit` column has **nothing in the config but `shard_mode: explicit`** — it is what a user gets. + +| model | tied | `auto` ns | `explicit` ns | Δ | rep spread A/B | + dense flag | +| -------------------- | :--: | ----------: | ------------: | ----------: | -------------: | -----------: | +| `mistral-7b` | no | 69,990,530 | 69,604,164 | **−0.552%** | 0.01% / 0.04% | — | +| `qwen3-8b` | no | 83,941,679 | 83,490,576 | **−0.537%** | 0.08% / 0.13% | — | +| `llama2-7b` | no | 70,094,537 | 69,748,502 | **−0.494%** | 0.03% / 0.02% | — | +| `gemma2-2b` | yes | 171,538,409 | 170,719,766 | **−0.477%** | 0.05% / 0.04% | — | +| `gemma-2b` | yes | 88,434,470 | 88,272,771 | **−0.183%** | 0.01% / 0.01% | — | +| `mixtral-8x7b` | no | 353,006,821 | 353,151,437 | +0.041% | 0.03% / 0.07% | — | +| `deepseek3-16b` (L8) | no | 230,885,239 | 231,292,724 | +0.176% | 0.02% / 0.02% | **+0.021%** | +| `gemma3-4b` (L18) | yes | 88,235,432 | 88,944,184 | +0.803% | 0.28% / 0.05% | **−0.101%** | + +**Median −0.330%. Explicit is faster on five of eight and inside the rep spread on a sixth.** The +two models where it is not — deepseek and gemma3 — are exactly the two that want +`dense_weight_grad_in_kernel_order` (§4.8), and with that flag set the worst regression across all +eight models is **mixtral's +0.041%, which is inside mixtral's own 0.07% rep spread**. + +Compare against where this started. The same eight models under `explicit` before any of the fixes +in this document: qwen3 **+3.25%**, gemma3-L16 **+2.02%**, deepseek **+0.92%**. The change is +almost entirely the LM-head default (§8 item 01) plus, on deepseek, the YaRN fix (§5.5). + +Four things to keep in mind before quoting this table: + +1. **It is one geometry and one mesh.** Four chips, fsdp 4, seq 1024, pdbs 1. It is above the AR→RS + gate, which is the point, but it is not a production model and not a production topology. +2. **Rows under ~0.2% are ties, not wins.** The rep spreads are in the table for exactly that + reason; mixtral's +0.041% and deepseek's +0.021% are reported as parity. +3. **gemma3's depth must be a multiple of 6.** At L16 the same sweep reads **+2.020%** for + `explicit` and +0.017% with the dense flag, because 16 // 6 = 2 leaves four layers unrolled + alongside the scan (§4.7). L18 is the honest number. +4. **The dense flag is not a general recommendation.** On the six models not shown in its column it + costs 0.06–0.60% (§4.8, §8 item 02). + +### 5.5 A mode-independent 1.4%: the YaRN frequency table was rebuilt inside every layer + +Not a `shard_mode` finding — it was found while looking for one, it helps both modes equally, and it +is the largest single win in this document. + +`YarnRotaryEmbedding.freqs_cis` was a `@property` that materialized the whole +`[max_position_embeddings, head_dim/2]` complex table and then indexed the `max_target_length` rows +the step actually reads. The table is *traced*, not a constant, so XLA cannot hoist or fold it: it is +rebuilt on every call site. On deepseek that is `f32[163840, 32]` — 5.24 M rows-worth of `cos`/`sin` +per build, of which 1024 rows (0.6%) are used — and the call site is inside the scanned layer. + +Per 3-step profile on `deepseek3-16b`, d16 geometry, `auto` arm: + +``` + (f32[163840,32], f32[163840,32]) fusion(f32[32] %multiply_add_fusion), kind=kLoop + 96 ops 35,842.7 us <- building the table + f32[1024,32] fusion(f32[163840,32] %gte, s32[1024] %broadcast_clamp_fusion), kind=kCustom + 168 ops 854.8 us <- reading 0.6% of it +``` + +Both drop to **zero** ops post-fix. With the buffer traffic they pulled along the whole-profile +device-op total falls 6,386,779 → 6,250,836 µs (**−2.13%**), of which −63,315 µs is `copy-start` and +−33,664 µs is the `while` aggregate. + +The fix (`embeddings.py`) is to compute the rows directly from `position` rather than build a table +to index — row `p` is `exp(1j · p · corrected_freqs)`, so `freqs_cis_at(position)` is one +`einsum("bs,h->bsh")` and a `cis`. The full table is kept as a member for reference and is what the +unit test compares against, bit for bit. + +A second, smaller part of the same fix: writing `exp(1j·θ)` as `complex(cos θ, sin θ)` rather than +`jnp.exp(1j·θ)`. They are bit-identical, but the latter goes through XLA's overflow-safe complex +`exponential` expansion, which evaluates `exp(real(1j·θ)) == exp(0)` over the whole tensor — and +under explicit it evaluates it **twice**, because the `Sharding` custom-call on the broadcast `1j` +stops the simplifier commuting the constant to the right, so the reassociation that lets CSE merge +the two `exponential`s never fires. This is the one part of the finding that is mode-specific. + +Measured (deepseek d16 L8, median of 3–8 reps, per arm, same build before and after): + +| arm | before | after | Δ | +| ------------------------- | -------------: | -------------: | ------: | +| `auto` | 234,120,034 ns | 230,891,079 ns | −1.379% | +| `explicit` | 236,217,057 | 233,015,668 | −1.355% | +| `explicit` + LM-head flag | 234,508,220 | 231,270,774 | −1.381% | +| `explicit` + both flags | 234,227,659 | 230,918,121 | −1.413% | + +**The win is the same size in every arm**, which is the check that it is not a sharding effect. Its +one consequence for this document is on deepseek's *margin*: `explicit` + both flags moves from ++0.046% to **+0.012%** against `auto`, because the fix removes slightly more from explicit than from +auto. + +Scope: this reaches every model with a `YarnRotaryEmbedding` and a `max_position_embeddings` much +larger than `max_target_length` — deepseek2/3 (163,840) most of all. `LlamaVisionRotaryEmbedding` +was deliberately left alone: it consumes the whole table by design, and `llama4` is not on the +explicit whitelist anyway. + ______________________________________________________________________ ## 6. What explicit sharding buys you @@ -989,7 +1196,36 @@ L6, +378.7 at L12. Caveat: 20 of auto's 25 and 8 of explicit's 13 all-reduces ne profile's `XLA Ops` line at all; the omission is symmetric, but the ledger only accounts for the 5 timed collectives in each run. -### 6.3 Errors instead of silent pessimization +### 6.3 The rotated reduce-scatter: explicit's largest and most stable win + +The single biggest opcode-level difference in explicit's favour at d16 geometry is one that no +earlier revision of this document noticed, because the shrink was too small to produce it. + +Under `auto`, the MLP weight-gradient reduce-scatter lands on a scatter dimension that is not +shard-aligned, and XLA implements the fixup as a **ring of `collective-permute`s** — 480 instances +of + +``` +(bf16[24,8192], bf16[24,8192], u32[], u32[]) collective-permute-start(bf16[24,8192] %slice), + source_target_pairs={{0,1},{1,2},{2,3}} +``` + +per 3-step profile. Explicit emits **zero to fifteen** of them: + +| model (d16) | `auto` count / span | `explicit` count / span | +| ----------- | ------------------: | ----------------------: | +| mistral | 495 / 29,360 µs | 15 / 753 µs | +| qwen3 | 480 / 14,571 µs | 0 / 0 µs | +| llama2 | 495 / 13,397 µs | 15 / 122 µs | +| gemma3-L18 | 540 / 15,165 µs | 270 / 8,636 µs | + +Read the spans with §3 caveat 1 in mind — `collective-permute-start` is an async span and overlaps +compute, so this is not 29 ms off mistral's clock. The **count** is the hard fact: explicit deletes +480 of 495 of these ops, in every untied model measured, at every rep. It is the largest term on +explicit's side of the ledger and the most stable, and it is the reason mistral and qwen3 come out +ahead once the LM-head transpose is out of the way. + +### 6.4 Errors instead of silent pessimization A mismatched ZeRO-1 layout is a hard type error under explicit; under auto GSPMD "papers over it with an extra collective" (`src/maxtext/configs/types.py:4548-4560`, the repo's own words). Explicit turns @@ -1062,27 +1298,31 @@ ______________________________________________________________________ ### For users choosing a mode today -- **Tied-embedding models (`logits_via_embedding: true`):** explicit is already as fast or faster — - including gemma3, once it is deep enough to scan (−0.58% at 12 layers). Nothing to wait for. -- **gemma3 specifically:** never benchmark it below 6 layers. `scan_length = num_layers // 6`, so a - 4-layer proxy unrolls the whole decoder and measures a different program (§4.7). -- **Untied dense models at production depth/sequence/batch:** the penalty is a few tenths of a - percent. Choose on capability. If you need `check_vma`, `fused_qkv`, or context parallelism on - non-TPU attention, you must use `auto`. -- **Wide models (emb ≥ 4096):** do *not* assume scale hides the cost — it is the configuration with - the largest absolute penalty measured. Profile before committing. +- **Most models, most geometries: just set `shard_mode: explicit` and write nothing else.** The + LM-head flag now defaults on where it can act, and with it explicit is faster on five of the eight + onboarded models and inside the rep spread on a sixth (§5.4). Choose on capability, not on this + table. +- **gemma3 and deepseek: also set `dense_weight_grad_in_kernel_order: true`.** These are the two + models where explicit still loses by default (+0.803% and +0.176%), and the flag takes gemma3 to + −0.101% and deepseek to +0.021%. It is inert under `auto` and does not touch the checkpoint. Do + not set it on anything else without measuring — it costs 0.06–0.60% on the other six (§4.8). +- **gemma3 specifically:** never benchmark it below 6 layers, and prefer a multiple of 6. + `scan_length = num_layers // 6`, so a 4-layer proxy unrolls the whole decoder and measures a + different program (§4.7), and any remainder is unrolled alongside the scan. +- **Wide models (emb ≥ 4096):** the +1.49% at emb 4096 in §5.3 predates the LM-head flag and has not + been re-measured with it on. Profile before committing. - **MoE with gradient accumulation and `data > 1`:** explicit is the only mode that can express deferred expert-weight all-reduce. This is the case explicit exists for — though note that path is unexercised by anything measured here, so its benefit is a code-level claim, not a measurement. -- **Pure data parallelism:** prefer `auto`, or set `lm_head_weight_grad_in_kernel_order: true` if you - are committed to explicit — that closes the whole +9.3% gap (§8 item 1). +10% is not noise. -- **Any untied model on explicit:** `lm_head_weight_grad_in_kernel_order: true` is checkpoint-safe - and inert under `auto`, so the only reason not to set it is that you have measured it neutral for - your geometry and prefer fewer knobs. +- **Pure data parallelism:** the +10% in §5.2 predates the flag, and the flag's mechanism is exactly + the one that dominated there (§8 item 1 measured the same fix closing the whole +9.3% gap). It + should now be a non-issue, but it has not been re-measured; verify rather than assume. +- **If you need `check_vma`, `fused_qkv`, or context parallelism on non-TPU attention:** you must + use `auto`. These are hard blocks, not slowdowns (§7.1). - **Pure tensor parallelism:** explicit is mildly better (−0.2%), for the reason in §5.2. The advantage may not survive at a real embedding width. -- **Small models / debugging runs:** explicit costs a few percent; usually irrelevant, but do not - benchmark a 4-layer proxy and extrapolate. +- **Small models / debugging runs:** do not benchmark a 4-layer proxy and extrapolate — at emb 512 + the penalty is a benchmark artifact in both directions (§5.3). ### For MaxText, ranked by value @@ -1100,8 +1340,8 @@ ______________________________________________________________________ transpose back into the dot; under explicit that fold never happens. A hand-written `custom_vjp` contracts `dk` straight into the kernel's own axis order instead, so there is no transpose to fold. `_dot_general_kernel_ordered_grad` in `linears.py` implements it, `DenseGeneral`'s - `weight_grad_in_kernel_order` gates it, and `lm_head_weight_grad_in_kernel_order: false` in - `base.yml` wires it to the untied head. It self-disables under `auto` (where the fold does happen) + `weight_grad_in_kernel_order` gates it, and `lm_head_weight_grad_in_kernel_order` in `base.yml` + wires it to the untied head. It self-disables under `auto` (where the fold does happen) and under quantization (where it does not own the dot). **Why this and not the transposed kernel.** They reach the same layout, but only the `custom_vjp` @@ -1170,7 +1410,75 @@ ______________________________________________________________________ Note that raw `reduce-scatter` *counts* are a poor acceptance test on their own — the causal control in §4.2 values that mechanism at ≤28% of the penalty. -02. **Do not bother eliding redundant `reshard` barriers — tested, worth exactly zero.** The obvious + **Update (2026-09-03): the flag is now on by default where it can act.** The table above is + seven configs at two geometries; the case for a default needed the flag to be safe on *every* + model, and it is: + + | model (d16, above gate) | `explicit` | `explicit` + flag | + | ----------------------- | ---------: | ----------------: | + | qwen3-8b | +3.252% | **−0.405%** | + | deepseek3-16b (L8) | +0.920% | **+0.164%** | + | llama2-7b | −0.061% | **−0.468%** | + | mistral-7b | −0.186% | **−0.548%** | + | mixtral-8x7b | −1.393% | +0.009% | + + Plus two controls on mixtral, the one model where the flag does not help: at **L12 / L14 / L16** + it reads +0.120% / −0.151% / −0.038%, oscillating around zero rather than trending; and + shrinking the vocab **8× (32000 → 4096)** leaves the spread unchanged (`explicit` −1.580% vs + −1.441%, with-flag −0.081% at both). Mixtral's −1.393% is therefore an MoE-scan scheduling + windfall that has nothing to do with the LM head's size, and the flag costs it nothing — it + lands on `auto`'s schedule rather than losing to it. **Across every model, depth and vocab + measured, the flag has never been worse than `auto`.** So it resolves to on for an untied model + under `shard_mode: explicit` and off everywhere else; `base.yml` writes `None` and + `resolve_lm_head_weight_grad_in_kernel_order` in `types.py` does the resolution, rejecting an + explicit `true` on a tied head rather than silently ignoring it. + +02. **Ship `dense_weight_grad_in_kernel_order` — but leave it off, and say why (§4.8).** The same + `custom_vjp` wired to every other `DenseGeneral`: the three `MlpBlock` projections + (`linears.py:650,670,689`), the four attention projections (`attentions.py:700,737,779,835`) and + the nine MLA / Indexer projections (`attention_mla.py:141,157,177,814,830,852,868,893`). + + Unlike the LM-head flag this one is **not** a missed optimization with one sign. Under explicit + XLA is free to choose the dense weight gradient's layout, because the `Sharding` custom-calls + decouple it from the stored parameter's; the flag declines that freedom and reproduces `auto`'s + layout op-for-op. Measured at d16, on top of the LM-head default: + + | model | default | + dense flag | | + | ---------- | ------: | -----------: | ---------- | + | gemma3-L18 | +0.690% | **−0.267%** | **use it** | + | deepseek | +0.164% | **−0.003%** | **use it** | + | mixtral | +0.009% | +0.065% | leave off | + | gemma2-2b | −0.518% | −0.327% | leave off | + | gemma-2b | −0.213% | +0.165% | leave off | + | llama2 | −0.468% | +0.013% | leave off | + | mistral | −0.548% | −0.049% | leave off | + | qwen3 | −0.405% | +0.192% | leave off | + + (That is the one build carrying both arms for all eight models, so it is quoted whole rather + than mixed with §5.4's. The later, cleaner build of §5.4 reproduces both recommendations at the + same magnitude: gemma3 +0.803% → **−0.101%**, deepseek +0.176% → **+0.021%**.) + + The discriminator is the scanned stack depth — the two winners have gradient stacks of 3 and 1, + the six losers all have 16 — but that is a proxy read off layouts, not a proven mechanism, and + two points below the line against six above is not enough to fit a threshold on. Encoding it as + a default would be fitting an XLA layout heuristic. **Leave it off, document the two models, and + revisit if a third data point ever lands between 3 and 16.** + +03. **Do not build a lookup table you index once (§5.5) — and do not expect an A/B to find it.** + `YarnRotaryEmbedding.freqs_cis` materialized `f32[163840, 32]` inside every layer of every step + to read 1024 rows out of it; computing the rows straight from `position` is **−1.4% of step + time in both modes**. It survived a full mode-comparison sweep precisely because both modes paid + it equally. The detector that found it is cheap and worth keeping: flag any op whose output + element count is ≥ 8× its largest input's. Run over all eight models' profiles it reproduced + this hit at exactly the reported magnitude and found **nothing else above ~0.06%** — the + `timescale` `@property` family, the splash-attention masks (numpy at trace time), + `generate_attention_mask`'s iota masks (`dot_product` path only, not selected on TPU), + `PositionalEmbedding._compute_embeddings` (once per step, outside the scan), the MoE one-hots, + `DeepSeekV4RotaryEmbedding.inv_freq` and all of `src/maxtext/utils/` are clean. The residual + `broadcast(f32[] 0)` at 0.8–1.2% of step is JAX's `lax.scan` `ys` output buffers, not a MaxText + bug. + +04. **Do not bother eliding redundant `reshard` barriers — tested, worth exactly zero.** The obvious companion fix to item 1 is to stop emitting a barrier when the reshard is a no-op (the value's aval already carries the requested sharding), on the theory that fewer barriers means more folds. It was implemented and measured, and it does not work: @@ -1188,39 +1496,66 @@ ______________________________________________________________________ elided. Recorded here so the next person does not spend the week: **the barrier count is not the lever, the gradient orientation is.** -03. **Add `shard_mode` to the golden suites (§7.3).** Until CI compares modes, every regression here +05. **Add `shard_mode` to the golden suites (§7.3).** Until CI compares modes, every regression here can silently return. `tests/integration/hlo_diff_test.py` is the cheapest place to start; assert on `after_codegen` reduce-scatter counts, not on static collective byte totals. -04. **Any HLO byte-counter used to compare modes must descend into +06. **Any HLO byte-counter used to compare modes must descend into `kind=kCustom, calls=%all-reduce-scatter.*` fusions** and report the post-slice shape. Counting the `all-reduce` inside the fusion at its pre-slice shape reports auto's wins as if they were losses, which is how the "identical bytes" claim survived a full sweep here. -05. **Re-examine the `check_vma` ban (§7.1).** JAX 0.11.1 accepts `check_vma=True` on an all-Explicit +07. **Re-examine the `check_vma` ban (§7.1).** JAX 0.11.1 accepts `check_vma=True` on an all-Explicit mesh; the restriction appears to be MaxText-imposed. Lifting it would give the MoE models on the explicit whitelist back a documented optimization. -06. **Validate `fused_qkv` + explicit and context-parallel + explicit** in `types.py` so users get a +08. **Validate `fused_qkv` + explicit and context-parallel + explicit** in `types.py` so users get a MaxText message instead of a raw JAX trace error. -07. **Fix `gcs_utils.upload_dump()`** to no-op (or write locally) for non-`gs://` paths — today it +09. **Fix `gcs_utils.upload_dump()`** to no-op (or write locally) for non-`gs://` paths — today it crashes after training and destroys the profile. -08. **Make explicit mode not defeat remat CSE on unrolled blocks (§4.7).** With `remat_policy: full` +10. **Make explicit mode not defeat remat CSE on unrolled blocks (§4.7).** With `remat_policy: full` and no scan, auto CSEs the recomputed forward against the original (61 remat-tagged instructions of 11,209); explicit cannot (1,013 of 12,452), costing 279.5 µs. This bites any config where a block is unrolled — gemma3 below 6 layers today, but also `scan_layers: false` runs generally. -09. **Deduplicate the tied-embedding bf16 cast (§4.6).** `embeddings.py:235-240`'s `.T` makes explicit +11. **Deduplicate the tied-embedding bf16 cast (§4.6).** `embeddings.py:235-240`'s `.T` makes explicit materialize the 128 MB embedding shard in bf16 **twice** (+80–82 µs on every gemma). Hoisting the cast above the transpose, or reusing one converted value for both consumers, is a pure win and is independent of everything else in this document. -10. **Investigate the tuple-repack all-reduce tax (§6.2)** — the one measured explicit cost that grows +12. **Investigate the tuple-repack all-reduce tax (§6.2)** — the one measured explicit cost that grows with depth (+112 → +219 → +379 µs from 4 to 12 layers on gemma3). -11. **Delete or rewire `validate_shard_mode`** (§7.4) — as dead code it actively misinforms. +13. **Delete or rewire `validate_shard_mode`** (§7.4) — as dead code it actively misinforms. + +### Levers that were tried and do not work + +Recorded so the next person does not spend the time. Each was implemented or measured, not merely +reasoned about, and each survived an adversarial attempt to rescue it. + +| candidate | verdict | +| --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| Stop pinning the logits output under explicit | 0-line HLO diff. The barriers are JAX's, not MaxText's (§4.5). | +| Elide no-op `reshard` barriers | Op-identical executable on all three configs (item 04). | +| `lm_head_kernel_transposed` (store the kernel `[vocab, embed]`) | Works, but slower on 4 of 7 configs, no checkpoint path, and re-initializes the model. Removed (item 01). | +| Reshard the RMSNorm scale under explicit outside TP | Implemented and reverted — negative on every non-TP config. | +| Sliced/prefetched MoE expert weights via MSA | No measurable effect; the expert all-gathers are already overlapped. | +| Remove the rotated-reduce-scatter `collective-permute` fixup | Backwards: it is an `auto`-mode op that explicit already deletes (§6.3). There was nothing on the explicit side to remove, and it is not caused by `_permuted_sharding`. | +| `param_scan_axis: 1` + XLA tiling as the dense-flag discriminator | Does not separate the winners from the losers; every model in the sweep uses the same value. | +| 3-D kernels / MLA low-rank / gemma3's attention pattern / contracting-axis-vs-FSDP, as the dense-flag discriminator | None separate the sign. gemma3-L18's win is carried by plain 2-D `[embed, mlp]` `MlpBlock` kernels — the same sites that lose on llama2 — and deepseek's by its length-1 dense-layer scan (4,640 µs of `f32[512,1,8192]` copies) plus 500 µs on the MoE shared MLP, not by MLA (MLA's 4-D gradients move ~170 µs). | +| Rematerializing RoPE sin/cos in the backward pass | Mode-neutral: identical to 0.1 µs in both arms. | +| `LlamaVisionRotaryEmbedding`'s full-table build | Consumes the whole table by design, and `llama4` is not on the explicit whitelist — unreachable. | +| Deduplicating the out-projection all-gather | The "extra" all-gather is the same op scheduled differently; no bytes to remove. | +| The gemma2 padded-reduce-scatter `concatenate` | Sign flips by rep; within noise. | +| deepseek's two remaining duplicated dense-layer weight all-gathers (`bf16[2048,1,8,192]` 371 µs / 12 ops, `bf16[2048,1,576]` 189 µs / 12 ops, both `%copy-done` duplicating an existing `%convert_element_type` all-gather) | Real, explicit-only, and worth **0.009%**. Not worth chasing. | + +A related null result worth stating positively: a systematic hunt for *any* remaining explicit-only +cost on llama2 / mistral / qwen3 at d16 found **nothing above ~0.10% of step**. The largest positive +is llama2's scan gradient-stack `dynamic-update-slice` at +68 µs/step (+0.098%), and it has an +identical signature and op count (960 / 960) in both arms with a sign that flips by model — i.e. it +is scheduling around a shared op, not an explicit-mode defect. ______________________________________________________________________ @@ -1278,3 +1613,26 @@ checkpointing off, throughout. - **3 profiled steps per run** for most configs (4 runs per mode only for the d16 point). Adequate given the 0.07% noise floor, but not a large sample — and the module-span lead gap (§3 caveat 3) puts a floor of a few tenths of a percent on any single-pair comparison. + +Added with the 2026-09-03 revision: + +- **§5.4 is one geometry.** emb 2048 / mlp 8192 / 16 layers / seq 1024 / pdbs 1 / fsdp 4 on four + chips. It is above the AR→RS gate, which is what makes it the right table to quote, but it is not + a production model and it is not a production mesh. Nothing here was re-measured at emb 4096, + seq ≥ 4096, pure DP, pure TP or multi-host with the new default on — the §5.2 and §5.3 rows are + all pre-flag and should be read as upper bounds. +- **3 reps per arm, medians compared.** Several §5.4 deltas (mixtral +0.009%, deepseek's +0.012% + over 8 reps) are inside the noise floor and are reported as ties, not wins. Only the ≥ 0.2% rows + should be read as directional. +- **The dense flag's discriminator is fitted to eight points**, with two on one side of the + threshold and six on the other, and it is a layout proxy rather than a proven mechanism (§4.8). + It is documented as a per-model recommendation for exactly the two models it was measured on. +- **The YaRN finding was measured on deepseek only.** It is the only model in the sweep whose + `max_position_embeddings` (163,840) dwarfs `max_target_length` by enough for the table build to + dominate; the fix is unconditional, so other YaRN models get whatever their ratio is worth, which + was not measured. +- **`deepseek3-671b-batchsplit` newly resolves the LM-head flag to `true`** and was verified by + inspection only — the flag's permutation is the identity for a 2-D `[embed_vocab, vocab]` kernel, + and the `reduced=`/`unreduced=` tags live on decoder-layer weight gradients inside + `deepseek_batchsplit.py`'s `shard_map`s, never on the LM head. That config does not fit on the + hardware used here, so it has not been run. From 5b97257025763dfc642a58618245093ce2f35fff Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Thu, 3 Sep 2026 21:57:26 +0000 Subject: [PATCH 7/9] Implement the kernel-order flags with jax.sharding.auto_axes The hand-written custom_vjp reordered the weight-gradient dot's operands itself and permuted the result back, to work around the sharding barrier JAX puts on every primitive result under an explicit mesh. Wrapping the forward dot in an auto_axes region instead removes the barrier for the length of that one dot, so Shardy re-propagates the operands' shardings exactly as shard_mode: auto would and XLA performs the fold on its own. Same mechanism, 114 fewer lines, and gradients are now bit-identical to the default rule rather than agreeing to rounding. Step times agree with the custom_vjp to within 0.15 pp on all five models measured both ways. Outside a mesh there is no barrier to remove, so eval_shape reports no sharding and the region is skipped. --- src/maxtext/layers/linears.py | 165 ++++++++-------------------------- tests/unit/linears_test.py | 163 +++++++++++++++++++-------------- 2 files changed, 133 insertions(+), 195 deletions(-) diff --git a/src/maxtext/layers/linears.py b/src/maxtext/layers/linears.py index d78880afaf..afdb5a1dae 100644 --- a/src/maxtext/layers/linears.py +++ b/src/maxtext/layers/linears.py @@ -85,118 +85,35 @@ def _compute_dot_general(inputs, kernel, kernel_axes, axis, contract_ind, matmul return dot_general(inputs, kernel, ((axis, contract_ind), ((), ())), precision=matmul_precision) -def _aval_sharding(value): - """The sharding carried by `value`'s type, or None when it carries none. - - Under `shard_mode: explicit` every aval knows its own sharding, so this is how - a hand-written backward pass pins its outputs to the same layout the default - transpose rule would have produced. - """ - try: - sharding = jax.typeof(value).sharding - except Exception: # pylint: disable=broad-except - return None - # A fully replicated aval still needs its `out_sharding` spelled out: the - # gradient dots below contract over sharded token axes, which JAX refuses to - # resolve on its own. - return sharding if getattr(sharding, "spec", None) is not None else None - - -def _permuted_sharding(value, produced_axes): - """`value`'s sharding with its axes reordered into `produced_axes` order. - - The gradient dots emit the operand's axes in `produced_axes` order and only - then permute them back, so the `out_sharding` handed to the dot has to be - permuted the same way. For every layout this path is enabled for - `produced_axes` is the identity and this is just `_aval_sharding`. +def _dot_general_in_auto_axes(inputs, kernel, dimension_numbers, precision, out_sharding): + """`lax.dot_general` traced inside an auto-sharding region. + + Under `shard_mode: explicit` JAX wraps every primitive result in a `Sharding` + custom-call, and XLA's algebraic simplifier will not fold + `transpose(dot(A, B))` back into one dot across it. The transpose autodiff + leaves on the weight gradient therefore survives to codegen, and the gradient + has to be relaid out before it can be written back to the parameter -- + worth +0.4% to +4.9% of step time depending on the model + (docs/guides/optimization/shard_mode_performance.md section 4.8). + + `jax.sharding.auto_axes` removes the barrier for the length of the dot: + inside the region the operands' shardings leave the type system and Shardy + re-propagates them, which is exactly what `shard_mode: auto` does, so XLA + folds the transpose back into the gradient dot and the gradient comes out in + the kernel's stored order. The rest of the model keeps its shardings in its + types. No arithmetic changes, so the gradients are bit-identical to the + default rule. """ - sharding = _aval_sharding(value) - produced_axes = tuple(produced_axes) - if sharding is None or produced_axes == tuple(range(len(produced_axes))): - return sharding - spec = sharding.spec - return sharding.update(spec=jax.sharding.PartitionSpec(*(spec[a] for a in produced_axes))) - - -def _restore_axis_order(value, produced_axes): - """Permute `value` so that its axis j is the operand's original axis j. - - `produced_axes[i]` names the original axis that dot_general placed at position - i. For every layout this path is enabled for the mapping is already the - identity, so the transpose is elided rather than merely folded later. - """ - produced_axes = tuple(produced_axes) - if produced_axes == tuple(range(len(produced_axes))): - return value - return lax.transpose(value, tuple(produced_axes.index(j) for j in range(len(produced_axes)))) - - -@functools.partial(jax.custom_vjp, nondiff_argnums=(2, 3, 4)) -def _dot_general_kernel_ordered_grad(inputs, kernel, dimension_numbers, precision, out_sharding): - """`lax.dot_general` whose weight gradient is emitted in stored-kernel layout. - - JAX's default transpose rule builds the kernel gradient as - `transpose(dot(g, inputs))`. XLA normally folds that transpose back into the - dot, but under `shard_mode: explicit` every dot output carries a `Sharding` - custom-call and algebraic simplification will not fold across it. The - transpose then survives to codegen, which leaves the gradient sharded on its - minor-most dimension -- the one position where XLA declines to fuse the - weight-gradient all-reduce into a reduce-scatter. See - docs/guides/optimization/shard_mode_performance.md for the measurements. - - Writing the backward pass by hand avoids the transpose instead of relying on - XLA to undo it: `dk` is contracted straight into the kernel's own axis order. - That reassociates the sum, so the gradients match the default rule to float - rounding (~1e-7 relative) rather than bit for bit. - """ - return lax.dot_general(inputs, kernel, dimension_numbers, precision=precision, out_sharding=out_sharding) - - -def _dot_general_kernel_ordered_grad_fwd(inputs, kernel, dimension_numbers, precision, out_sharding): - out = lax.dot_general(inputs, kernel, dimension_numbers, precision=precision, out_sharding=out_sharding) - return out, (inputs, kernel) - - -def _dot_general_kernel_ordered_grad_bwd(dimension_numbers, precision, out_sharding, res, g): - """Backward pass for `_dot_general_kernel_ordered_grad`.""" - del out_sharding # Constrains the forward output only. - inputs, kernel = res - (lhs_contract, rhs_contract), (lhs_batch, rhs_batch) = dimension_numbers - if lhs_batch or rhs_batch: - raise ValueError("_dot_general_kernel_ordered_grad does not support batch dimensions.") - lhs_free = tuple(i for i in range(inputs.ndim) if i not in lhs_contract) - rhs_free = tuple(i for i in range(kernel.ndim) if i not in rhs_contract) - n_lhs_free = len(lhs_free) - - # dk contracts the free input axes against g's leading axes, which correspond - # to them one for one. The result carries the kernel's contracting axes (in - # ascending operand order) followed by its free axes -- the kernel's own order - # whenever it is stored in_features-major. - dk_axes = tuple(rhs_contract[lhs_contract.index(a)] for a in sorted(lhs_contract)) + rhs_free - dk = lax.dot_general( - inputs, - g, - ((lhs_free, tuple(range(n_lhs_free))), ((), ())), - precision=precision, - out_sharding=_permuted_sharding(kernel, dk_axes), - ) - dk = _restore_axis_order(dk, dk_axes) - - # dx contracts g's trailing axes against the kernel's free axes, likewise - # paired one for one. - dx_axes = lhs_free + tuple(lhs_contract[rhs_contract.index(a)] for a in sorted(rhs_contract)) - dx = lax.dot_general( - g, - kernel, - ((tuple(range(n_lhs_free, g.ndim)), rhs_free), ((), ())), - precision=precision, - out_sharding=_permuted_sharding(inputs, dx_axes), - ) - dx = _restore_axis_order(dx, dx_axes) - return dx, dk - - -_dot_general_kernel_ordered_grad.defvjp(_dot_general_kernel_ordered_grad_fwd, _dot_general_kernel_ordered_grad_bwd) + dot = functools.partial(lax.dot_general, dimension_numbers=dimension_numbers, precision=precision) + if out_sharding is None: + # `auto_axes` has to be told what the region produces. Ask JAX for the + # sharding the same dot would have carried without the region -- outside a + # mesh that is `None`, and there is no barrier to remove in the first place. + out_sharding = jax.eval_shape(dot, inputs, kernel).sharding + if out_sharding is None: + return dot(inputs, kernel) + # pylint: disable-next=too-many-function-args (pylint reads the decorator overload of `auto_axes`) + return jax.sharding.auto_axes(dot, out_sharding=out_sharding)(inputs, kernel) def _compute_dot_general_nnx( @@ -223,9 +140,7 @@ def _compute_dot_general_nnx( out_sharding = truncate_out_sharding(out_sharding, out_ndim) if weight_grad_in_kernel_order: - return _dot_general_kernel_ordered_grad( - inputs, kernel, ((axis, contract_ind), ((), ())), matmul_precision, out_sharding - ) + return _dot_general_in_auto_axes(inputs, kernel, ((axis, contract_ind), ((), ())), matmul_precision, out_sharding) return dot_general( inputs, kernel, ((axis, contract_ind), ((), ())), precision=matmul_precision, out_sharding=out_sharding @@ -279,16 +194,14 @@ def __init__( transpose XLA emits for a single combined 2-axis all-gather. debug_sharding: when True, log the logical/physical sharding of the two-stage all-gather constraints to the sharding dump files. - weight_grad_in_kernel_order: compute the weight gradient with a hand-written - backward pass that contracts straight into the kernel's stored axis order, - rather than letting autodiff emit `transpose(dot(...))` and relying on XLA - to fold the transpose away. See + weight_grad_in_kernel_order: trace the dot inside a `jax.sharding.auto_axes` + region so that XLA folds `transpose(dot(...))` back into the gradient dot + and the weight gradient comes out in the kernel's stored axis order. See docs/guides/optimization/shard_mode_performance.md section 4.3 for why - that fold does not happen under `shard_mode: explicit`. The stored kernel - and its initialization are untouched, so no checkpoint conversion is - needed; the gradients differ from the default rule only by floating-point - reassociation. Only takes effect under `shard_mode: explicit`; ignored - when the layer is quantized. + that fold does not happen under `shard_mode: explicit` otherwise. The + stored kernel, its initialization and the arithmetic are all untouched, + so gradients are bit-identical to the default rule. Only takes effect + under `shard_mode: explicit`; ignored when the layer is quantized. rngs: RNG state for initialization in nnx. """ self.in_features_shape = canonicalize_tuple(in_features_shape) @@ -444,9 +357,9 @@ def __call__( contract_ind = tuple(range(0, len(self.axis))) quant_dot_general = self.quant_dot_general if slice_bounds is None else None - # The hand-written weight gradient only pays off where XLA refuses to fold - # the transpose away, and it has to own the whole dot to do it -- so it is - # off under auto sharding and under any quantized path. + # The auto-axes region only buys anything where XLA refuses to fold the + # transpose away, and it has to own the whole dot to do it -- so it is off + # under auto sharding and under any quantized path. weight_grad_in_kernel_order = ( self.weight_grad_in_kernel_order and self.shard_mode == ShardMode.EXPLICIT diff --git a/tests/unit/linears_test.py b/tests/unit/linears_test.py index c720b9a3a5..0700194b01 100644 --- a/tests/unit/linears_test.py +++ b/tests/unit/linears_test.py @@ -15,6 +15,7 @@ """Tests for linears.py.""" import functools +import re import sys import unittest from unittest.mock import MagicMock, patch @@ -182,6 +183,16 @@ def test_slice_bounds_invalid(self): with self.assertRaisesRegex(ValueError, "slice_bounds .* must be valid and within"): layer(inputs, slice_bounds=(0, 100)) + @staticmethod + def _explicit_mesh(): + """A mesh whose axes carry their shardings in the type system. + + The flag only does anything under one of these -- that is where JAX puts a + sharding barrier between the weight-gradient dot and its transpose -- so + every test of the mechanism has to run inside one. + """ + return jax.make_mesh((1, 1), ("fsdp", "tensor"), axis_types=(jax.sharding.AxisType.Explicit,) * 2) + def _make_kernel_ordered_pair(self, in_features_shape, out_features_shape, axis=-1, use_bias=False): """Builds a default layer and a weight_grad_in_kernel_order twin sharing weights. @@ -211,70 +222,72 @@ def test_weight_grad_in_kernel_order_forward_matches(self): cases = ((4, 8, -1), (4, (2, 8), -1), ((2, 3), 8, (1, 2))) for in_features_shape, out_features_shape, axis in cases: with self.subTest(in_features_shape=in_features_shape, out_features_shape=out_features_shape): - default, ordered = self._make_kernel_ordered_pair(in_features_shape, out_features_shape, axis=axis) - shape = (2, 3) if axis == -1 else (5,) - shape += (in_features_shape,) if isinstance(in_features_shape, int) else tuple(in_features_shape) - inputs = jax.random.normal(jax.random.PRNGKey(0), shape) - np.testing.assert_array_equal(default(inputs), ordered(inputs)) + with jax.set_mesh(self._explicit_mesh()): + default, ordered = self._make_kernel_ordered_pair(in_features_shape, out_features_shape, axis=axis) + shape = (2, 3) if axis == -1 else (5,) + shape += (in_features_shape,) if isinstance(in_features_shape, int) else tuple(in_features_shape) + inputs = jax.random.normal(jax.random.PRNGKey(0), shape) + np.testing.assert_array_equal(default(inputs), ordered(inputs)) def test_weight_grad_in_kernel_order_gradients_match(self): - """Reordering the contraction must not change the gradient beyond float noise.""" + """The region moves no arithmetic, so every gradient must come out bit-identical.""" cases = ((4, 8, -1), (4, (2, 8), -1), ((2, 3), 8, (1, 2))) for in_features_shape, out_features_shape, axis in cases: with self.subTest(in_features_shape=in_features_shape, out_features_shape=out_features_shape): - default, ordered = self._make_kernel_ordered_pair(in_features_shape, out_features_shape, axis=axis) - shape = (2, 3) if axis == -1 else (5,) - shape += (in_features_shape,) if isinstance(in_features_shape, int) else tuple(in_features_shape) - inputs = jax.random.normal(jax.random.PRNGKey(0), shape) + with jax.set_mesh(self._explicit_mesh()): + default, ordered = self._make_kernel_ordered_pair(in_features_shape, out_features_shape, axis=axis) + shape = (2, 3) if axis == -1 else (5,) + shape += (in_features_shape,) if isinstance(in_features_shape, int) else tuple(in_features_shape) + inputs = jax.random.normal(jax.random.PRNGKey(0), shape) - def loss(layer, x): - return jnp.sum(jnp.sin(layer(x))) + def loss(layer, x): + return jnp.sum(jnp.sin(layer(x))) - grad_default = nnx.grad(loss)(default, inputs).kernel[...] - grad_ordered = nnx.grad(loss)(ordered, inputs).kernel[...] - self.assertEqual(grad_ordered.shape, ordered.kernel[...].shape) - np.testing.assert_allclose(grad_default, grad_ordered, rtol=1e-5, atol=1e-6) + grad_default = nnx.grad(loss)(default, inputs).kernel[...] + grad_ordered = nnx.grad(loss)(ordered, inputs).kernel[...] + self.assertEqual(grad_ordered.shape, ordered.kernel[...].shape) + np.testing.assert_array_equal(grad_default, grad_ordered) - grad_x_default = jax.grad(functools.partial(loss, default))(inputs) - grad_x_ordered = jax.grad(functools.partial(loss, ordered))(inputs) - np.testing.assert_allclose(grad_x_default, grad_x_ordered, rtol=1e-5, atol=1e-6) + grad_x_default = jax.grad(functools.partial(loss, default))(inputs) + grad_x_ordered = jax.grad(functools.partial(loss, ordered))(inputs) + np.testing.assert_array_equal(grad_x_default, grad_x_ordered) def test_weight_grad_in_kernel_order_lm_head_shape(self): - """The LM head's 2-D kernel is the only shape the config flag enables. + """The LM head's 2-D kernel is the shape the config flag enables by default. - Reordering the contraction reassociates the sum, so the gradient agrees to - float rounding rather than exactly. The stored kernel -- and therefore the - initialization -- is untouched, so the difference stays at rounding level - instead of changing the training trajectory. + Bit-identical here matters more than anywhere else: the head's gradient + feeds the largest parameter in most models, so a reassociated sum would show + up in the training trajectory rather than staying at rounding level. """ for dtype in (jnp.float32, jnp.bfloat16): with self.subTest(dtype=jnp.dtype(dtype).name): - default, ordered = self._make_kernel_ordered_pair(8, 32) - inputs = jax.random.normal(jax.random.PRNGKey(0), (2, 5, 8), dtype) - weights = jax.random.normal(jax.random.PRNGKey(1), (2, 5, 32), dtype) + with jax.set_mesh(self._explicit_mesh()): + default, ordered = self._make_kernel_ordered_pair(8, 32) + inputs = jax.random.normal(jax.random.PRNGKey(0), (2, 5, 8), dtype) + weights = jax.random.normal(jax.random.PRNGKey(1), (2, 5, 32), dtype) - def loss(layer, x, weights=weights): - return jnp.sum((layer(x) * weights).astype(jnp.float32)) + def loss(layer, x, weights=weights): + return jnp.sum((layer(x) * weights).astype(jnp.float32)) - grad_default = nnx.grad(loss)(default, inputs).kernel[...] - grad_ordered = nnx.grad(loss)(ordered, inputs).kernel[...] - scale = float(jnp.abs(grad_default).max()) - self.assertLess(float(jnp.abs(grad_default - grad_ordered).max()), 1e-5 * scale) + grad_default = nnx.grad(loss)(default, inputs).kernel[...] + grad_ordered = nnx.grad(loss)(ordered, inputs).kernel[...] + np.testing.assert_array_equal(grad_default, grad_ordered) def test_weight_grad_in_kernel_order_composes_with_remat(self): # The LM head sits outside the scanned stack, but training still wraps it in - # remat, and a custom_vjp that breaks under checkpointing is unusable. - default, ordered = self._make_kernel_ordered_pair(4, 8) - inputs = jax.random.normal(jax.random.PRNGKey(0), (2, 3, 4)) + # remat, and the region is re-traced on the recomputed forward pass. + with jax.set_mesh(self._explicit_mesh()): + default, ordered = self._make_kernel_ordered_pair(4, 8) + inputs = jax.random.normal(jax.random.PRNGKey(0), (2, 3, 4)) - def grads(layer): - return nnx.grad(nnx.remat(lambda l, x: jnp.sum(jnp.sin(l(x)))))(layer, inputs).kernel[...] + def grads(layer): + return nnx.grad(nnx.remat(lambda l, x: jnp.sum(jnp.sin(l(x)))))(layer, inputs).kernel[...] - np.testing.assert_allclose(grads(default), grads(ordered), rtol=1e-5, atol=1e-6) + np.testing.assert_array_equal(grads(default), grads(ordered)) def test_weight_grad_in_kernel_order_ignored_under_auto(self): - # Under auto sharding XLA folds the transpose itself, so the custom_vjp buys - # nothing and must not take over the dot. + # Under auto sharding there is no barrier for the region to remove -- XLA + # folds the transpose itself -- so the dot must be left alone. layer = linears.DenseGeneral( in_features_shape=4, out_features_shape=8, @@ -286,11 +299,13 @@ def test_weight_grad_in_kernel_order_ignored_under_auto(self): reference = linears.DenseGeneral( in_features_shape=4, out_features_shape=8, shard_mode=ShardMode.AUTO, rngs=nnx.Rngs(params=0) ) - np.testing.assert_array_equal(layer(inputs), reference(inputs)) + with patch.object(linears, "_dot_general_in_auto_axes") as spy: + np.testing.assert_array_equal(layer(inputs), reference(inputs)) + spy.assert_not_called() def test_weight_grad_in_kernel_order_yields_to_quantization(self): - # The custom_vjp has to own the whole dot, so it must step aside for a - # quantized one rather than wrap it. + # The region has to own the whole dot, so it must step aside for a quantized + # one rather than wrap it. mock_quant = MagicMock() mock_quant.quant_mode = None mock_instance = mock_quant.dot_general_cls.return_value.return_value @@ -306,35 +321,45 @@ def test_weight_grad_in_kernel_order_yields_to_quantization(self): rngs=self.rngs, ) inputs = jax.random.normal(jax.random.PRNGKey(0), (2, 3, 4)) - with patch.object(linears, "_dot_general_kernel_ordered_grad") as spy: + with patch.object(linears, "_dot_general_in_auto_axes") as spy: layer(inputs) spy.assert_not_called() - def test_weight_grad_in_kernel_order_emits_no_transpose(self): - """The mechanism itself: no transpose on the weight-gradient path.""" - default, ordered = self._make_kernel_ordered_pair(4, 8) - inputs = jax.random.normal(jax.random.PRNGKey(0), (2, 3, 4)) + def test_weight_grad_in_kernel_order_drops_the_sharding_barrier(self): + """The mechanism itself: nothing between the weight-gradient dot and its transpose. + + Autodiff still writes the transpose either way. What the flag changes is + whether a `sdy.sharding_constraint` sits in between -- XLA's algebraic + simplifier will not fold `transpose(dot(A, B))` back into a single dot + across one, so the transpose survives to codegen and the gradient has to be + relaid out before it can be written back to the parameter. + """ + with jax.set_mesh(self._explicit_mesh()): + default, ordered = self._make_kernel_ordered_pair(4, 8) + inputs = jax.random.normal(jax.random.PRNGKey(0), (2, 3, 4)) - def count_transposes(layer): - def loss(l, x): - return jnp.sum(jnp.sin(l(x))) + def transpose_operand_op(layer): + """The op producing the value the weight-gradient transpose consumes.""" - text = nnx.jit(nnx.grad(loss)).lower(layer, inputs).as_text() - return text.count("stablehlo.transpose") + def loss(l, x): + return jnp.sum(jnp.sin(l(x))) - self.assertEqual(count_transposes(ordered), 0) - self.assertGreater(count_transposes(default), 0) + text = nnx.jit(nnx.grad(loss)).lower(layer, inputs).as_text() + defining_op = dict(re.findall(r"(%\w+) = ([\w.]+)", text)) + transposed = re.findall(r"= stablehlo\.transpose (%\w+)", text) + self.assertEqual(len(transposed), 1) + return defining_op[transposed[0]] - @staticmethod - def _explicit_mesh(): - return jax.make_mesh((1, 1), ("fsdp", "tensor"), axis_types=(jax.sharding.AxisType.Explicit,) * 2) + self.assertEqual(transpose_operand_op(ordered), "stablehlo.dot_general") + self.assertEqual(transpose_operand_op(default), "sdy.sharding_constraint") def test_weight_grad_in_kernel_order_with_replicated_kernel(self): - """A replicated kernel still has to spell out the gradient's sharding. + """A replicated kernel still has to spell out what the region produces. - The gradient dot contracts over the sharded token axes, which JAX refuses to - resolve on its own -- omitting `out_sharding` there raises instead of - defaulting to replicated. + `auto_axes` has no default for `out_sharding`, and the shape it is handed + here is the awkward one: the forward dot's output is sharded over a batch + axis the kernel knows nothing about, so the answer has to come from the + operands rather than from the kernel's own sharding. """ mesh = self._explicit_mesh() with jax.set_mesh(mesh): @@ -352,7 +377,7 @@ def loss(x, k, dot): return jnp.sum(jnp.sin(dot(x, k))) ordered = functools.partial( - linears._dot_general_kernel_ordered_grad, # pylint: disable=protected-access + linears._dot_general_in_auto_axes, # pylint: disable=protected-access dimension_numbers=dimension_numbers, precision=None, out_sharding=None, @@ -367,10 +392,10 @@ def loss(x, k, dot): def test_weight_grad_in_kernel_order_with_permuted_axes(self): """Dimension numbers that do not pair operand axes in ascending order. - No `DenseGeneral` layout builds these today, but the backward pass permutes - its dots' outputs back into operand order and the `out_sharding` it hands - those dots has to be permuted the same way or the custom_vjp returns a - cotangent whose type does not match its primal. + No `DenseGeneral` layout builds these today, but they are the case where the + region's operands come back out permuted, so they check that leaving and + re-entering the type system lands every cotangent on the sharding its primal + had. """ mesh = self._explicit_mesh() with jax.set_mesh(mesh): @@ -392,7 +417,7 @@ def loss(x, k, dot): return jnp.sum(jnp.sin(dot(x, k))) ordered = functools.partial( - linears._dot_general_kernel_ordered_grad, # pylint: disable=protected-access + linears._dot_general_in_auto_axes, # pylint: disable=protected-access dimension_numbers=dimension_numbers, precision=None, out_sharding=out_sharding, From 6a54684a705ea735abc352300dc9697b4414c03f Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Thu, 3 Sep 2026 21:57:37 +0000 Subject: [PATCH 8/9] Turn dense_weight_grad_in_kernel_order on under explicit sharding The flag shipped default-off because its sign looked model-dependent: worth ~0.9 pp on gemma3 and ~0.16 pp on deepseek, and a 0.06-0.60% cost on the other six models. Every one of those eight points was measured at 16 layers, which is a scan depth where an unrelated XLA layout pathology inflates both sharding modes by ~25% and swamps the effect being read. Re-measured at depths where that pathology does not fire, the flag has one sign everywhere and the two flags are additive. Across 15 (model, depth) pairs the LM-head flag alone leaves a median +0.463% penalty against auto, worst +1.038%; with this one on as well the median is +0.005% and the full range is -0.183% to +0.134%. On the eight-model matrix at healthy depth, explicit with nothing written out but shard_mode now reads median +0.000% against auto, worst +0.133%, and wins outright on gemma-2b (-0.685%) and gemma3-4b (-0.374%). Also rewrites both flags' descriptions around the auto_axes region that now implements them. --- src/maxtext/configs/base.yml | 35 +++++++++++++-------------- src/maxtext/configs/types.py | 47 +++++++++++++++++++++++------------- tests/unit/pyconfig_test.py | 18 ++++++++++++++ 3 files changed, 65 insertions(+), 35 deletions(-) diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 82dd903898..b9829499fa 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -200,14 +200,14 @@ logits_via_embedding: false normalize_embedding_logits: true # whether to normalize pre-softmax logits if logits_via_embedding is true logits_dot_in_fp32: false # whether to use fp32 in logits_dense or shared_embedding dot product for stability cast_logits_to_fp32: true # whether to cast the logits to fp32. the higher precision is generally beneficial, but it can vary slightly. -# Compute the untied LM head's weight gradient with a hand-written backward pass that contracts straight into the -# kernel's stored orientation, instead of letting autodiff emit transpose(dot(...)) and hoping XLA folds it away. -# Under shard_mode: explicit that fold never happens, because the Sharding custom-call on every dot output is a -# barrier algsimp will not cross, so the gradient is left sharded on its minor-most dimension and loses the -# all-reduce -> reduce-scatter fusion (see docs/guides/optimization/shard_mode_performance.md). The stored kernel -# and its initialization are untouched, so no checkpoint conversion is needed; gradients differ from the default -# rule only by floating-point reassociation. No effect under shard_mode: auto, where XLA folds the transpose itself, -# and none on a tied head, which has no kernel of its own -- so None means on exactly where it can do something. +# Trace the untied LM head's dot inside a jax.sharding.auto_axes region so that XLA folds the transpose back into +# the gradient dot and the weight gradient comes out in the kernel's stored orientation. Without the region that +# fold never happens under shard_mode: explicit, because the Sharding custom-call on every dot output is a barrier +# algsimp will not cross, so the gradient is left sharded on its minor-most dimension and loses the all-reduce -> +# reduce-scatter fusion (see docs/guides/optimization/shard_mode_performance.md). The stored kernel, its +# initialization and the arithmetic are untouched, so gradients are bit-identical to the default rule and no +# checkpoint conversion is needed. No effect under shard_mode: auto, where XLA folds the transpose itself, and none +# on a tied head, which has no kernel of its own -- so None means on exactly where it can do something. lm_head_weight_grad_in_kernel_order: None # on for an untied model under shard_mode: explicit if None float32_qk_product: false # in dot_product attention, whether to cast to fp32 the inputs to qk product float32_logits: false # in dot_product attention, whether to cast to fp32 the inputs to softmax @@ -719,16 +719,15 @@ shard_optimizer_over_data: false # when dense MLP weight matrices are sharded on both fsdp and fsdp-transpose axes, use two separate all-gather calls dense_fsdp_use_two_stage_all_gather: false -# The in-loop counterpart of lm_head_weight_grad_in_kernel_order: compute the attention q/k/v/out and MLP wi/wo -# weight gradients with a hand-written backward pass that contracts straight into each kernel's stored orientation. -# Unlike the LM head flag this is a trade, not a fix, so it defaults off. Under shard_mode: explicit XLA is free to -# pick the gradient's layout -- the Sharding custom-calls decouple it from the stored parameter's -- and this flag -# declines that freedom, reproducing shard_mode: auto's layout op for op. Measured, that is worth roughly -0.3% on -# gemma3 and -0.2% on deepseek, whose scanned gradient stacks are 3 and 1 deep, and costs 0.06%-0.60% on the models -# whose stacks are 16 deep (see docs/guides/optimization/shard_mode_performance.md section 4.8). Turn it on for -# gemma3 and deepseek; measure before turning it on for anything else. Stored kernels and their initialization are -# untouched, so no checkpoint conversion is needed. No effect under shard_mode: auto or on quantized layers. -dense_weight_grad_in_kernel_order: false +# The in-loop counterpart of lm_head_weight_grad_in_kernel_order: trace the attention q/k/v/out and MLP wi/wo dots +# inside a jax.sharding.auto_axes region so that XLA folds the transpose back into the gradient dot and each weight +# gradient comes out in its kernel's stored orientation. Same barrier as the LM head flag, same default rule -- None +# means on wherever it can act, i.e. under shard_mode: explicit. Without it explicit is 0.33%-1.48% slower than auto +# at every measured layer count that is not a multiple of 8; with it explicit lands within 0.2% of auto everywhere, +# and the two flags are additive (see docs/guides/optimization/shard_mode_performance.md section 4.8). Stored kernels, +# their initialization and the arithmetic are untouched, so gradients are bit-identical to the default rule and no +# checkpoint conversion is needed. No effect under shard_mode: auto or on quantized layers. +dense_weight_grad_in_kernel_order: None # on under shard_mode: explicit if None # Unless explicitly specified, the number of TPU slices is automatically determined. It should only be set for # disaggregated reinforcement learning workloads using multiple slices. For ahead of time compilation, diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 39a532df3c..ce4aac3dbb 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -607,11 +607,11 @@ class LogitsAndLoss(BaseModel): lm_head_weight_grad_in_kernel_order: bool | None = Field( None, description=( - "Compute the untied LM head's weight gradient directly in the kernel's stored axis order via a " - "custom_vjp, so autodiff emits no transpose after the gradient dot. Recovers the weight-gradient " - "reduce-scatter that shard_mode=explicit otherwise loses, leaving the stored kernel and its " - "initialization untouched. No effect under shard_mode=auto. None means on for an untied model " - "under shard_mode=explicit and off everywhere else." + "Trace the untied LM head's dot inside a jax.sharding.auto_axes region so that XLA folds the " + "transpose back into the gradient dot and the weight gradient comes out in the kernel's stored " + "axis order. Recovers the weight-gradient reduce-scatter that shard_mode=explicit otherwise loses, " + "leaving the stored kernel, its initialization and the arithmetic untouched. No effect under " + "shard_mode=auto. None means on for an untied model under shard_mode=explicit and off everywhere else." ), ) final_logits_soft_cap: None | NonNegativeFloat = Field( @@ -1268,19 +1268,18 @@ class LayoutAndSharding(BaseModel): False, description="Use two separate All-Gather calls for dense MLP weights sharded on both FSDP and FSDP-transpose.", ) - dense_weight_grad_in_kernel_order: bool = Field( - False, + dense_weight_grad_in_kernel_order: bool | None = Field( + None, description=( - "Compute the per-layer attention and MLP projection weight gradients directly in each kernel's stored " - "axis order via a custom_vjp, so autodiff emits no transpose after the gradient dot. This is the " - "in-loop counterpart of lm_head_weight_grad_in_kernel_order, but unlike that flag it is a trade rather " - "than a fix, so it defaults off. Under shard_mode=explicit the Sharding custom-calls decouple the " - "gradient's layout from the stored parameter's, leaving XLA free to pick it; this flag declines that " - "freedom and reproduces shard_mode=auto's layout op for op. Measured, that is worth roughly -0.3% on " - "gemma3 and -0.2% on deepseek, whose scanned gradient stacks are 3 and 1 deep, and costs 0.06%-0.60% on " - "the models whose stacks are 16 deep (docs/guides/optimization/shard_mode_performance.md section 4.8). " - "Measure before turning it on for anything else. Stored kernels and their initialization are untouched. " - "No effect under shard_mode=auto or on quantized layers." + "Trace the per-layer attention and MLP projection dots inside a jax.sharding.auto_axes region so that " + "XLA folds the transpose back into the gradient dot and each weight gradient comes out in its kernel's " + "stored axis order. This is the in-loop counterpart of lm_head_weight_grad_in_kernel_order and, like it, " + "None means on wherever it can act: shard_mode=explicit, where the Sharding custom-call on every dot " + "output blocks that fold. Without it explicit is 0.33%-1.48% slower than auto at every measured layer " + "count that is not a multiple of 8; with it explicit lands within 0.2% of auto everywhere " + "(docs/guides/optimization/shard_mode_performance.md section 4.8). Stored kernels, their initialization " + "and the arithmetic are untouched, so gradients are bit-identical to the default rule. No effect under " + "shard_mode=auto, where XLA folds the transpose itself, or on quantized layers." ), ) internal_compile: bool = Field( @@ -3313,6 +3312,20 @@ def resolve_lm_head_weight_grad_in_kernel_order(self) -> "MaxTextConfig": ) return self + @model_validator(mode="after") + def resolve_dense_weight_grad_in_kernel_order(self) -> "MaxTextConfig": + """Resolve the in-loop counterpart of the LM-head flag. + + Same barrier, same shape of fix, same default rule: on wherever it can act. + The two are independent and additive -- on llama2 at 18 layers, explicit + costs +1.098% against auto with neither flag, +0.621% with this one alone, + +0.429% with the LM-head flag alone and +0.005% with both + (docs/guides/optimization/shard_mode_performance.md section 4.8). + """ + if self.dense_weight_grad_in_kernel_order is None: + self.dense_weight_grad_in_kernel_order = self.shard_mode == ShardMode.EXPLICIT + return self + @model_validator(mode="after") def set_derived_and_validate_values(self) -> "MaxTextConfig": """ diff --git a/tests/unit/pyconfig_test.py b/tests/unit/pyconfig_test.py index 9b42f52736..1fb729433b 100644 --- a/tests/unit/pyconfig_test.py +++ b/tests/unit/pyconfig_test.py @@ -278,6 +278,24 @@ def resolve(**kwargs): with self.assertRaisesRegex(Exception, "only applies to the untied LM head"): resolve(shard_mode="explicit", logits_via_embedding=True, lm_head_weight_grad_in_kernel_order=True) + def test_dense_weight_grad_in_kernel_order_default(self): + """The in-loop counterpart follows the same rule, minus the tied-head carve-out.""" + + def resolve(**kwargs): + return pyconfig.initialize( + [os.path.join(MAXTEXT_PKG_DIR, "train.py"), get_test_config_path()], + skip_jax_distributed_system=True, + **kwargs, + ).dense_weight_grad_in_kernel_order + + self.assertTrue(resolve(shard_mode="explicit")) + self.assertFalse(resolve(shard_mode="auto")) + # Unlike the LM head, every model has these projections, tied or not. + self.assertTrue(resolve(shard_mode="explicit", logits_via_embedding=True)) + # Writing it out still wins over the default, in both directions. + self.assertFalse(resolve(shard_mode="explicit", dense_weight_grad_in_kernel_order=False)) + self.assertTrue(resolve(shard_mode="auto", dense_weight_grad_in_kernel_order=True)) + def test_resolve_config_path(self): self.assertEqual(resolve_config_path("foo"), os.path.join("src", "foo")) self.assertEqual(resolve_config_path(__file__), __file__) From 169aac7724283e82d979be8621d6dbd192889caa Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Thu, 3 Sep 2026 21:57:50 +0000 Subject: [PATCH 9/9] Correct the shard_mode guide for the scan-depth layout pathology Walking llama2 from 2 to 24 layers, ns-per-layer is flat at ~3.3 M except at stack lengths 2, 4, 8, 16 and 24, where it jumps 25-34%. At those lengths XLA assigns layout {2,0,1} to the scanned parameter stack, so each dynamic-update-slice writes a degenerate T(1,128)-tiled slice: 8.16 ms of a 36.5 ms step at L8, identical in both sharding modes. param_scan_axis: 0 recovers it (-23.6% at L8, -24.9% at L16) and it survives at production geometry. Every benchmark in the guide before this revision used a stack length of 4, 8 or 16, so several "explicit wins" rows are noise inside a 25% inflation. The originals are kept and marked; new section 4.9 characterizes the pathology and new section 5.6 replaces section 5.4 with the eight-model matrix re-measured on the shipped defaults at depths where it does not fire. --- .../optimization/shard_mode_performance.md | 457 ++++++++++++++++-- 1 file changed, 419 insertions(+), 38 deletions(-) diff --git a/docs/guides/optimization/shard_mode_performance.md b/docs/guides/optimization/shard_mode_performance.md index 7ecf65e2c6..19a15a1949 100644 --- a/docs/guides/optimization/shard_mode_performance.md +++ b/docs/guides/optimization/shard_mode_performance.md @@ -24,6 +24,13 @@ because two of the corrections reverse a conclusion that looked solid at the tim Revised 2026-09-03: the whole matrix re-measured at above-gate geometry (§5.4), the LM-head flag turned **on by default** where it can act, a second flag added for the dense-layer gradients (§4.8), and a mode-independent 1.4% win found and shipped along the way (§5.5). +Revised again 2026-09-03 after a depth sweep: **every "explicit wins" row in this document sits at a +scanned-stack depth where an unrelated XLA layout pathology inflates both modes by ~25% (§4.9)**. +With those depths excluded explicit needs *both* kernel-order flags to reach parity, so +`dense_weight_grad_in_kernel_order` now also defaults on under explicit (§4.8), and both flags are +now implemented with `jax.sharding.auto_axes` rather than a hand-written `custom_vjp` (§8 item 01). +The eight-model matrix was then re-run on that build at healthy depths: **§5.6 supersedes §5.4 and +is the table to quote.** `shard_mode` (`src/maxtext/configs/base.yml:554`) selects how MaxText expresses sharding: @@ -142,6 +149,15 @@ ______________________________________________________________________ > (L12/L14/L16), at any vocab (4096 or 32000) — while explicit without it is +3.25% on qwen3. > It is therefore now tri-state and resolves to **on** wherever it can act: `shard_mode: explicit` and an untied head. Writing it out is only needed to reproduce a measurement. + > *Revised again (2026-09-03).* The mechanism is unchanged; the implementation is not. The + > hand-written `custom_vjp` has been replaced by a `jax.sharding.auto_axes` region around the + > forward dot, which drops the barrier for the length of that one dot and lets Shardy + > re-propagate the operands' shardings exactly as `shard_mode: auto` would — so XLA performs the + > fold itself instead of MaxText hand-writing around it. Same step times to within 0.15 pp on all + > five models measured both ways, 114 fewer lines, gradients now **bit-identical** to the default + > rule rather than agreeing to rounding, and it works at all 17 call sites including MLA + > (§8 item 01). + 11. **With that default, explicit is faster than `auto` on five of eight onboarded models and inside the run-to-run spread on a sixth.** The one number a user needs is §5.4: at emb 2048 / mlp 8192 / 16 layers, with nothing in the config but `shard_mode`, the median across the eight models is @@ -151,6 +167,20 @@ ______________________________________________________________________ (§4.8): it is worth 0.90 pp on gemma3 and 0.16 pp on deepseek, and costs 0.06–0.60% on the other six. With those two per-model settings applied, no model regresses by more than 0.041%. + > *Superseded (2026-09-03).* Both halves of this item are contaminated by §4.9. **Sixteen layers + > is one of the depths where the layout pathology fires**, and that pathology is worth ~25% of + > step time — enough that scheduling noise inside it swamps a 0.4% sharding effect and, on five + > of eight models, flips its sign. Re-measuring llama2 at L16 with the pathology removed + > (`param_scan_axis: 0`) turns the "explicit wins by 0.494%" of §5.4 into **explicit loses by + > 1.482%**, and the dense flag takes it to **+0.006%**. Across 15 (model, depth) pairs at + > *healthy* depths the LM-head flag alone leaves a median **+0.463%** penalty (worst +1.038% on + > gemma2) while both flags together give a median **+0.005%** and a full range of + > **−0.183% … +0.134%**. So the second flag is not a two-model carve-out at all — it is the other + > half of the fix, it now defaults on under explicit too, and the honest summary of both is + > *parity with `auto`, reliably*, rather than a win. The eight-model matrix re-run on that build + > at healthy depths is **§5.6**: median **+0.000%**, worst **+0.133%**, and the only two rows + > outside their own rep spread are `gemma-2b` (−0.685%) and `gemma3-4b` (−0.374%), both wins. + 12. **The largest single win found here is not a sharding fix at all.** `YarnRotaryEmbedding` built its whole `[max_position_embeddings, head_dim/2]` frequency table inside every layer of every step in order to read `max_target_length` rows out of it — `f32[163840, 32]` on deepseek, 0.6% @@ -165,12 +195,27 @@ ______________________________________________________________________ > different matrix from the same seed at the flipped shape — perturbs the loss ~35× more than > explicit's own noise. It has been removed rather than shipped alongside a strictly better flag. +13. **The biggest number found in this whole investigation is not a sharding number either: a + scanned stack whose length is 8, 16, 24 (or 2 or 4) costs ~25% of step time, in both modes.** + `param_scan_axis: 1` stacks each parameter as `[in, L, out]`. At those lengths XLA assigns the + gradient stack layout `{2,0,1}` instead of `{2,1,0}`, so the per-iteration + `dynamic-update-slice` writes a degenerate `T(1,128)`-tiled slice: **8.16 ms of a 36.5 ms step** + on llama2 at L8. Walking depth 2→24 the ns/layer trend is flat at ~3.3 M except at exactly those + lengths, where it jumps 25–34%. `param_scan_axis: 0` removes it (**L8 −23.6%, L16 −24.9%**) and + costs 0.18–0.28% at healthy depths. It survives at production geometry (+14.1% at L8, +11.5% at + L16). See §4.9 — and note that **every benchmark in this document before §4.9 used a stack length + of 4, 8 or 16**, so it is the most consequential methodology finding here. + **Practical guidance:** as of this revision, `shard_mode: explicit` with nothing else written out is -the faster mode on most onboarded models at realistic geometry (§5.4), so choose on capability -(§6, §7) and treat performance as a tie-breaker rather than a tax. Two carve-outs remain: **gemma3 -and deepseek want `dense_weight_grad_in_kernel_order: true`** (§4.8), and **pure data parallelism** -was +10% before the LM-head fix — re-verify it for your geometry rather than assuming the fix -covers it, since the DP measurement predates the default flip. +**at parity** with `auto` on all eight onboarded models at realistic geometry and healthy depth — +median **+0.000%**, worst case **+0.133%** against a 0.16% rep spread, and two models (`gemma-2b` +−0.685%, `gemma3-4b` −0.374%) where explicit wins outright (§5.6; §4.8 gives the same picture over +15 (model, depth) pairs, median +0.005%, worst +0.134%). Choose on capability (§6, §7) and treat +performance as a non-issue rather than a tax. Both kernel-order flags now default on where they can +act; nothing needs to be written out. Two things to still verify for yourself: **pure data +parallelism** was +10% before the LM-head fix and has not been re-measured, and if your +`base_num_decoder_layers` is a multiple of 8 you are paying §4.9's ~25% in both modes and should try +`param_scan_axis: 0`. ______________________________________________________________________ @@ -840,6 +885,164 @@ below the threshold against six above is not enough to fit one. **That is why th default-off and the LM-head flag does not** — the LM-head transpose is a missed fold with one sign, this is a layout trade with two. +> *Superseded (2026-09-03) — and the paragraph above contains its own refutation.* "A 16-deep stack +> forces a relayout anyway, because the optimizer wants `T(1,128)` on a tensor the scan writes with +> `T(8,128)`" is not a property of 16-deep stacks. It is §4.9: the degenerate layout XLA picks for a +> stack whose length is a multiple of 8, in **both** sharding modes, at a cost of ~25% of step time. +> Every "the flag loses here" data point sat inside it. Read on. + +#### The resolution: measure at a depth where nothing else is happening + +Re-running the same comparison at depths that are *not* multiples of 8 — where the `dynamic-update-slice` +is a well-tiled `{2,1,0}` write and the step is 25% cheaper — reverses the conclusion. Fifteen +(model, depth) pairs, d16 geometry, 2 reps each, medians, both columns against that config's own +`auto`: + +| model | L | `explicit` (LM-head flag only) | `explicit` + dense flag | +| ---------- | --: | -----------------------------: | ----------------------: | +| llama2 | 12 | +0.622% | −0.016% | +| llama2 | 14 | +0.391% | +0.002% | +| llama2 | 18 | +0.429% | +0.005% | +| llama2 | 20 | +0.463% | +0.005% | +| mistral | 12 | +0.598% | +0.051% | +| mistral | 14 | +0.919% | −0.007% | +| qwen3 | 12 | +0.354% | +0.044% | +| qwen3 | 14 | +0.397% | +0.052% | +| qwen3 | 18 | +0.534% | +0.008% | +| gemma2 | 12 | +0.950% | −0.095% | +| gemma2 | 14 | +1.038% | −0.183% | +| gemma3 | 18 | +0.803% | −0.101% | +| mixtral | 12 | −0.115% | +0.134% | +| mixtral | 14 | −0.053% | +0.052% | +| deepseek | 12 | −0.348% | +0.061% | +| **median** | | **+0.463%** | **+0.005%** | + +**Range: `explicit` alone −0.348% … +1.038%; with the dense flag −0.183% … +0.134%.** The flag is +not a two-model carve-out and it is not a coin flip — it is a *variance eliminator*. It costs at +most 0.19 pp anywhere and it removes a penalty that is otherwise present on 12 of 15 configs. The +one model where it gives something up, mixtral, gives up 0.06–0.25 pp; the six models where the +original table said it "costs 0.06–0.60%" were reading scheduling noise inside the §4.9 pathology. + +The two flags are also **independent and additive**, which is the check that they are two instances +of one mechanism rather than one effect measured twice. At L18, all four combinations, against the +same `auto`: + +| config | llama2 L18 | qwen3 L18 | +| ---------------------------- | ----------: | ----------: | +| `explicit`, neither flag | +1.098% | +5.579% | +| dense flag only | +0.621% | +4.853% | +| LM-head flag only | +0.429% | +0.534% | +| **both flags** (the default) | **+0.005%** | **+0.008%** | + +qwen3 is the clean case: its untied 151,936-entry head dominates, so the LM-head flag is worth 5 pp +and the dense flag 0.7 pp; on llama2 the two are the same size. Neither subsumes the other, and the +sum of the two individual gains predicts the joint one to within 0.03 pp on llama2 and 0.2 pp on +qwen3. **Both flags therefore now resolve to on under `shard_mode: explicit`.** + +One disclosure that does not fit the pattern and is not being papered over: **on mixtral the +LM-head flag forfeits a reproducible ~1.4%.** Explicit with *neither* flag reads −1.436% / −1.448% / +−1.522% at L12 / L14 / L16 against `auto`, and turning the LM-head flag on lands it back on `auto`'s +schedule (−0.115% / −0.053% / −0.012%). That is an MoE-scan scheduling windfall, not an LM-head +effect — shrinking mixtral's vocab 8× leaves the spread unchanged (§8 item 01) — and deepseek, the +other MoE model, has the *opposite* sign (the LM-head flag takes it from +0.947% to +0.176% at L8). +One model in one direction is not a rule to encode in a default, so the default stands and the cost +is stated here. A mixtral user who has measured this may set `lm_head_weight_grad_in_kernel_order: false`. + +______________________________________________________________________ + +### 4.9 Not a sharding effect at all: `param_scan_axis: 1` degenerates when the stack length is a multiple of 8 + +This section exists because it invalidates measurements elsewhere in this document, not because it +has anything to do with `shard_mode`. The effect is **identical in both modes** — it was found by +sweeping `auto` alone. + +Walk llama2 down the depth ladder at d16 geometry (emb 2048, mlp 8192, seq 1024, pdbs 1, fsdp 4), +2 reps per point, and step time is not proportional to depth: + +| L | `auto` ns | ns/layer | `explicit` | `explicit` + dense | +| -----: | ----------: | ------------: | ----------: | -----------------: | +| **2** | 11,007,529 | **5,503,764** | **−0.902%** | −0.050% | +| 3 | 12,494,770 | 4,164,923 | +0.027% | +0.002% | +| **4** | 19,376,215 | **4,844,054** | **−0.517%** | +0.083% | +| 5 | 18,731,819 | 3,746,364 | +0.465% | +0.074% | +| 6 | 21,837,597 | 3,639,600 | +0.343% | +0.037% | +| 7 | 24,953,852 | 3,564,836 | +0.330% | +0.018% | +| **8** | 36,536,639 | **4,567,080** | **−0.444%** | −0.034% | +| 9 | 31,091,571 | 3,454,619 | +0.342% | +0.004% | +| 10 | 34,169,574 | 3,416,957 | +0.409% | +0.004% | +| 12 | 40,345,816 | 3,362,151 | +0.622% | −0.016% | +| 14 | 46,583,332 | 3,327,381 | +0.391% | +0.002% | +| **16** | 70,094,537 | **4,380,909** | **−0.494%** | — | +| 18 | 58,852,896 | 3,269,605 | +0.429% | +0.005% | +| 20 | 65,004,700 | 3,250,235 | +0.463% | +0.005% | +| **24** | 104,650,808 | **4,360,450** | **−0.542%** | −0.024% | + +Two patterns, and they are the same pattern. The ns/layer column decays smoothly from 4.16 M at L3 +to 3.25 M at L20 — **except** at L2, L4, L8, L16 and L24, where it is 25–34% above the trend. And +those five depths are **exactly** the five where `explicit` "beats" `auto`. Everywhere else explicit +is +0.33% to +0.62%, monotonically, with a rep spread under 0.1%. + +**The HLO says what it is.** Diffing the L8 and L9 `auto` dumps (`opdiff.py`), the scanned gradient +stack changes layout: + +``` +L9 f32[9,512,8192]{2,1,0:T(8,128)} <- normal, row-major, 8x128 tiles +L8 f32[8,512,8192]{2,0,1:T(8,128)} <- middle dimension made minor +``` + +and with it every per-iteration write into that stack: + +``` +L8 dynamic-update-slice(f32[8,512,8192]{2,0,1}, f32[1,512,8192]{2,0,1:T(1,128)S(1)} %copy, ...) + three families, 34,480 + 31,730 + 31,676 us = 97,886 us / profile + = 8.16 ms of a 36.5 ms step = 22% of the step +``` + +`T(1,128)` is a degenerate tile — one row per tile instead of eight — so the write moves 8× the +minimum number of tiles. The same three families are present, at the same cost, in the `explicit` +and `explicit`+flags dumps: this is XLA's layout assignment reacting to the stack's leading +dimension, and sharding mode has no input to it. + +**`param_scan_axis: 0` removes it.** That config stacks parameters as `[L, in, out]` instead of +`[in, L, out]`: + +| llama2, d16 | `param_scan_axis: 1` | `param_scan_axis: 0` | Δ | +| ----------- | -------------------: | -------------------: | -----------: | +| L8 | 36,535,103 | 27,917,743 | **−23.587%** | +| L9 | 31,091,571 | 31,003,302 | −0.284% | +| L16 | 70,096,329 | 52,622,643 | **−24.928%** | +| L18 | 58,852,896 | 58,747,091 | −0.180% | + +So it is worth ~24% at the bad depths and ~0.2% at the good ones — the two are the same measurement, +because `param_scan_axis: 0`'s L8 and L16 sit on the healthy ns/layer trend that `param_scan_axis: 1` +only reaches at non-multiples of 8. It is **not** an artifact of the shrunk geometry either. At +`prod` (16 query heads, seq 4096) the same jump is there: + +| llama2, prod | step ns | ns/layer | vs the next depth up | +| ------------ | ----------: | --------: | -------------------: | +| L8 | 79,070,204 | 9,883,776 | **+14.07%** | +| L9 | 77,978,783 | 8,664,309 | — | +| L16 | 151,238,509 | 9,452,407 | **+11.45%** | +| L18 | 152,662,296 | 8,481,239 | — | + +**What this costs the rest of this document.** Every headline comparison before this section was run +at L4 (§5.1, §5.2), L8 (deepseek) or L16 (§5.4, §4.8) — all inside the pathology. The clearest +correction is §5.4's llama2 row. Re-run at L16 with the pathology removed: + +| llama2, L16, `param_scan_axis: 0` | step ns | vs `auto` | +| --------------------------------- | ---------: | ----------: | +| `auto` | 52,622,643 | — | +| `explicit` (LM-head flag only) | 53,402,334 | **+1.482%** | +| `explicit` + dense flag | 52,625,548 | **+0.006%** | + +§5.4 reports the same model at the same depth as **−0.494% for explicit**. The sign is opposite. A +0.4% sharding effect measured on top of a 25% layout artifact is measuring the artifact's scheduling +noise, and this is why the flag-default conclusion changed (§4.8). + +This has not been chased further. Whether `param_scan_axis: 0` should become the MaxText default is +a question about checkpoint layout, optimizer-state layout and every model config, not a question +about sharding, and it wants its own investigation (§8 item 03). + ______________________________________________________________________ ## 5. Results @@ -870,9 +1073,9 @@ copies\] by the per-op decomposition in §4.4. (§4.7) — a different code path from the model at any realistic depth. Re-run at 12 layers it is **−0.58%**. Read the matrix as 9 models plus one methodology finding. -> **This table is superseded by §5.4** and is kept only because §4 refers to its dumps. Every row is +> **This table is superseded by §5.6** and is kept only because §4 refers to its dumps. Every row is > below XLA's minor-most AR→RS gate (§4.2) and none has the LM-head fix, so it measures the worst -> case of a mode that no longer exists. Quote §5.4 instead. +> case of a mode that no longer exists. Quote §5.6 instead. > **Removed from this table:** the "Δ collective bytes" and "collective count" columns that appeared > in the first revision of this document. Both are artifacts — see §3 caveats 1 and 2. Explicit's @@ -1066,6 +1269,14 @@ Four things to keep in mind before quoting this table: 4. **The dense flag is not a general recommendation.** On the six models not shown in its column it costs 0.06–0.60% (§4.8, §8 item 02). +> **Superseded (2026-09-03): read this table only for its `auto` column.** Sixteen layers is one of +> the depths where §4.9's layout pathology fires, and it is worth ~25% of step time in both modes. +> Six of these eight rows are at L16 (gemma3 is at L18 and deepseek at L8 — L8 is also inside the +> pathology, so only gemma3's row is clean). Re-running llama2 at L16 with `param_scan_axis: 0` +> gives `explicit` **+1.482%**, not the −0.494% above; the caveat 4 note about the dense flag costing +> 0.06–0.60% is the same artifact. The replacement table, at a healthy depth and with both flags +> defaulted on, is **§5.6**. + ### 5.5 A mode-independent 1.4%: the YaRN frequency table was rebuilt inside every layer Not a `shard_mode` finding — it was found while looking for one, it helps both modes equally, and it @@ -1121,6 +1332,57 @@ larger than `max_target_length` — deepseek2/3 (163,840) most of all. `LlamaVis was deliberately left alone: it consumes the whole table by design, and `llama4` is not on the explicit whitelist anyway. +### 5.6 The table to quote — eight models, shipped defaults, at a depth where §4.9 is not firing + +This replaces §5.4. Same eight models, same geometry (emb 2048 / mlp 8192 / seq 1024 / pdbs 1 / +fsdp 4, four chips), same protocol (3 reps per arm, medians), re-run on the build this document +describes. Two things changed: every model now sits at a scanned-stack length that is **not** a +multiple of 8 (§4.9), and the `explicit` column has **nothing in the config but +`shard_mode: explicit`** — both `lm_head_weight_grad_in_kernel_order` and +`dense_weight_grad_in_kernel_order` resolve to on by themselves (§8 items 01–02). Sorted by margin. + +| model | layers | tied | `auto` ns | `explicit` ns | Δ | rep spread A/B | +| --------------- | -----: | :--: | ----------: | ------------: | ----------: | -------------: | +| `gemma-2b` | 18 | yes | 78,786,920 | 78,247,099 | **−0.685%** | 0.08% / 0.06% | +| `gemma3-4b` | 18 | yes | 88,335,850 | 88,005,221 | **−0.374%** | 0.21% / 0.31% | +| `mixtral-8x7b` | 14 | no | 309,329,429 | 309,088,916 | −0.078% | 0.08% / 0.07% | +| `llama2-7b` | 18 | no | 58,858,748 | 58,857,796 | −0.002% | 0.02% / 0.04% | +| `gemma2-2b` | 18 | yes | 150,702,912 | 150,706,999 | +0.003% | 0.06% / 0.02% | +| `mistral-7b` | 18 | no | 58,432,480 | 58,443,356 | +0.019% | 0.02% / 0.01% | +| `qwen3-8b` | 18 | no | 72,851,689 | 72,914,405 | +0.086% | 0.18% / 0.24% | +| `deepseek2-16b` | 12 | no | 360,257,035 | 360,735,487 | +0.133% | 0.16% / 0.05% | + +**Median +0.000%. Worst case +0.133%, against a 0.16% rep spread on the same row.** Six of the eight +rows are inside their own rep spread in both directions — that is the definition of a tie. The two +that are not are `gemma-2b` and `gemma3-4b`, and `explicit` wins both. + +The three models this revision set out to fix, before and after (before = §5.4, `explicit` with only +the LM-head flag defaulted on, at the depth §5.4 used): + +| model | §5.4 (LM-head flag only) | §5.6 (both flags, healthy depth) | +| --------------- | -----------------------: | -------------------------------: | +| `gemma3-4b` | +0.803% | **−0.374%** | +| `mixtral-8x7b` | +0.041% | **−0.078%** | +| `deepseek2-16b` | +0.176% | **+0.133%** | + +gemma3 is the one where the flag does real work: 1.18 pp, and it changes the sign. mixtral and +deepseek were already inside noise in §5.4 and are still inside noise here; what the dense flag buys +on those two is not a win but the absence of a tail — across the 15 healthy (model, depth) pairs in +§4.8 the worst `explicit` row goes from +1.038% to +0.134% when it is on. + +Four caveats before quoting this: + +1. **mixtral is at 14 layers and deepseek at 12, not 18.** At 18 both run out of HBM at this + geometry — deepseek asks for 43.56 G of HLO temporaries against 31.24 G available — so the sweep + fails outright rather than producing a slow number. 14 and 12 are both healthy stack lengths + under §4.9, which is the property that matters; §4.8 measures the same two models at the same two + depths with all four flag combinations if you want the decomposition. +2. **gemma3's depth must be a multiple of 6** (§4.7): `scan_length = num_layers // 6`, and at any + other depth the remainder is unrolled alongside the scan and dominates the comparison. +3. **It is one geometry and one mesh.** Four chips, fsdp 4, seq 1024, pdbs 1. Above the AR→RS gate + (§4.2/§5.3), which is the point, but not a production model and not a production topology. +4. **Rows under ~0.15% are ties, not results.** The rep spreads are in the table for that reason. + ______________________________________________________________________ ## 6. What explicit sharding buys you @@ -1298,14 +1560,22 @@ ______________________________________________________________________ ### For users choosing a mode today -- **Most models, most geometries: just set `shard_mode: explicit` and write nothing else.** The - LM-head flag now defaults on where it can act, and with it explicit is faster on five of the eight - onboarded models and inside the rep spread on a sixth (§5.4). Choose on capability, not on this - table. -- **gemma3 and deepseek: also set `dense_weight_grad_in_kernel_order: true`.** These are the two - models where explicit still loses by default (+0.803% and +0.176%), and the flag takes gemma3 to - −0.101% and deepseek to +0.021%. It is inert under `auto` and does not touch the checkpoint. Do - not set it on anything else without measuring — it costs 0.06–0.60% on the other six (§4.8). +- **Most models, most geometries: just set `shard_mode: explicit` and write nothing else.** Both + kernel-order flags now default on where they can act, and with them explicit is at parity with + `auto` on every onboarded model at every healthy depth measured: median **+0.000%** and worst case + **+0.133%** across all eight models on the shipped defaults (§5.6), median +0.005% across the 15 + (model, depth) pairs of §4.8. Choose on capability. +- **Do not write `dense_weight_grad_in_kernel_order` out any more.** The "gemma3 and deepseek only" + advice this line used to carry was an artifact of benchmarking at 16 layers (§4.9); the flag is + now on by default under explicit and there is no model it should be turned off for. +- **`base_num_decoder_layers` a multiple of 8?** You are paying ~25% of step time to an XLA layout + pathology, in *either* mode, and `param_scan_axis: 0` recovers it (§4.9). Unrelated to + `shard_mode`, larger than everything else in this document, and not yet validated as a default — + measure it on your config. +- **mixtral only:** explicit with **neither** kernel-order flag is ~1.4% faster than `auto` at + L12–L16, and the shipped default gives that up to land on parity (−0.078% at L14, §5.6). If you + run mixtral and can measure it, `lm_head_weight_grad_in_kernel_order: false` is worth trying; no + other model behaves this way, and deepseek behaves the opposite way (§4.8). - **gemma3 specifically:** never benchmark it below 6 layers, and prefer a multiple of 6. `scan_length = num_layers // 6`, so a 4-layer proxy unrolls the whole decoder and measures a different program (§4.7), and any remainder is unrolled alongside the scan. @@ -1335,14 +1605,31 @@ ______________________________________________________________________ > `lm_head_kernel_transposed`, storing the kernel as `[vocab, embed]`: that one works, but it is > dominated on every axis by what follows, so it has been removed rather than shipped beside it. - Since the barrier cannot be removed, remove the *need* for the fold. Autodiff's default transpose - rule builds the weight gradient as `transpose(dot(g, inputs))` and relies on `algsimp` folding the - transpose back into the dot; under explicit that fold never happens. A hand-written `custom_vjp` - contracts `dk` straight into the kernel's own axis order instead, so there is no transpose to - fold. `_dot_general_kernel_ordered_grad` in `linears.py` implements it, `DenseGeneral`'s - `weight_grad_in_kernel_order` gates it, and `lm_head_weight_grad_in_kernel_order` in `base.yml` - wires it to the untied head. It self-disables under `auto` (where the fold does happen) - and under quantization (where it does not own the dot). + Since the barrier cannot be removed globally, remove it *locally*, for the length of the one dot + that needs the fold. Autodiff's default transpose rule builds the weight gradient as + `transpose(dot(g, inputs))` and relies on `algsimp` folding the transpose back into the dot; + under explicit that fold never happens because a `Sharding` custom-call sits between the two. + Tracing the forward dot inside a `jax.sharding.auto_axes` region takes its operands' shardings + out of the type system for that dot only — Shardy then re-propagates them exactly as + `shard_mode: auto` does — so the barrier is never emitted there, XLA folds the transpose itself, + and the gradient comes out in the kernel's stored order. `_dot_general_in_auto_axes` in + `linears.py` is the whole implementation, `DenseGeneral`'s `weight_grad_in_kernel_order` gates + it, and `lm_head_weight_grad_in_kernel_order` in `base.yml` wires it to the untied head. It + self-disables outside an explicit mesh (where there is no barrier) and under quantization (where + it does not own the dot). + + > *Implementation note (2026-09-03).* This shipped first as a hand-written `custom_vjp` + > (`_dot_general_kernel_ordered_grad`) that contracted `dk` straight into the kernel's axis order, + > so that no transpose was ever emitted. The `auto_axes` region reaches the same place by letting + > XLA do the fold, and is better on four counts: it is **114 lines shorter**, its gradients are + > **bit-identical** to the default rule rather than agreeing to rounding (it reassociates + > nothing — the numerical-equivalence table below becomes vacuous for column D), it has no + > batch-dimension restriction so it reaches all 17 call sites including MLA and the MoE shared + > MLP, and it cannot drift from JAX's transpose rule because it *is* JAX's transpose rule. The + > two were measured against each other on all five models at both flag settings and agree to + > within **0.15 pp** everywhere (llama2 −0.494 vs −0.449, qwen3 −0.537 vs −0.443, mixtral +0.041 + > vs +0.083, gemma3 +0.803 vs +0.958, deepseek +0.176 vs +0.202). The `custom_vjp` has been + > deleted. **Why this and not the transposed kernel.** They reach the same layout, but only the `custom_vjp` recovers the *collective*. On llama2/shrink/fsdp4 the kernel-order dump gains @@ -1433,8 +1720,8 @@ ______________________________________________________________________ `resolve_lm_head_weight_grad_in_kernel_order` in `types.py` does the resolution, rejecting an explicit `true` on a tied head rather than silently ignoring it. -02. **Ship `dense_weight_grad_in_kernel_order` — but leave it off, and say why (§4.8).** The same - `custom_vjp` wired to every other `DenseGeneral`: the three `MlpBlock` projections +02. **Ship `dense_weight_grad_in_kernel_order` — on by default under explicit (§4.8).** The same + region wired to every other `DenseGeneral`: the three `MlpBlock` projections (`linears.py:650,670,689`), the four attention projections (`attentions.py:700,737,779,835`) and the nine MLA / Indexer projections (`attention_mla.py:141,157,177,814,830,852,868,893`). @@ -1464,7 +1751,51 @@ ______________________________________________________________________ a default would be fitting an XLA layout heuristic. **Leave it off, document the two models, and revisit if a third data point ever lands between 3 and 16.** -03. **Do not build a lookup table you index once (§5.5) — and do not expect an A/B to find it.** + > *Reversed (2026-09-03).* The "discriminator" was §4.9. Stack depth 16 is one of the lengths + > where `param_scan_axis: 1` degenerates and the step gets 25% more expensive in **both** modes; + > the six models that "lost" were all measured there, and what was being read as a layout trade + > was scheduling noise on top of an artifact. Measured at depths that are not multiples of 8 — + > 15 (model, depth) pairs across seven models — the flag never costs more than 0.19 pp and + > removes a +0.35% to +1.04% penalty on 12 of 15 (§4.8, resolution table). It is therefore now + > tri-state like its LM-head counterpart and resolves to **on** under `shard_mode: explicit`; + > `resolve_dense_weight_grad_in_kernel_order` in `types.py` does the resolution. There is no + > tied-head carve-out here — every model has these projections. + > + > The two flags are independent and additive: at L18, explicit costs +1.098% (llama2) / +5.579% + > (qwen3) with neither, +0.621% / +4.853% with this one alone, +0.429% / +0.534% with the LM-head + > flag alone, and **+0.005% / +0.008%** with both. + > + > One model pays for the *LM-head* default: mixtral is −1.4% against `auto` with **neither** flag + > at L12/L14/L16 and lands on parity with both. That is an MoE-scan scheduling windfall unrelated + > to the head (an 8× vocab shrink leaves it unchanged) and deepseek's MoE has the opposite sign, + > so it is disclosed rather than encoded. A mixtral user who measures it can set + > `lm_head_weight_grad_in_kernel_order: false`. + +03. **Investigate `param_scan_axis: 0` as the default — it is worth ~25% at any depth that is a + multiple of 8, in both modes (§4.9).** This is not a sharding item and it is the largest number + in this document. `param_scan_axis: 1` stacks scanned parameters as `[in, L, out]`; when `L` is + a multiple of 8 (also at L2 and L4) XLA lays the gradient stack out `{2,0,1}` and every + per-iteration `dynamic-update-slice` writes a `T(1,128)`-tiled slice — one row per tile. + Measured on llama2 at d16: **L8 −23.6%, L16 −24.9%** from flipping to `param_scan_axis: 0`, + against −0.18…−0.28% at L9/L18 where the pathology is absent. It persists at `prod` geometry + (+14.1% / +11.5%). The ns/layer trend across L2…L24 is flat except at exactly those depths. + + What is *not* established: whether `param_scan_axis: 0` is safe as a global default. It changes + the on-disk parameter-stack layout, so it interacts with checkpoint compatibility and with every + optimizer-state layout, and it has only been measured on one model family, one geometry and one + chip count. The tractable next steps are (a) reproduce the ns/layer discontinuity on a second + model family and a second topology, (b) check whether the degenerate layout is reachable through + an XLA layout hint instead of a config change, and (c) work out whether a checkpoint written at + one `param_scan_axis` can be read at the other. **Until then, the actionable advice is narrower: + do not benchmark anything at a layer count that is a multiple of 8, and if your production depth + is one, try the flag.** + + The one thing this item settles immediately is a measurement rule for the rest of this document: + §5.1/§5.2 (L4), §5.4 (L16, plus deepseek at L8) and §4.8's original table (L16) are all inside + the pathology, and a sub-1% sharding delta measured there is measuring the artifact's scheduling + noise. §4.8 and §5.6 are the depths where that is not true. + +04. **Do not build a lookup table you index once (§5.5) — and do not expect an A/B to find it.** `YarnRotaryEmbedding.freqs_cis` materialized `f32[163840, 32]` inside every layer of every step to read 1024 rows out of it; computing the rows straight from `position` is **−1.4% of step time in both modes**. It survived a full mode-comparison sweep precisely because both modes paid @@ -1478,7 +1809,7 @@ ______________________________________________________________________ `broadcast(f32[] 0)` at 0.8–1.2% of step is JAX's `lax.scan` `ys` output buffers, not a MaxText bug. -04. **Do not bother eliding redundant `reshard` barriers — tested, worth exactly zero.** The obvious +05. **Do not bother eliding redundant `reshard` barriers — tested, worth exactly zero.** The obvious companion fix to item 1 is to stop emitting a barrier when the reshard is a no-op (the value's aval already carries the requested sharding), on the theory that fewer barriers means more folds. It was implemented and measured, and it does not work: @@ -1496,39 +1827,39 @@ ______________________________________________________________________ elided. Recorded here so the next person does not spend the week: **the barrier count is not the lever, the gradient orientation is.** -05. **Add `shard_mode` to the golden suites (§7.3).** Until CI compares modes, every regression here +06. **Add `shard_mode` to the golden suites (§7.3).** Until CI compares modes, every regression here can silently return. `tests/integration/hlo_diff_test.py` is the cheapest place to start; assert on `after_codegen` reduce-scatter counts, not on static collective byte totals. -06. **Any HLO byte-counter used to compare modes must descend into +07. **Any HLO byte-counter used to compare modes must descend into `kind=kCustom, calls=%all-reduce-scatter.*` fusions** and report the post-slice shape. Counting the `all-reduce` inside the fusion at its pre-slice shape reports auto's wins as if they were losses, which is how the "identical bytes" claim survived a full sweep here. -07. **Re-examine the `check_vma` ban (§7.1).** JAX 0.11.1 accepts `check_vma=True` on an all-Explicit +08. **Re-examine the `check_vma` ban (§7.1).** JAX 0.11.1 accepts `check_vma=True` on an all-Explicit mesh; the restriction appears to be MaxText-imposed. Lifting it would give the MoE models on the explicit whitelist back a documented optimization. -08. **Validate `fused_qkv` + explicit and context-parallel + explicit** in `types.py` so users get a +09. **Validate `fused_qkv` + explicit and context-parallel + explicit** in `types.py` so users get a MaxText message instead of a raw JAX trace error. -09. **Fix `gcs_utils.upload_dump()`** to no-op (or write locally) for non-`gs://` paths — today it +10. **Fix `gcs_utils.upload_dump()`** to no-op (or write locally) for non-`gs://` paths — today it crashes after training and destroys the profile. -10. **Make explicit mode not defeat remat CSE on unrolled blocks (§4.7).** With `remat_policy: full` +11. **Make explicit mode not defeat remat CSE on unrolled blocks (§4.7).** With `remat_policy: full` and no scan, auto CSEs the recomputed forward against the original (61 remat-tagged instructions of 11,209); explicit cannot (1,013 of 12,452), costing 279.5 µs. This bites any config where a block is unrolled — gemma3 below 6 layers today, but also `scan_layers: false` runs generally. -11. **Deduplicate the tied-embedding bf16 cast (§4.6).** `embeddings.py:235-240`'s `.T` makes explicit +12. **Deduplicate the tied-embedding bf16 cast (§4.6).** `embeddings.py:235-240`'s `.T` makes explicit materialize the 128 MB embedding shard in bf16 **twice** (+80–82 µs on every gemma). Hoisting the cast above the transpose, or reusing one converted value for both consumers, is a pure win and is independent of everything else in this document. -12. **Investigate the tuple-repack all-reduce tax (§6.2)** — the one measured explicit cost that grows +13. **Investigate the tuple-repack all-reduce tax (§6.2)** — the one measured explicit cost that grows with depth (+112 → +219 → +379 µs from 4 to 12 layers on gemma3). -13. **Delete or rewire `validate_shard_mode`** (§7.4) — as dead code it actively misinforms. +14. **Delete or rewire `validate_shard_mode`** (§7.4) — as dead code it actively misinforms. ### Levers that were tried and do not work @@ -1538,12 +1869,12 @@ reasoned about, and each survived an adversarial attempt to rescue it. | candidate | verdict | | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | | Stop pinning the logits output under explicit | 0-line HLO diff. The barriers are JAX's, not MaxText's (§4.5). | -| Elide no-op `reshard` barriers | Op-identical executable on all three configs (item 04). | +| Elide no-op `reshard` barriers | Op-identical executable on all three configs (item 05). | | `lm_head_kernel_transposed` (store the kernel `[vocab, embed]`) | Works, but slower on 4 of 7 configs, no checkpoint path, and re-initializes the model. Removed (item 01). | | Reshard the RMSNorm scale under explicit outside TP | Implemented and reverted — negative on every non-TP config. | | Sliced/prefetched MoE expert weights via MSA | No measurable effect; the expert all-gathers are already overlapped. | | Remove the rotated-reduce-scatter `collective-permute` fixup | Backwards: it is an `auto`-mode op that explicit already deletes (§6.3). There was nothing on the explicit side to remove, and it is not caused by `_permuted_sharding`. | -| `param_scan_axis: 1` + XLA tiling as the dense-flag discriminator | Does not separate the winners from the losers; every model in the sweep uses the same value. | +| `param_scan_axis: 1` + XLA tiling as the dense-flag discriminator | Correct as stated — every model uses the same value, so it cannot separate models — but the conclusion drawn from it was wrong. It separates *depths*, and that is the whole discriminator: §4.9, worth ~25% of step time in both modes. | | 3-D kernels / MLA low-rank / gemma3's attention pattern / contracting-axis-vs-FSDP, as the dense-flag discriminator | None separate the sign. gemma3-L18's win is carried by plain 2-D `[embed, mlp]` `MlpBlock` kernels — the same sites that lose on llama2 — and deepseek's by its length-1 dense-layer scan (4,640 µs of `f32[512,1,8192]` copies) plus 500 µs on the MoE shared MLP, not by MLA (MLA's 4-D gradients move ~170 µs). | | Rematerializing RoPE sin/cos in the backward pass | Mode-neutral: identical to 0.1 µs in both arms. | | `LlamaVisionRotaryEmbedding`'s full-table build | Consumes the whole table by design, and `llama4` is not on the explicit whitelist — unreachable. | @@ -1557,6 +1888,14 @@ is llama2's scan gradient-stack `dynamic-update-slice` at +68 µs/step (+0.098%) identical signature and op count (960 / 960) in both arms with a sign that flips by model — i.e. it is scheduling around a shared op, not an explicit-mode defect. +> *Postscript (2026-09-03).* That last sentence is right about the sharding question and badly +> understated about everything else. The shared op it dismisses is §4.9's degenerate +> `dynamic-update-slice`, and at d16/L16 it is not +68 µs of *delta* but **8.2 ms of step**, ~22% of +> the whole thing, paid identically by both modes. A/B diffing two modes cannot see a cost both +> modes pay — the same blind spot that hid the YaRN table (§5.5). Both were found only by looking at +> one mode's absolute profile, which is now the second entry in a very short list of things a +> mode-comparison sweep structurally cannot find. + ______________________________________________________________________ ## 9. Reproducibility @@ -1583,6 +1922,9 @@ that cost time to discover: profile is flushed (§8 item 6). - Clear `jax_cache_dir`. A persistent-cache hit skips compilation and therefore suppresses the HLO dump, while producing an identical executable and step time. +- The geometry above is the *shrink* used by §5.1–§5.2 and by §4's dumps. Everything from §4.8 + onward — including §5.6 — uses `base_emb_dim=2048 base_mlp_dim=8192 base_num_query_heads=8 base_num_kv_heads=8 head_dim=128` (referred to as **d16**) with `base_num_decoder_layers` varied + per row, plus `base_moe_mlp_dim=8192` on mixtral. §5.6 writes out nothing else but `shard_mode`. Timings are `train_step` durations read off `/device:TPU:0` in the `.xplane.pb`, median of the 3 profiled steps; HLO facts are from `*after_optimizations.txt` and `*after_codegen*`. The TPU is a @@ -1621,18 +1963,57 @@ Added with the 2026-09-03 revision: a production model and it is not a production mesh. Nothing here was re-measured at emb 4096, seq ≥ 4096, pure DP, pure TP or multi-host with the new default on — the §5.2 and §5.3 rows are all pre-flag and should be read as upper bounds. + - **3 reps per arm, medians compared.** Several §5.4 deltas (mixtral +0.009%, deepseek's +0.012% over 8 reps) are inside the noise floor and are reported as ties, not wins. Only the ≥ 0.2% rows should be read as directional. + - **The dense flag's discriminator is fitted to eight points**, with two on one side of the threshold and six on the other, and it is a layout proxy rather than a proven mechanism (§4.8). It is documented as a per-model recommendation for exactly the two models it was measured on. + + > *Withdrawn (2026-09-03).* There is no discriminator; all eight points were measured inside + > §4.9's layout pathology. See §4.8's resolution table. + - **The YaRN finding was measured on deepseek only.** It is the only model in the sweep whose `max_position_embeddings` (163,840) dwarfs `max_target_length` by enough for the table build to dominate; the fix is unconditional, so other YaRN models get whatever their ratio is worth, which was not measured. + - **`deepseek3-671b-batchsplit` newly resolves the LM-head flag to `true`** and was verified by inspection only — the flag's permutation is the identity for a 2-D `[embed_vocab, vocab]` kernel, and the `reduced=`/`unreduced=` tags live on decoder-layer weight gradients inside `deepseek_batchsplit.py`'s `shard_map`s, never on the LM head. That config does not fit on the hardware used here, so it has not been run. + + > *Amended (2026-09-03).* It now resolves **both** flags to `true`. The dense flag reaches the + > decoder-layer projections, which is where `deepseek_batchsplit.py`'s `reduced=`/`unreduced=` + > tags live. Those tags are applied to gradients inside a `shard_map`, and the `auto_axes` region + > wraps only the forward `dot_general` in `DenseGeneral` — a different program point — so the two + > should not interact. That reasoning has not been run on hardware either. + +Added with the second 2026-09-03 revision: + +- **§4.9's pathology is characterized on one model family.** llama2 at d16 and at `prod`, plus the + consequence visible in gemma3's and deepseek's short scans. The rule "stack length a multiple of 8 + (also 2 and 4)" is read off 15 depths on one model; it has not been checked on a second topology, + a second chip generation, or a non-Adam optimizer, and no XLA-side explanation for *why* layout + assignment picks `{2,0,1}` at those lengths has been established. +- **The dense flag's new default rests on 15 (model, depth) pairs at one geometry.** Seven models, + depths 12–20, d16, 2 reps each. Its worst measured cost is +0.134% (mixtral L12). It has not been + measured at `prod`, at emb 4096, under TP or DP, or on any depth above 20. +- **The `auto_axes` and `custom_vjp` implementations were compared at L16** — inside the pathology. + Their agreement (≤0.15 pp on five models, both flag settings) is therefore an agreement between + two noisy measurements; the claim it supports is "the swap is not a regression", not "the swap is + worth +0.05%". The bit-exactness claim is separate and is asserted by unit test, not by benchmark. +- **mixtral's −1.4% no-flag win is unexplained.** It reproduces at L12, L14 and L16, survives an 8× + vocab shrink, and has the opposite sign on the other MoE model in the sweep. It is disclosed as a + cost of the shipped default rather than attributed. +- **§5.6's two MoE rows are at different depths from the other six.** mixtral is at 14 layers and + deepseek at 12 because at 18 both exceed HBM at this geometry (deepseek asks 43.56 G of HLO + temporaries against 31.24 G available). Both are healthy stack lengths under §4.9, so the + comparison is sound, but the eight rows are not all the same program size and the table should not + be read as a cross-model ranking — only each row against its own `auto`. +- **Nothing in §5.6 is multi-host.** Every number in this document is a single 4-chip v5e host. The + collective topology that explicit's `reduced=`/`unreduced=` tags exist to exploit (§6.1) only + appears at scale, and so does the failure mode that would most plausibly break parity.