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

src/maxtext/layers/nnx_scan.py

@@ -103,6 +103,7 @@ def apply_scanned_layers(
     remat_policy: Callable[..., Any] | None = None,
     prevent_cse: bool = True,
     unroll: int = 1,
+    static_unroll: bool = False,
 ) -> Any:
   """Applies stacked NNX layers using ``jax.lax.scan``.
 
@@ -114,6 +115,12 @@ def apply_scanned_layers(
   Externally managed per-layer state, such as KV caches, is not supported by
   this scan path.
 
+  ``static_unroll`` swaps ``jax.lax.scan`` for a Python loop over the stacked
+  layers. The traced program is identical apart from the loop structure, but the
+  per-layer rematerialization residuals stay separate values instead of being
+  stacked on a leading axis, which is what an enclosing scan needs when the
+  policy offloads them to pinned host.
+
   Note:
     This is a minimal, model-agnostic scan primitive. The other NNX decoder
     paths instead go through ``NNXDecoder._apply_layers_sequentially``, a heavier
@@ -189,10 +196,31 @@ def scan_body(current_carry, scanned_state):
     )
     return next_carry, (new_params, updated_rest)
 
-  scan_fn = jax.checkpoint(scan_body, policy=remat_policy, prevent_cse=prevent_cse) if remat else scan_body
-  final_carry, (scanned_new_params, scanned_rest) = jax.lax.scan(
-      scan_fn, carry, (params, rest), length=length, unroll=unroll
-  )
+  if static_unroll:
+    # ``prevent_cse=False`` relies on the scan boundary to keep XLA from folding a
+    # rematerialized forward back into the original one. Straight-line code has no
+    # such boundary, so leaving CSE enabled would undo the rematerialization: on
+    # qwen3-next-80b at 8k tokens it costs 3.4GB of HBM.
+    prevent_cse = True
+
+  body_fn = jax.checkpoint(scan_body, policy=remat_policy, prevent_cse=prevent_cse) if remat else scan_body
+
+  if static_unroll:
+    # 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``.
+    final_carry = carry
+    per_iteration = []
+    for i in range(length):
+      slice_i = jax.tree.map(lambda x, i=i: x[i], (params, rest))
+      final_carry, outputs = body_fn(final_carry, slice_i)
+      per_iteration.append(outputs)
+    scanned_new_params, scanned_rest = jax.tree.map(lambda *xs: jnp.stack(xs), *per_iteration)
+  else:
+    final_carry, (scanned_new_params, scanned_rest) = jax.lax.scan(
+        body_fn, carry, (params, rest), length=length, unroll=unroll
+    )
 
   if param_scan_axis != 0:
     scanned_new_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, param_scan_axis), scanned_new_params)

src/maxtext/models/qwen3.py

@@ -1538,6 +1538,7 @@ def _scan_local_layers(self, y, layer_kwargs):
         remat=remat,
         remat_policy=self.remat_policy_fn if remat else None,
         prevent_cse=maxtext_utils.should_prevent_cse_in_remat(self.config) if remat else True,
+        static_unroll=remat and maxtext_utils.offload_needs_static_unroll(self.config),
     )
 
   def _scan_global_layer(self, y, layer_kwargs):
@@ -1639,6 +1640,9 @@ def __call__(
     cfg = self.config
     inputs = carry
     inputs = self._maybe_shard_with_logical(inputs, self.activation_axis_names)
+    # Name the block input so remat policies can save or offload it, as in Gemma4's
+    # scannable block. The per-layer inputs inside the block carry the same name.
+    inputs = checkpoint_name(inputs, "decoder_layer_input")
 
     layer_kwargs = {
         "decoder_segment_ids": decoder_segment_ids,
@@ -1794,6 +1798,9 @@ def __call__(
     if isinstance(inputs, tuple):
       inputs = inputs[0]
     inputs = self._maybe_shard_with_logical(inputs, self.activation_axis_names)
+    # The residual entering the layer is the natural remat checkpoint: it is cheap to
+    # keep (or offload) and everything else in the layer can be recomputed from it.
+    inputs = checkpoint_name(inputs, "decoder_layer_input")
     residual = inputs
 
     # First LayerNorm, applied before the attention block.

src/maxtext/utils/maxtext_utils.py

@@ -255,6 +255,40 @@ def get_save_and_offload_names(config) -> tuple[list[str], list[str]]:
   return [], []
 
 
+def offload_needs_static_unroll(config) -> bool:
+  """Whether a rematerialized inner layer loop must be a Python loop rather than a scan.
+
+  Hybrid models (Gemma4, Qwen3-Next) scan over *blocks*, and each block scans over
+  its own layers and rematerializes them itself. When the remat policy offloads a
+  tensor, the block's inner scan emits that residual as a pinned-host stacked
+  output and the outer block scan stacks it once more. XLA's host-offload pass
+  cannot pair the copy-start/copy-done across those two levels, and the compile
+  fails post-optimization (``Bitcast cannot have different memory spaces``, or an
+  ``async-start ... operand`` shape mismatch between the ``S(5)`` and default
+  memory spaces).
+
+  Unrolling the inner loop in Python keeps every offloaded residual exactly one
+  scan level deep -- the same shape a homogeneous model such as Llama2 produces --
+  while leaving the outer block loop rolled, so XLA still pipelines one block's
+  host copies at a time. ``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.
+
+  Only applied when the policy actually offloads something. It grows the block
+  body by the number of layers per block, which costs compile time that runs
+  keeping every checkpointed tensor in HBM should not pay.
+
+  Losing the inner loop also costs some HBM, because the rematerialized forward
+  is no longer confined to one loop iteration, so the unroll only pays for itself
+  when the policy offloads enough. On qwen3-next-80b at 8k tokens the unroll
+  costs 2.5GB and ``minimal_offloaded`` moves 16.9GB off device, a 13.2GB win;
+  offloading only ``decoder_layer_input`` moves 1.2GB against a 4.3GB unroll and
+  comes out 1.3GB behind.
+  """
+  _, offload_names = get_save_and_offload_names(config)
+  return bool(offload_names)
+
+
 def load_compiled(config, partial_train, state, execution_devices):
   """# Loading a serialized compiled train step function."""
 

