Onboard Qwen3.5 and Qwen3-Next to explicit sharding and ZeRO-1 gradient accumulation

[NuojCheng]

Description

Onboards the qwen3_5 and qwen3_next decoder blocks to shard_mode=explicit, and
adds ZeRO-1 + efficient gradient accumulation coverage on top of them.

The two are the hybrid members of the Qwen3 family: both interleave GatedDeltaNet
linear attention with a full-attention layer every
inhomogeneous_layer_cycle_interval layers, on top of a fully-MoE block with a shared
expert. Qwen3.5 subclasses the Qwen3-Next sublayers outright, so almost all of the work
is shared. Two commits:

1. Onboard Qwen3.5 to explicit sharding and ZeRO-1 gradient accumulation — the
   shared sublayers plus Qwen3_5DecoderLayer.
2. Onboard Qwen3-Next to explicit sharding — Qwen3NextDecoderLayer and
   Qwen3NextScannableBlock, which is the only structural difference between the two.

The decoder layers

Qwen3_5DecoderLayer, Qwen3NextDecoderLayer and Qwen3NextScannableBlock all still
expressed their activation layouts with nn.with_logical_constraint, so
shard_mode=explicit rejected both blocks at config validation time. They now use the
same pattern the Qwen3, Qwen2 and DeepSeek layers already use: cache the physical
NamedShardings in __init__, pin every sublayer output through
maybe_shard_with_logical, and thread out_sharding / intermediate_sharding into
the norms, both attention variants and the MoE block. Under ShardMode.AUTO the
callees ignore those arguments and the constraints become no-ops, matching the pattern
cdfade207 established for Qwen3 so XLA keeps the fusions it picks today.

Qwen3NextScannableBlock is the part that is new relative to Qwen3.5, which runs one
cycle as a plain Python loop. Qwen3-Next puts its linear-attention layers inside
nnx_scan.apply_scanned_layers and its lone full-attention layer inside a
trip-count-one jax.lax.scan, and both scans require the carry's layout to be
invariant across iterations. That holds because the decoder layer now returns the
layout it was handed, so nothing extra is needed inside the scans themselves.

The shared sublayers

• Qwen3NextGatedDeltaNet is the interesting one. Explicit sharding cannot infer a
  layout across its reshapes, its head repeat or the jax.shard_map boundary, so each
  intermediate is pinned to the logical axes the auto path already asks GSPMD for (a
  new _explicit_activation_shardings helper returns the flat / head / state layouts,
  all None under AUTO). A_log and dt_bias are stored replicated but broadcast
  against (B, S, H_v) activations whose head axis is sharded, so they are resharded
  onto the head axis first — the same fix _align_scale_with_normalized_axis applies
  to the norm scales. Finally, shard_map manualises the mesh axes it is given and
  will not insert a reshard for an operand whose layout differs from in_specs, so its
  six operands are resharded to match before the call. The existing
  with_sharding_constraint on mixed_qkvz is routed through maybe_shard_with_name,
  since under EXPLICIT with_sharding_constraint is an assertion rather than a hint
  and rejects any layout it would have to change.
• Qwen3NextRMSNorm was accepting shard_mode, kernel_axes and
  parameter_memory_host_offload and then silently dropping them on the floor; it now
  forwards them to the inner RMSNorm. No caller passes the latter two today, so that
  part is a no-op fix.
• Qwen3NextRMSNormGated had no way to take an out_sharding; it now accepts one
  and forwards it.
• Attention forwards config.shard_mode to the hybrid q/k norms.

Qwen3NextGatedDeltaNet, Qwen3NextFullAttention and Qwen3NextSparseMoeBlock all
gained an out_sharding keyword (plus intermediate_sharding on the MoE block).
Nothing inside the jax.shard_map body changed.

Deliberate gaps

Two new config guards reject combinations that have not been onboarded. Both come
from the shared gated-delta-net path, so both apply to qwen3_5 and qwen3_next:

• sparse_matmul=False — RoutedMoE.dense_matmul is still not onboarded to explicit
  sharding. This gap is shared with qwen3_moe and kimi and is left for a follow-up.
• context parallelism (ici/dcn_context_parallelism,
  ici/dcn_context_usp_ulysses_parallelism) — the gated-delta-net short convolution
  left-pads the sequence by gdn_conv_kernel_dim - 1 and slices the result back, which
  explicit sharding cannot express on a sharded sequence axis. shard_mode=auto still
  works there.

The Qwen3-VL / Qwen3-Omni encoders remain unsupported; they are covered by the existing
use_multimodal guard.

Known issue, not fixed here

shard_mode=explicit + shard_optimizer_over_data=True + ici_expert_parallelism>1
fails for every MoE model, not just the two here:

ShardingTypeError: add got incompatible shardings for broadcasting:
  ('expert', None, None), (('data', 'expert'), None, None)

add_data_to_sharding prepends data to the optimizer mu/nu spec while the gradients
keep the plain params spec. This reproduces on unmodified main with mixtral-8x7b at
DP2 x EP2 with ZeRO-1 and 4 accumulation steps (it passes under shard_mode=auto). The
ZeRO-1 coverage added here is therefore data-parallel only, which is the same coverage
every other onboarded MoE model has. A fix is left for a follow-up.

Tests

