fix(qwen3-next): make activation offload work under the nested block scan

[NuojCheng]

Stacked on #4964 — please review/merge that first; this PR targets its branch, so the diff here is just the two files of new logic plus tests.

Summary

Offloading a rematerialization checkpoint to pinned host does not work on Qwen3-Next today. decoder_layer_input=offload is silently ignored, and the named offload presets (minimal_offloaded, qkv_proj_offloaded) fail to compile. This PR fixes both, which is what makes remat_policy=minimal_offloaded usable on the model: 69.74 GB → 56.54 GB of HBM at 8k tokens, a 19% reduction.

The two bugs

1. The checkpoint name was never attached. Qwen3NextDecoderLayer never called checkpoint_name(inputs, "decoder_layer_input"), so the remat policy had nothing to match and both decoder_layer_input=device and decoder_layer_input=offload were no-ops. An AOT compile of qwen3-next-80b-a3b produces byte-identical memory stats for the two settings (temp_size_in_bytes=46,779,974,304 for both). Only the unrelated dense Qwen3DecoderLayer had the name, via AttentionWithNorm.

The fix attaches it in the layer and in Qwen3NextScannableBlock, mirroring what Gemma4DecoderLayer and Gemma4ScannableBlock already do.

2. Once named, offload does not compile. The block's inner local-layer scan emits the offloaded residual as a pinned-host stacked output, and the outer block scan stacks it a second time. XLA's host-offload pass cannot pair the copy-start with the copy-done across two scan levels, and the compile fails post-optimization with either Bitcast cannot have different memory spaces or an async-start expects the shape of operand 0 to match the async shape mismatch between the S(5) and default memory spaces. Confirmed on the real model: the offending buffer is bf16[12,3,1,1024,2048]{...S(5)} = [12 blocks, 3 local layers, batch, seq, emb]. A 4-layer config (outer scan length 1) compiles; 8 and 12 layers fail.

The fix

apply_scanned_layers grows a static_unroll mode that replaces jax.lax.scan with a Python loop over the stacked layers. Each layer's residual then stays a separate value instead of being stacked on a leading axis, so the outer block scan stacks it exactly once — the same single-level shape a homogeneous model such as Llama2 produces, which XLA handles. The outer block loop stays rolled, so host copies still pipeline one block at a time.

Two things that look like they should work but don't:

• jax.lax.scan(..., unroll=n) is not a substitute. It removes the loop but still stacks the per-iteration outputs, so the extra dimension survives and the compile still fails.
• Breaking the nesting at the outer block loop instead does compile, but it costs about 7 GB of HBM and 4x the compile time, which defeats the purpose of offloading. The inner loop is the right place to break it.

The unrolled path also forces prevent_cse=True. should_prevent_cse_in_remat returns False whenever scan_layers=True, because the scan boundary is what stops XLA folding a rematerialized forward back into the original one. Straight-line code has no such boundary, so leaving CSE enabled silently undoes the rematerialization — worth 3.4 GB of HBM on its own.

The whole mechanism is gated on the policy actually offloading something (maxtext_utils.offload_needs_static_unroll), so runs that keep every checkpointed tensor in HBM pay nothing: same scan, same compile time as before.

Measurements

AOT compiles of qwen3-next-80b-a3b on v5p-64, max_target_length=8192, per_device_batch_size=1. HBM is temp_size_in_bytes from compiled.memory_analysis().

remat_policy	HBM	host	compile
minimal	69.74 GB	0	54 s
minimal_offloaded	56.54 GB	16.91 GB	109 s

Isolating the two effects under minimal: the static unroll costs 2.48 GB (69.74 → 72.22 GB with offload disabled), and the offload returns 15.68 GB, close to 1:1 on the 16.91 GB moved off device.

The practical value is that offload buys minimal-level recompute at close to full-remat memory: 56.54 GB with minimal's FLOPs, against 69.74 GB for minimal on device or 46.78 GB for a policy that recomputes everything.

