From 36d4288e41c2c5ef35e1f6ff3e4dafc51885a55e Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Tue, 25 Aug 2026 02:38:35 +0000 Subject: [PATCH] fix(qwen3-next): make activation offload work under the nested block scan Offloading a remat checkpoint to pinned host was broken on Qwen3-Next, for two independent reasons. First, `Qwen3NextDecoderLayer` never attached the `decoder_layer_input` checkpoint name, so `decoder_layer_input=device` and `decoder_layer_input=offload` were both silent no-ops -- an AOT compile of qwen3-next-80b produces byte-identical memory stats for the two. Only the unrelated dense `Qwen3DecoderLayer` had the name, via `AttentionWithNorm`. Attach it in the layer and in `Qwen3NextScannableBlock`, mirroring Gemma4. Second, once the name is attached the compile fails. The block's inner local-layer scan emits the offloaded residual as a pinned-host stacked output and the outer block scan stacks it again; XLA's host-offload pass cannot pair the copy-start/copy-done across two levels and it fails post-optimization with `Bitcast cannot have different memory spaces` or an `async-start ... operand` shape mismatch between the `S(5)` and default memory spaces. `apply_scanned_layers` grows a `static_unroll` mode that replaces `jax.lax.scan` with a Python loop, so each layer's residual stays a separate value rather than being stacked on a leading axis. The outer block scan then stacks it exactly once -- the same single-level shape a homogeneous model such as Llama2 produces -- and the outer loop stays rolled, so host copies still pipeline one block at a time. `jax.lax.scan(..., unroll=n)` is not a substitute: it drops the loop but still stacks the per-iteration outputs, so the extra dimension survives. Breaking the nesting at the outer block loop instead works too, but costs 7GB of HBM and 4x the compile time. Straight-line code also has no scan boundary, so the mode 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; without a loop, leaving CSE on undoes the rematerialization, worth 3.4GB on its own. The unroll is gated on the policy actually offloading something, so runs that keep every checkpointed tensor in HBM pay nothing. AOT on qwen3-next-80b-a3b, v5p-64, 8k tokens, per_device_batch_size=1: remat_policy=minimal 69.74GB HBM, 0GB host, 54s compile remat_policy=minimal_offloaded 56.54GB HBM, 16.91GB host, 109s compile Offloading only `decoder_layer_input` is a much weaker lever on this model (48.07GB vs 46.78GB): a `jax.checkpoint` always keeps its own inputs as residuals regardless of policy, so the tensor is small relative to what the unroll costs. The named offload presets, which move whole projections, are where offloading pays. This is documented on `offload_needs_static_unroll`. --- src/maxtext/layers/nnx_scan.py | 36 ++++++++++++++-- src/maxtext/models/qwen3.py | 7 +++ src/maxtext/utils/maxtext_utils.py | 34 +++++++++++++++ tests/unit/maxtext_utils_test.py | 23 ++++++++++ tests/unit/nnx_decoders_test.py | 68 ++++++++++++++++++++++++++++++ tests/unit/train_compile_test.py | 27 ++++++++++++ 6 files changed, 191 insertions(+), 4 deletions(-) diff --git a/src/maxtext/layers/nnx_scan.py b/src/maxtext/layers/nnx_scan.py index 19d49502bd..6f5e147e06 100644 --- a/src/maxtext/layers/nnx_scan.py +++ b/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) diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index 94f070f2bb..5140228f9c 100644 --- a/src/maxtext/models/qwen3.py +++ b/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. diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index ba76420997..6f6cea7176 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ b/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.""" diff --git a/tests/unit/maxtext_utils_test.py b/tests/unit/maxtext_utils_test.py index ed18ca6256..2971d65e61 100644 --- a/tests/unit/maxtext_utils_test.py +++ b/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() diff --git a/tests/unit/nnx_decoders_test.py b/tests/unit/nnx_decoders_test.py index 4041aff7d6..a6f40c9db3 100644 --- a/tests/unit/nnx_decoders_test.py +++ b/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.""" diff --git a/tests/unit/train_compile_test.py b/tests/unit/train_compile_test.py index 95eccd4843..ab75c85a67 100644 --- a/tests/unit/train_compile_test.py +++ b/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"