Add explicit sharding support for the Qwen3 model family

[NuojCheng]

Description

Onboards the Qwen3 family to explicit sharding (shard-in-types), following the pattern established for llama2 (#2470) and deepseek (#2783). The qwen3, qwen3_moe and qwen3_custom_moe decoder blocks — 21 of the 23 checked-in qwen3 model configs — now run with shard_mode=explicit, including ZeRO-1 optimizer sharding combined with gradient accumulation.

Model changes

AttentionWithNorm is shared by the dense, MoE and custom-MoE Qwen3 layers, so most of the work lands there:

• It builds out_sharding / mlp_intermediate_sharding from the logical axis rules and threads them through the norms, the attention call and the MLP/MoE block. Without an explicit output sharding the attention out-projection has no legal layout to infer under explicit axes.
• shard_mode is passed into the pre/post attention RMSNorms.
• nn.with_logical_constraint is replaced with maybe_shard_with_logical, which reshards under explicit axes and constrains under auto, so the same call site serves both modes.

qwen3_custom.py gets the same treatment for its latent_norm and layer_up_projection paths.

RMSNorm fix

The RMSNorm scale is stored with the norm logical axis (→ tensor). That is correct for a norm over the embedding axis, but wrong for Qwen3's per-head QK norm: it normalizes head_dim, which is unsharded, while the neighbouring heads axis is the one on tensor. The multiply therefore placed tensor on two axes of the result and JAX rejected it:

ShardingTypeError: dot_general operation with inputs: bf16[128@tensor], bf16[4,256,4@tensor,128]
produces an illegally sharded result: bf16[128@tensor,4,256,4@tensor]

The scale is now realigned to the activation's last-dim sharding before the multiply. This runs only under ShardMode.EXPLICIT and is a no-op when the two already agree, so no other model's layout changes. The alternative — relabelling the QK-norm kernel_axes — would have churned every model's golden sharding descriptors.

Config guards

• use_multimodal is rejected under explicit sharding: the Qwen3-VL/Omni vision and audio encoders are not onboarded (a reshape in the patch embed fails).
• qwen3_next and qwen3_5 remain unsupported — they use gated-delta-net linear attention and are a separate piece of work.

Tests

Explicit sharding is only a change in how layouts are expressed, so it must be numerically neutral. The new tests run the same tiny model under both shard modes and compare per-step losses:

• train_tests.py::test_tpu_qwen3_explicit_sharding_matches_auto — all three decoder blocks, each paired with the parallelism that stresses it most (tensor parallelism for the dense QK norm, expert parallelism for the MoE dispatch, FSDP for the unscanned custom MoE).
• train_tests.py::test_tpu_qwen3_zero1_gradient_accumulation — shard_optimizer_over_data + gradient_accumulation_steps=8 under explicit sharding versus a plain auto + GA baseline, for the dense and MoE decoders. This is what exercises the reduced/unreduced gradient PartitionSpec labels.
• pyconfig_test.py::test_explicit_sharding_qwen3_decoder_support — the accept/reject config guards.
• One smoke test for the custom-MoE block under explicit sharding.

Verification

Run on a v6e-4:

• Every auto-vs-explicit pair matches bit-for-bit (e.g. dense TP-4 189.796 / 189.196 / 188.858 in both modes; MoE EP-4 8.082 / 8.065 / 8.055). Beyond the configurations kept in the test file, parity was also confirmed manually for DP2×FSDP2, unscanned layers, and the custom-MoE latent up-projection path.
• ZeRO-1 + GA matches the auto + GA baseline exactly for both MoE decoders and to ~1e-5 for dense (ZeRO-1 reassociates the gradient all-reduce, so the assertion allows rtol=1e-4).
• No regressions: the llama2 test_tpu_zero1_gradient_accumulation still passes, and 185 sharding/config unit tests pass. The qwen3-0.6b golden input-sharding dump was regenerated; the diff is purely additive (two activation entries the decoder now reports through maybe_shard_with_logical).

Checklist

• 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 introduces support for explicit sharding in Qwen3 decoders (including MoE and custom MoE variants) across MaxText. It adds helper utilities to align normalization scales with normalized axes, updates model definitions to propagate sharding constraints, and includes comprehensive integration and unit tests to verify that explicit sharding produces identical loss trajectories to auto-sharding. The review feedback highlights two issues: a critical bug in normalizations.py where the non-existent jax.typeof function is called (which will raise an AttributeError at runtime), and a typo in train_smoke_test.py where a raw string is used instead of an f-string, preventing the interpolation of dataset_path.

[gemini-code-assist]

The use of jax.typeof(y) and jax.typeof(scale) will raise an AttributeError at runtime because jax.typeof is not a valid JAX function. Instead, you should access the .sharding attribute directly on the arrays/tracers (y.sharding and scale.sharding).

Additionally, to prevent potential IndexError or TypeError crashes when running on single-device or fully replicated setups (where spec can be None or empty P()), we should add defensive checks before indexing activation_spec and scale_spec.

Suggested change

	activation_spec = jax.typeof(y).sharding.spec
	scale_spec = jax.typeof(scale).sharding.spec
	if scale_spec[-1] == activation_spec[-1]:
	return scale
	return jax.sharding.reshard(scale, jax.sharding.PartitionSpec(activation_spec[-1]))
	activation_spec = y.sharding.spec
	scale_spec = scale.sharding.spec
	activation_last_axis = activation_spec[-1] if (activation_spec is not None and len(activation_spec) > 0) else None
	scale_last_axis = scale_spec[-1] if (scale_spec is not None and len(scale_spec) > 0) else None
	if scale_last_axis == activation_last_axis:
	return scale
	return jax.sharding.reshard(scale, jax.sharding.PartitionSpec(activation_last_axis))

[gemini-code-assist]

The dataset_path argument is passed as a raw string r"dataset_path={self.dataset_path}" instead of an f-string. This prevents Python from interpolating self.dataset_path and passes the literal string "{self.dataset_path}". While this test currently passes because dataset_type=synthetic ignores the dataset path, it is a copy-paste typo that should be corrected to an f-string for correctness and consistency.

Suggested change

	r"dataset_path={self.dataset_path}",
	f"dataset_path={self.dataset_path}",

[codecov]

Codecov Report

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

Files with missing lines	Patch %	Lines
src/maxtext/models/qwen3_custom.py	75.00%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!