Remove the explicit-sharding LM head transpose via lm_head_kernel_transposed

[NuojCheng]

Description

Adds lm_head_kernel_transposed (default false), a config flag that stores the untied LM head kernel as [vocab, embed] instead of [embed, vocab]. This removes the single largest source of the shard_mode: explicit performance penalty. On a 4-way-DP v5p-8 configuration it takes explicit mode from +9.34% slower than auto to −0.23% — i.e. slightly faster than auto — with the same numerics.

Why

shard_mode: explicit is measurably slower than shard_mode: auto across the onboarded models. Tracing it through the HLO, the root cause is a missed algebraic simplification, not a communication-schedule difference.

On an all-AxisType.Explicit mesh, JAX's lower_with_sharding_in_types (jax/_src/interpreters/mlir.py) annotates the output of every sharding-in-types primitive (~60 of them) with a custom_call_target="Sharding" op. It has no equality short-circuit, so it emits the annotation even when the sharding is unchanged. This is a JAX lowering behaviour — it is not caused by any with_sharding_constraint / reshard / out_sharding= pin in MaxText. (I verified this: removing the out_sharding= pin on the logits produces a 0-line HLO diff, and 0 of the 1,276 Sharding custom calls in the module originate from it.)

One of those Sharding calls lands between the LM head weight-gradient dot and the transpose that follows it. XLA's algebraic simplifier normally folds transpose(dot(A, B)) → dot(B, A), but the custom call sits in the middle and acts as an optimization barrier, so the fold never fires. The window matters: algsimp runs in simplification-1 (passes 0008–0009), while the partitioner runs at 0012–0016 — once the module is partitioned the fold is permanently unavailable.

The unfolded transpose costs two things:

1. Mechanism A — lost reduce-scatter. The weight gradient's all-reduce can no longer fuse into a reduce-scatter. This only bites at small embedding dims: the AR→DS fusion is rejected when the per-shard minor-most extent is below ~512 bytes (one 128-lane × 32-bit vreg row), independent of dtype. Worth up to ~28% of the total penalty, and only on the emb=512 configurations.
2. Mechanism B — entry relayout copies. The gradient materializes as [V, embed] while the parameter is [embed, V], so layout assignment inserts relayout copies on the entry parameters. Six f32[emb/n, 32000] copies, on every configuration measured. This is the dominant, universal cost.

The fix

Rather than trying to make the fold fire through the barrier, store the kernel in the orientation the gradient naturally comes out in. DenseGeneral gains a general-purpose kernel_transposed option:

• The kernel is allocated as out_features + in_features and dot_general contracts on its trailing axes. dot_general emits batch + lhs-free + rhs-free dims, so the output shape and order are unchanged.
• kernel_axes is permuted alongside the shape, so callers keep passing logical axes in the usual in..., out... order and sharding= keeps describing the same axis it did before.
• Autodiff then produces the weight gradient already in the kernel's orientation. There is no transpose to fold, so the barrier is irrelevant.

The untied LM head opts in via lm_head_kernel_transposed, wired through both the linen (decoders.apply_output_head) and nnx (nnx_decoders) paths.

Results

v5p-8, 2 reps per configuration, step time in µs:

config	auto	explicit [E,V]	explicit [V,E]	penalty before	penalty after
dp4	4430.2	4844.0	4420.1	+9.34%	−0.23%
emb4096	72200.8	73330.0	72601.8	+1.56%	+0.56%

That is −8.75% and −0.99% off explicit-mode step time, recovering 102% and 64% of the gap respectively.

The structural acceptance test passes on both: the f32[emb/n, 32000] relayout copies drop from 6 to 0 (matching auto), and the entry layout triple flips from f32[1024,32000] to f32[32000,1024].

Why it defaults to false

