Onboard the Llama4 decoder to explicit sharding

[NuojCheng]

Description

Onboards the Llama4 decoder to shard_mode=explicit, following the same pattern as the Qwen3, Mistral and DeepSeek blocks that are already supported.

Llama4DecoderLayer and Llama4ScannableBlock now route their activation constraints through maybe_shard_with_logical instead of nn.with_logical_constraint, cache the NamedShardings for the activation and MLP-intermediate layouts in __init__, and thread out_sharding / intermediate_sharding into the attention, RMSNorm, dense MLP and RoutedAndSharedMoE sublayers. Under ShardMode.AUTO every one of those arguments is ignored and GSPMD infers exactly the layouts it does today, so this is a no-op for existing runs.

RoutedAndSharedMoE already accepts both sharding arguments (it is shared with DeepSeek), so no MoE changes were needed.

The rotary embedding fix

LLaMARotaryEmbedding — which Llama4 reaches via rope_type: "llama3.1" — could not run under explicit sharding at all. It built the lax.select predicate by tiling a 1-D parity mask up to the full input shape:

jax.lax.select(jnp.tile(mask, inputs.shape[:-1] + (1,)), shifted_right, shifted_left)

That produces a replicated [B, S, N, H] predicate, and sharding-in-types requires the predicate to match the cases in sharding as well as shape, so the trace fails with:

ShardingTypeError: select 'which' must be scalar or have the same sharding as cases,
got P(None, None, None, None) and P('expert', None, None, None)

Broadcasting the same 1-D mask through jnp.where selects identical elements and carries no sharding of its own. The two formulations are mathematically identical — verified elementwise-equal in the forward pass, in the gradient, and in bf16 — but they fuse differently in XLA, and swapping unconditionally shifted llama3.1 losses by ~1e-5 in ShardMode.AUTO. Since llama3.1 is out of scope here, the rewrite is gated on shard_mode so AUTO keeps byte-identical behaviour; this was confirmed by re-running the auto path against pristine main.

Testing

• tests/unit/pyconfig_test.py — decoder_block=llama4 is accepted under explicit sharding, and use_multimodal=True is still rejected (only the text stack is onboarded; the Llama4 vision encoder is not).
• tests/unit/train_compile_test.py — two AOT compiles: llama4-17b-16e with ZeRO-1 + gradient accumulation on v5p-256, and llama4-17b-128e at FSDP 64 × expert 8 on v5p-1024.
• tests/integration/train_tests.py — the integration model uses 4 layers with interleave_moe_layer_step=2, which covers all four Llama4 layer variants in one run: chunked+dense, chunked+MoE, chunked+dense and global (NoPE) + MoE.

Measured on a local TPU v7-8, 3 steps of synthetic data:

Configuration	auto	explicit	max rel. diff
tensor parallelism	[8.119888, 8.105545, 8.097639]	[8.119888, 8.105545, 8.097639]	bit-for-bit
expert parallelism	[8.119880, 8.105589, 8.097667]	[8.119880, 8.105585, 8.097659]	4.7e-7
ZeRO-1 + grad accum (vs. unsharded auto)	[8.120317, 8.107924, 8.101201]	[8.120396, 8.107988, 8.101174]	9.6e-6

The ZeRO-1 test carries @pytest.mark.skip_on_tpu7x for consistency with the other ZeRO-1 tests (b/517509898); the numbers above were produced by temporarily lifting that marker, and the marker was restored afterwards.

Regression sweep: the pre-existing qwen3, qwen3_moe, qwen3_custom_moe, mistral and mixtral explicit-vs-auto parity tests, and the existing AOT explicit/ZeRO-1 compiles, all still pass. pyink is clean and pylint reports only pre-existing R0917/W0613 findings on untouched signatures.

Known gaps

• RoutedMoE.dense_matmul (sparse_matmul=False) is still not onboarded to explicit sharding — a gap Llama4 shares with qwen3_moe and mixtral.
• Llama4 supports top-1 routing only (num_experts_per_tok=1), so the top-k dispatch path is not exercised.
• The Llama4 vision encoder is untouched; use_multimodal remains rejected under explicit sharding.

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

# Unit
JAX_PLATFORMS=cpu pytest tests/unit/pyconfig_test.py -k explicit_sharding
JAX_PLATFORMS=cpu pytest tests/unit/train_compile_test.py -k "explicit or zero1 or llama4"

# Integration, on a TPU host
pytest tests/integration/train_tests.py -k "explicit or zero1"

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• 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.

[gemini-code-assist]

Code Review

This pull request adds support for explicit sharding in the Llama4 model. Key changes include updating the LLaMA embedding shift logic to use jnp.where under explicit sharding to avoid predicate sharding issues, configuring physical shardings and logical constraints in the Llama4 decoder, and adding comprehensive integration, unit, and AOT compilation tests to verify correctness. There are no review comments, and I have no additional feedback to provide.

[codecov]

Codecov Report

❌ Patch coverage is 8.00000% with 23 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/models/llama4.py	9.52%	19 Missing ⚠️
src/maxtext/layers/embeddings.py	0.00%	4 Missing ⚠️

📢 Thoughts on this report? Let us know!