tests/unit/maxtext_utils_test.py

@@ -1766,5 +1766,28 @@ def test_non_offloading_policies_return_empty(self):
       self.assertEqual(maxtext_utils.get_save_and_offload_names(self._cfg(policy)), ([], []))
 
 
+class TestOffloadNeedsStaticUnroll(unittest.TestCase):
+  """Tests for maxtext_utils.offload_needs_static_unroll (pure config logic, no device needed)."""
+
+  _cfg = staticmethod(TestGetSaveAndOffloadNames._cfg)  # pylint: disable=protected-access
+
+  def test_offloading_policy_needs_a_static_unroll(self):
+    """An offloaded residual cannot cross two nested scans, so the inner loop must be a Python loop."""
+    cfg = self._cfg("custom", tensors_to_offload=["decoder_layer_input"])
+    self.assertTrue(maxtext_utils.offload_needs_static_unroll(cfg))
+
+  def test_preset_offload_policies_also_need_it(self):
+    """The named offload presets go through the same save/offload split, so they unroll too."""
+    for policy in ("qkv_proj_offloaded", "minimal_offloaded"):
+      self.assertTrue(maxtext_utils.offload_needs_static_unroll(self._cfg(policy)))
+
+  def test_device_only_policies_keep_the_inner_scan(self):
+    """Unrolling costs compile time, so policies that offload nothing must not pay it."""
+    self.assertFalse(maxtext_utils.offload_needs_static_unroll(self._cfg("full")))
+    self.assertFalse(
+        maxtext_utils.offload_needs_static_unroll(self._cfg("custom", tensors_on_device=["decoder_layer_input"]))
+    )
+
+
 if __name__ == "__main__":
   unittest.main()

tests/unit/nnx_decoders_test.py

@@ -1160,6 +1160,74 @@ def test_decoder_keeps_layers_past_the_last_whole_block(self):
     )
 
 
