Update Qwen3.5 scan-remat pattern by re-using Qwen3-Next's nested block scan

[NuojCheng]

What

Extends the Qwen3-Next scanned-block work from #4964 to Qwen3.5.

Qwen3.5 repeats the exact same hybrid attention period as Qwen3-Next — inhomogeneous_layer_cycle_interval = 4, i.e. three GatedDeltaNet layers followed by one full-attention layer — so rather than porting the logic, this makes Qwen3.5 reuse it.

The two architectures now share the decoder layer as well as the block. Qwen3_5DecoderLayer subclasses Qwen3NextDecoderLayer and overrides only three factory hooks — _make_full_attention, _make_linear_attention, _make_mlp — so the forward pass, both LayerNorms, the residual structure, the is_full_attention_layer derivation and the MoE load-balance handling are inherited rather than duplicated. The hooks keep the submodule attribute names (input_layernorm, attention, post_attention_layernorm, mlp), so the checkpoint parameter paths are identical either way. Qwen3_5ScannableBlock likewise subclasses Qwen3NextScannableBlock and overrides only _make_decoder_layer.

models/qwen3_5.py goes from 266 to 127 lines. Qwen3.5 picks up the inner scan over the homogeneous linear-attention layers, the trip-count-one lax.scan scheduling barrier around the full-attention layer, and per-sub-layer remat, for free.

Changes

• models/qwen3.py — extract sub-block construction from Qwen3NextDecoderLayer.__init__ into _make_full_attention() / _make_linear_attention() / _make_mlp(), and layer construction from Qwen3NextScannableBlock.__init__ into _make_decoder_layer(). These are the only points a sibling architecture has to diverge at. Also threads forced_routed_experts through the layer and the block (see below). No behavior change for Qwen3-Next.
• models/qwen3_5.py — Qwen3_5DecoderLayer and Qwen3_5ScannableBlock become thin subclasses that override the hooks. Qwen3NextDecoderLayer gains an explicit is_full_attention_layer kwarg, because inside a scan a sub-layer's position in the cycle is not recoverable from layer_idx; it still falls back to the (layer_idx + 1) % cycle == 0 derivation when unscanned. The MoE load-balance loss is assigned rather than sow-n — sow appends to a tuple, so the layer's Intermediate structure grows on every call, which a lax.scan body cannot express.
• layers/decoders.py, layers/nnx_decoders.py — route QWEN3_5 through the shared hybrid path. _init_scanned_qwen3_next / _apply_qwen3_next_scanned_blocks are renamed to ..._qwen3_hybrid_... and pick the block class by decoder_block, including for layers_remainder. forced_routed_experts_per_layer() is factored out of reshape_forced_routed_experts_for_scan() so the hybrid path can slice routing per sub-layer.
• layers/nnx_scan.py — apply_scanned_layers() gains an optional xs pytree of per-layer scan inputs. Callers that pass no xs keep the two-argument apply_fn signature; Gemma4, the only other caller, is unchanged.
• layers/moe.py — hold the FSDP expert-kernel all-gather inside the decoder loop when weight_dtype == dtype. Not part of the port; it is a pre-existing bug in the scanned MoE path that this PR would otherwise widen to Qwen3.5. See The weight_dtype=bfloat16 blow-up below.
• checkpoint_conversion/utils/param_mapping.py — remap Qwen3.5's scanned checkpoint mapping from the old per-cycle-position layers-layer_{i} prefixes to the local_layers / global_layer tree, mirroring what perf(qwen3-next): nested scan over hybrid attention with per-layer remat #4964 did for Qwen3-Next. Local layers nest [block][local], the global layer is flat [block].
• checkpoint_conversion/utils/tensor_handling.py — stacked_axes() was guarded on isinstance(mt_key, str), so Qwen3.5's composite (wi_0, wi_1) key (fed by one fused HF gate_up_proj tensor) fell through to the leading-axis MoE layout instead of the nested-scan layout and would have silently transposed those weights. Tuple keys now resolve via their first component.
• integration/vllm/weight_converter.py, integration/vllm/maxtext_vllm_rollout.py — teach the RL rollout weight sync to read the nested tree (see below).