What this does not buy

Offloading decoder_layer_input on its own is a weak lever on this model, and the PR does not pretend otherwise:

remat_policy=custom, pdbs	rolled + device	unrolled + device	unrolled + offload	net
1	46.78 GB	51.12 GB	48.07 GB (1.21 GB host)	+1.29 GB
4	92.05 GB	99.77 GB	92.66 GB (4.83 GB host)	+0.61 GB

The offload itself works and is roughly 1:1, but the unroll costs slightly more than it returns. Two reasons: jax.checkpoint always keeps its own inputs as residuals regardless of policy, so naming decoder_layer_input and setting it to device changes HBM by 163 KB — the only thing the name buys is the ability to offload it; and under custom almost everything else is recomputed, which makes the rematerialized region that the unroll has to inline large. Under minimal that region is small, which is why the unroll tax there is 2.48 GB rather than 4.34 GB.

The gap narrows with batch size but does not cross zero, and per_device_batch_size=8 does not fit on v5p-64 at all (108 GB required vs 95.74 GB available), so there is no headroom above 4 to find a crossover. This is documented on offload_needs_static_unroll so it does not have to be re-measured.

Testing

• tests/unit/nnx_decoders_test.py::TestQwen3NextDecoderLayerInputOffload — three jaxpr-level tests on the scanned Qwen3-Next decoder: the layers name their input, an offloading policy puts MemorySpace.Host in the gradient jaxpr, and a device-only policy does not.
• tests/unit/train_compile_test.py::test_qwen3_next_offloaded_decoder_layer_input — AOT regression test on v5p-64. Verified that it fails without the fix and passes with it.
• tests/unit/maxtext_utils_test.py::TestOffloadNeedsStaticUnroll — the gating logic, including that device-only policies keep the inner scan.
• Full tests/unit/nnx_decoders_test.py (60 passed, 2 skipped) and the qwen3-next and gemma4 AOT tests pass. pre-commit clean.

Follow-up

Gemma4 has the same nested-scan offload problem and currently pays the expensive workaround unconditionally: _apply_gemma4_scanned_blocks sets block_unroll = max(1, scan_length) for every run, offload or not. The same static_unroll=True argument would let it keep its block loop rolled. Left out of this PR because it needs its own AOT validation.

[gemini-code-assist]

Code Review

This pull request introduces a static_unroll option in apply_scanned_layers to replace jax.lax.scan with a Python loop when offloading rematerialized residuals to pinned host. This prevents nested scans from causing XLA compilation failures due to memory-space mismatches. Additionally, it names the decoder layer inputs in Qwen3 to enable offloading and adds corresponding unit and AOT compile tests. The feedback suggests clarifying a comment in nnx_scan.py regarding the Python loop and manual stacking of outputs to avoid confusion for future maintainers.

[gemini-code-assist]

The comment here is a bit confusing. It states that outputs are not stacked, but jnp.stack is called on line 217. The key distinction is that using a Python loop avoids jax.lax.scan's implicit handling of rematerialization residuals, which would otherwise stack them and cause issues with nested offloading. The layer outputs are then manually stacked to mimic lax.scan's behavior for the returned values.

To improve clarity for future maintainers, consider rephrasing the comment to focus on how this approach affects rematerialization residuals rather than the layer outputs.

Suggested change

	# A Python loop instead of ``jax.lax.scan``: the per-iteration outputs stay
	# separate values rather than being stacked into a leading-axis array. That
	# matters when this helper runs inside another scan and ``remat_policy``
	# offloads a residual to pinned host -- see ``offload_needs_static_unroll``.
	# A Python loop instead of ``jax.lax.scan``. This traces `length` separate calls
	# to `body_fn`, which prevents `jax.lax.scan` from implicitly stacking rematerialization
	# residuals. This is critical when this helper runs inside another scan and `remat_policy`
	# offloads a residual to pinned host -- see `offload_needs_static_unroll`.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!