+class TestQwen3NextDecoderLayerInputOffload(unittest.TestCase):
+  """`decoder_layer_input` must be nameable and offloadable on the Qwen3-Next stack.
+
+  Two things have to hold for `remat_policy=custom decoder_layer_input=offload` to
+  do anything. The layers have to attach the name (otherwise the policy matches
+  nothing and the setting is silently ignored), and the offloaded residual has to
+  survive lowering -- XLA cannot carry a pinned-host tensor out of two nested
+  scans, which is why the block loop unrolls when the policy offloads.
+  """
+
+  def _grad_jaxpr(self, num_decoder_layers=8, **overrides):
+    """Returns the printed jaxpr of a gradient through the scanned Qwen3-Next decoder."""
+    cfg = _make_config(
+        **{
+            **_QWEN3_NEXT_CONFIG,
+            "base_num_decoder_layers": num_decoder_layers,
+            "scan_layers": True,
+            **overrides,
+        }
+    )
+    mesh = _make_mesh(cfg)
+    decoder = NNXDecoder(config=cfg, mesh=mesh, model_mode=MODEL_MODE_TRAIN, rngs=nnx.Rngs(params=0, dropout=1))
+    embedding = Embed(
+        num_embeddings=cfg.vocab_size,
+        num_features=cfg.emb_dim,
+        dtype=cfg.dtype,
+        config=cfg,
+        mesh=mesh,
+        rngs=nnx.Rngs(params=0),
+    )
+    batch, seq = cfg.global_batch_size_to_train_on, cfg.max_target_length
+    ids = jnp.zeros((batch, seq), dtype=jnp.int32)
+    segment_ids = jnp.full((batch, seq), DECODING_ACTIVE_SEQUENCE_INDICATOR)
+    positions = jnp.broadcast_to(jnp.arange(seq)[None], (batch, seq))
+
+    decoder_graphdef, decoder_params, decoder_rest = nnx.split(decoder, nnx.Param, ...)
+    embed_graphdef, embed_params, embed_rest = nnx.split(embedding, nnx.Param, ...)
+
+    def loss(dec_params, emb_params, dec_rest, emb_rest):
+      dec = nnx.merge(decoder_graphdef, dec_params, dec_rest)
+      emb = nnx.merge(embed_graphdef, emb_params, emb_rest)
+      logits, _, _ = dec(
+          emb,
+          ids,
+          decoder_positions=positions,
+          decoder_segment_ids=segment_ids,
+          deterministic=True,
+          model_mode=MODEL_MODE_TRAIN,
+      )
+      return jnp.sum(logits**2)
+
+    jaxpr = jax.make_jaxpr(jax.grad(loss, argnums=(0, 1)))(decoder_params, embed_params, decoder_rest, embed_rest)
+    return jaxpr.pretty_print(use_color=False)
+
+  def test_layers_name_their_input_for_the_remat_policy(self):
+    """Without the name, `decoder_layer_input=device|offload` would be a no-op."""
+    self.assertIn("name=decoder_layer_input", self._grad_jaxpr(remat_policy="full"))
+
+  def test_offloading_policy_moves_the_residual_to_host(self):
+    """`MemorySpace.Host` is how a device_put to pinned host renders in a jaxpr."""
+    offloaded = self._grad_jaxpr(remat_policy="custom", decoder_layer_input="offload")
+    self.assertIn("MemorySpace.Host", offloaded)
+
+  def test_no_host_transfer_without_an_offloading_policy(self):
+    """Nothing should reach host memory when every checkpointed tensor stays on device."""
+    self.assertNotIn("MemorySpace.Host", self._grad_jaxpr(remat_policy="full"))
+
+
 class TestNNXDecoderDeepseekAndGemma4(unittest.TestCase):
   """Tests for Deepseek and Gemma4 specific decoder logic."""
 

tests/unit/train_compile_test.py

@@ -793,6 +793,33 @@ def test_qwen3_next(self):
         )
     )
 
+  def test_qwen3_next_offloaded_decoder_layer_input(self):
+    """AOT test for offloading a Qwen3-Next remat checkpoint to pinned host.
+
+    Qwen3-Next scans over blocks that scan (and rematerialize) their own layers.
+    An offloaded residual therefore leaves two nested scans, which XLA could not
+    lower until the block loop was unrolled -- the compile used to fail
+    post-optimization with a memory-space mismatch. Kept small: the unroll makes
+    compile time grow with the block count.
+    """
+    compiled_trainstep_file = "/tmp/test_qwen3_next_offloaded"
+    train_compile_main(
+        (
+            "",
+            get_test_config_path(),
+            f"compiled_trainstep_file={compiled_trainstep_file}",
+            "compile_topology=v5p-64",
+            "compile_topology_num_slices=1",
+            "model_name=qwen3-next-80b-a3b",
+            "override_model_config=true",
+            "base_num_decoder_layers=8",
+            "per_device_batch_size=1",
+            "max_target_length=1024",
+            "remat_policy=custom",
+            "decoder_layer_input=offload",
+        )
+    )
+
   def test_deepseek32(self):
     # test deepseek3.2 with sparse attention
     compiled_trainstep_file = "/tmp/test_deepseek32.pickle"