Results

AOT temp_size_in_bytes per device, using the config from train_compile_test.py::test_qwen3_5 (per_device_batch_size=1, max_target_length=1024, sparse_matmul, megablox, attention=flash, use_tokamax_splash, remat_policy=full — the base.yml default — on both sides). Both sides freshly compiled against main @ 457040f7:

model	topology	before	after	delta
qwen3.5-35b-a3b	tpu7x-16	30.36 GB	25.13 GB	−17.2% ✅
qwen3.5-35b-a3b	tpu7x-8	43.61 GB	44.11 GB	+1.1%
qwen3.5-397b-a17b	v5p-512	75.62 GB	46.01 GB	−39.2% ✅
qwen3.5-397b-a17b	v5p-128	128.05 GB (OOM)	116.25 GB (OOM)	−9%

Exact bytes: 35b/tpu7x-16 30359507040 → 25132249440; 35b/tpu7x-8 43612869024 → 44113347456; 397b/v5p-512 75622737696 → 46011566848. The 397b/v5p-128 OOM row is the original pre-rebase measurement and was not re-run.

The saving is a roughly constant 5–7 GB, not a constant percentage. Sweeping device count on identical silicon (qwen3.5-35b-a3b, tpu7x-N):

devices	before	after	delta
8	43.61 GB	44.11 GB	+0.50 GB
16	30.36 GB	25.13 GB	−5.23 GB
32	21.50 GB	15.42 GB	−6.08 GB
64	19.81 GB	14.70 GB	−5.11 GB
128	15.59 GB	8.51 GB	−7.08 GB

What shrinks with device count is the baseline, not the saving. The FSDP-sharded share of the working set falls as 1/N, while what this PR removes does not shard at all: sparse_matmul deliberately strips fsdp from the expert-kernel pspecs, so the gathered per-device MoE kernels are full size on every chip, and block-level remat keeps all four sub-layers' worth of them co-live. Per-sub-layer remat keeps one. The same absolute saving is therefore 17% of a 30 GB budget and 45% of a 16 GB one.

The cross-generation control agrees: v5p-128 and tpu7x-64 are both 64 JAX devices (v5p slice names count TensorCores) and start from quite different baselines — 16.16 GB and 19.81 GB — yet save 5.129 GB and 5.105 GB, within 0.5% of each other. v5p-512 (256 devices) saves 5.25 GB, 11.99 → 6.74 GB.

argument_size_in_bytes is unchanged in both directions (26.01 GB at 35b/tpu7x-16, 18.61 GB at 397b/v5p-512). Generated code size roughly halves on every row, including the ones where temp memory does not move — that is the nested scan replacing an unrolled cycle, and it is topology-independent: 35b/tpu7x-8 205 → 99 MB, 35b/tpu7x-16 210 → 99 MB, 397b/v5p-512 320 → 163 MB.

Real training on a v7-8 (qwen3.5-35b-a3b, synthetic data, ici_fsdp_parallelism=8, per_device_batch_size=1, max_target_length=1024, 6 steps) converges on both sides with matching loss curves (step 5: 7.056 before, 7.072 after — RNG-fork differences only). Memory there is a wash, 40.6 → 41.1 GB estimated temp, consistent with the tpu7x-8 row.

The layer-sharing refactor is free. Recompiling qwen3.5-35b-a3b at tpu7x-16 after collapsing Qwen3_5DecoderLayer onto Qwen3NextDecoderLayer and threading router replay gives temp_size_in_bytes=25132249440, generated_code_size_in_bytes=99265536 — byte-identical to the build before the refactor.

Test plan

All CPU-only, run individually against this branch:

suite	result
tests/unit/nnx_decoders_test.py	65 passed, 2 skipped (-k Qwen3: 17 passed, up from 7 on main)
tests/unit/moe_test.py	38 passed, 40 skipped (TPU-only)
tests/unit/router_replay_test.py	38 passed, 1 skipped
tests/unit/nnx_scan_test.py	4 passed
tests/unit/param_mapping_test.py	30 passed
tests/post_training/unit/weight_converter_test.py	27 passed

• The three Qwen3-Next test classes in nnx_decoders_test.py are parameterized by class attributes so the Qwen3.5 subclasses reuse them: nested-scan output matches a sequential unroll numerically, and the Linen and NNX parameter trees match, both for a whole number of periods and with a remainder. New: a block given forced routing that does not cover every sub-layer must raise.
• moe_test.py adds test_sparse_matmul_pins_expert_kernel_all_gather_when_the_upcast_is_an_identity, parameterized over (bf16, bf16), (f32, f32) and (f32, bf16). It asserts the barrier is applied exactly when the upcast is an identity, and — by object identity, since the stand-in kernels share a shape — that the sharding constraint consumes the barrier's outputs rather than the raw kernels, so a barrier whose results are dropped does not pass. Both matching-dtype cases fail without the moe.py change.
• param_mapping_test.py covers the local_layers / global_layer mapping shape and the composite-key axis fix.
• router_replay_test.py adds a case for the scan remainder block, where only the trailing layer's routing differs between two runs, so the losses can diverge only if the remainder replays its own slice. Reverting the qwen3.py / nnx_scan.py / nnx_decoders.py changes fails test_loss_fn_with_forced_routed_experts_scanned_qwen3_5 and test_per_layer_routing_is_not_broadcast_scanned.
• weight_converter_test.py adds a nested-layout source tree and asserts it converts byte-identically to the per-slot layout, plus a truncated cycle-slot axis being rejected. Both fail if the weight_converter.py change is reverted.
• tests/post_training/unit/vllm_rollout_unroll_test.py adds a nested local_layers + global_layer unroll case and a missing-slot-axis rejection. This file needs vllm installed to import the module under test, so it was verified in CI-equivalent form by executing the unroll helpers in isolation: the nested layout produces a dense layers_{0..7} with every slot on its own layer, and the legacy per-slot and homogeneous paths are unchanged. The same check fails against 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]

Code Review

This pull request implements hierarchical nested scans for Qwen3-Next and Qwen3.5 hybrid decoders, enabling per-layer rematerialization to optimize memory usage. It introduces scannable blocks that structure inhomogeneous attention cycles and updates checkpoint conversion utilities to support multi-axis stacked tensors with arbitrary nesting depths. Comprehensive unit tests are added to verify decoder parity and state restoration. The review feedback suggests updating the type annotations for mt_key from str to str | tuple[str, ...] in stacked_axes and _build_multi_axis_stacked_tensor to accurately reflect that composite keys can be represented as tuples.

[gemini-code-assist]

The mt_key parameter is annotated as str, but the function explicitly handles tuple inputs (lines 69-70). To ensure type safety and prevent static analysis warnings, the type annotation should be updated to str | tuple[str, ...].

Suggested change

	def stacked_axes(mt_key: str, config, depth: int) -> tuple:
	def stacked_axes(mt_key: str | tuple[str, ...], config, depth: int) -> tuple:

[gemini-code-assist]

Update the type annotation of mt_key to support tuple inputs as well, aligning with the updated stacked_axes signature.

Suggested change

	mt_key: str = "",
	mt_key: str | tuple[str, ...] = "",

[gemini-code-assist]

Update the type annotation of mt_key to support tuple inputs as well, aligning with the updated stacked_axes signature.

Suggested change

	mt_key: str = "",
	mt_key: str | tuple[str, ...] = "",

[codecov]

Codecov Report

❌ Patch coverage is 94.50549% with 10 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/nnx_decoders.py	86.48%	2 Missing and 3 partials ⚠️
src/maxtext/models/qwen3_5.py	86.36%	3 Missing ⚠️
...xtext/checkpoint_conversion/utils/param_mapping.py	95.45%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!