All commands below were run on a v4-8 TPU VM.

New AOT coverage (tests/unit/train_compile_test.py), one pair per decoder:

• test_qwen3_5_explicit_sharding — qwen3.5-397b-a17b on v5p-512, FSDP 32 x expert
  8, sparse_matmul/megablox/flash/tokamax-splash.
• test_qwen3_next_explicit_sharding — qwen3-next-80b-a3b on v5p-512, same mesh.
• test_qwen3_5_explicit_sharding_zero1 and test_qwen3_next_explicit_sharding_zero1 —
  four-layer variants on v5p-256, DP 128, gradient_accumulation_steps=4,
  shard_optimizer_over_data=True.

These are the large-scale checks: explicit sharding type-checks every operation's
layout instead of letting GSPMD infer one, so a missing out_sharding fails the trace
here rather than silently costing a collective at a scale a real test cannot reach.
Four layers is one full inhomogeneous_layer_cycle_interval, so both attention
variants are still covered while the model stays small enough to hold data-parallel
replicas of the parameters.

$ pytest tests/unit/train_compile_test.py -k "qwen3_5 or qwen3_next" --durations=0
84.11s test_qwen3_5_explicit_sharding          69.10s test_qwen3_5   (existing)
79.57s test_qwen3_5_explicit_sharding_zero1    54.63s test_qwen3_next (existing)
57.16s test_qwen3_next_explicit_sharding
48.74s test_qwen3_next_explicit_sharding_zero1
6 passed in 403s

Each new case costs about what the AUTO sibling next to it already costs.

New TPU integration coverage (tests/integration/train_tests.py), parameterized
over both hybrid decoders rather than duplicated:

• test_tpu_qwen3_hybrid_explicit_sharding_matches_auto — explicit vs auto. Following
  the convention test_tpu_qwen3_explicit_sharding_matches_auto already uses, each
  decoder is paired with the one parallelism that stresses a different half of the
  shared layer rather than running the cross-product: Qwen3.5 under expert parallelism,
  which shards the MoE dispatch, and Qwen3-Next under tensor parallelism, which shards
  the gated-delta-net head axis — the one the layer has to carry by hand across its
  reshapes, its head repeat and the shard_map boundary.
• test_tpu_qwen3_hybrid_zero1_gradient_accumulation — explicit + ZeRO-1 + 8
  accumulation steps against the auto, non-ZeRO-1 baseline.

$ pytest tests/integration/train_tests.py -k qwen3_hybrid
2 passed, 4 subtests passed in 124s

Both run with scan_layers on, so the Qwen3-Next subtests exercise the nested scans.
Qwen3-Next was additionally spot-checked with scan_layers=False (the unscanned
Qwen3NextDecoderLayer path) under both shard modes.

The existing _qwen3_losses / _mistral_losses helpers were near-identical, so they
are now thin delegates over a shared _losses(run_name, model_overrides, extra_args)
rather than more copies being added.

Numerical note. Unlike mistral and qwen3, the hybrid decoders are not
bit-for-bit between explicit and auto. activation_batch carries the expert axis, so
pinning it reassociates the backward reductions rather than leaving the layout
untouched — the forward pass is bit-for-bit and the drift only appears once gradients
flow. Measured over 20 steps it stays below 3.1e-5 relative and changes sign, i.e. it
is float noise rather than the two runs pulling apart, so the parity assertions use
rtol=1e-4. For the same reason each decoder's AUTO path itself moves by up to
2.8e-5 relative against main (its with_logical_constraint calls become AUTO
no-ops), which is the tradeoff cdfade207 deliberately made for Qwen3 to restore XLA
fusion.

New config-validation coverage (tests/unit/pyconfig_test.py): qwen3_5 and
qwen3_next moved into the accepted-decoder loop, and
test_explicit_sharding_gated_delta_net_unsupported_combinations covers the two new
guards for both decoders.

$ pytest tests/unit/pyconfig_test.py -k "explicit_sharding or zero1"
8 passed, 9 subtests passed in 10s

Regressions checked:

$ pytest tests/integration/train_tests.py \
    -k "qwen3_explicit or mistral_explicit or qwen3_zero1 or mistral_zero1 or gemma_zero1"
5 passed, 13 subtests passed in 199s

$ pytest tests/unit/qwen3_next_shared_expert_test.py tests/unit/qwen35_partial_mrope_test.py \
    tests/unit/router_replay_test.py tests/unit/sharding_test.py tests/unit/sharding_nnx_test.py \
    tests/unit/maxtext_utils_test.py tests/unit/param_mapping_test.py tests/unit/nnx_decoders_test.py \
    tests/unit/decoder_layer_model_mode_test.py
300 passed, 5 skipped, 5 subtests passed in 762s

• The three attention_test.py failures that show up when the CPU device count is
  raised above 1 are present unchanged on main.
• tests/integration/hlo_diff_test.py needs no reference-HLO regeneration: it covers
  deepseek3, llama3-8b and qwen3-1.7b, and no hunk in this PR is reachable from any of
  them (every qwen3.py change is inside a Qwen3Next* class, and the attentions.py
  change is inside the is_qwen3_hybrid branch).
• pylint reports 9.96/10 on the changed files, with all four remaining messages
  present unchanged on main.

Checklist

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

• 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]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!