The flag changes the on-disk checkpoint orientation, which means auto mode gets the transposed kernel too — and there it regresses, because in auto mode the fold does fire and the original orientation was already optimal: +3.96% llama2-fsdp4, +1.91% mixtral, +0.98% qwen3, +0.46%/+0.35% at emb 2048/4096. Under min(auto, explicit) that is a net loss in 5 of 7 configurations, so it would be wrong to flip the default now.

Follow-up: a restore-time transform that keeps the checkpoint canonically at [embed, vocab] and transposes on load only when shard_mode: explicit. That decouples the two modes, at which point this can default on for explicit without touching auto. I'd like to land the mechanism first so the follow-up is a small change.

Shortcomings / known limitations

• Checkpoint compatibility. This changes the on-disk shape of params-decoder-logits_dense-kernel, so it cannot be flipped on an existing run without transposing that array. Called out in base.yml, in the types.py field description, and in the DenseGeneral docstring.
• Rejected combinations (raise rather than silently mishandle): quant (the LM head matmul is never quantized anyway), slice_bounds (it slices the trailing output-feature axis, which a transposed kernel keeps at the front), and logits_via_embedding: true (there is no untied head to transpose) — the last enforced by a pydantic model_validator.
• Tunix vLLM sharding map is updated for the flipped kernel. This one is important: both the old (None, "model") and new ("model", None) specs are rank 2, so getting it wrong would mis-shard silently instead of raising.
• Static vLLM converter (integration/vllm/weight_converter.py): MODEL_TO_CONVERSION_RULES is looked up without a config at all four call sites, so it cannot be made config-aware cheaply. The combination is documented as unsupported in a comment rather than silently producing wrong weights.

Analysis doc

The full investigation — HLO pass-pipeline walkthrough, xprof attribution, the AR→DS fusion threshold, per-model measurements, and the alternatives considered — is added as docs/guides/optimization/shard_mode_performance.md, linked from the optimization guide grid and toctree.

Tests

Unit tests — 7 new tests in tests/unit/linears_test.py covering the kernel_transposed mechanism:

JAX_PLATFORMS=cpu pytest tests/unit/linears_test.py -q
# 20 passed, 1 skipped, 2 subtests passed

• test_kernel_transposed_shape_and_axes — kernel shape and kernel_axes are both permuted
• test_kernel_transposed_forward_matches — forward output is bit-identical to the default orientation given transposed weights (subtests over scalar and (2, 8) output shapes)
• test_kernel_transposed_gradient_matches — gradient is bit-identical modulo the transpose
• test_kernel_transposed_bias — bias shape/axes track the flipped kernel
• test_kernel_transposed_init_statistics_match — fan_in is computed over the same axes, so init std is unchanged (0.24941 both ways)
• test_kernel_transposed_rejects_quantization, test_kernel_transposed_rejects_slice_bounds

End-to-end — v5p-8, the numbers in the table above, plus the HLO acceptance check on relayout-copy count and the entry layout triple. Config validation verified in both directions (logits_via_embedding=True + the flag raises; the flag alone flips the kernel to (256, 128) through the real model init).

Lint — pyink --check --pyink-indentation=2 --line-length=122 clean on all touched Python; pylint --disable=R0401,R0917,W0201,W0613 clean (the two remaining types.py findings are pre-existing on main); mdformat --check --number clean on both docs.

Checklist

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request introduces a detailed performance analysis guide comparing explicit and automatic sharding modes, and implements a performance optimization via the lm_head_kernel_transposed configuration. This option allows storing the untied LM head kernel as [vocab, embed] instead of [embed, vocab], which prevents an expensive transpose operation under explicit sharding. The changes update DenseGeneral to support the transposed kernel layout, add corresponding configuration validation, adjust integration mappings for Tunix and vLLM, and include comprehensive unit tests. The single review comment regarding a date typo in the documentation has been filtered out, leaving no further feedback to provide.

[codecov]

Codecov Report

❌ Patch coverage is 92.59259% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/integration/tunix/utils.py	0.00%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!