Fix NNX MTP zero-loss, qwix MTP crash, and DeepSeek scan-axis sharding corruption

Three independent bugs in the NNX path, surfaced by running DeepSeek-V3 with
multi-token prediction.

mtp_loss was silently 0. mtp_losses/mtp_acceptance subclass nnx.Intermediate and
nnx type filters match by isinstance, so the generic nnx.pop(model,
nnx.Intermediate) in loss_fn consumed them before the dedicated pop ran. That pop
then returned empty and calculate_mtp_loss fell through to its 0.0 default. Pop
the subclasses first. Linen is unaffected; it keys mutable collections by name.

qwix quantization crashed with "roll requires ndarray, got NoneType". NNX modules
are stateful, so quantize_model must run a real forward pass, but only tokens,
positions and segment ids were passed and the MTP block reads the decoder targets
unconditionally. Pass dummy targets, restoring the semantics of the
is_initializing() guard the NNX port left commented out. Linen never traces a
forward during quantization, so it never hit this.

DeepSeek params were left replicated, tripping assert_params_sufficiently_sharded
once the crash above was fixed. _create_scanned_layers records each stack's axis
name in nnx.PARTITION_NAME ("dense_layers"/"moe_layers"), but the scan-axis
round-trip hardcoded the literal "layers". The name never matched, so the fallback
stripped a real logical axis instead of the scan axis and a bogus "layers" was
inserted: ('embed','dense_layers','mlp') became ('dense_layers','layers','mlp').
'embed' is the only carrier of fsdp, hence P(None, ...). This is neither an MTP
nor a qwix bug: it corrupts on every forward for multi-stack models, but is
normally invisible because the sharding snapshot is taken before the first train
step. qwix makes it fatal only because it quantizes inside model construction,
within the traced snapshot. Resolve the axis name per-variable from
nnx.PARTITION_NAME, and raise on inconsistent metadata rather than blindly
stripping an axis, as flax's own nnx.spmd.remove_axis does.

Verified on deepseek3-tiny (v6e-8) with sharding_tolerance untouched: nnx+mtp and
nnx+mtp+qwix now both report mtp_loss 1.224, matching the Linen reference. The
sharding metadata round-trip goes from 29 of 32 params corrupted to 0. gpt3-52k,
gpt3-52k+qwix, gemma2-2b and deepseek3-tiny with scan_layers=false all train clean.

Add regression tests for all three. multi_token_prediction_test.py drives the real
loss_fn and asserts a non-zero MTP loss; quantizations_nnx_mtp_test.py quantizes an
MTP model through qwix without crashing; maxtext_utils_nnx_test.py checks the
scan-axis helpers keep the real logical axis. Each fails on the pre-fix code.

Author: ecnal-cienet

src/maxtext/layers/nnx_decoders.py

@@ -998,13 +998,9 @@ def layer_fn(carry, scanned_vars):
       final_carry, scanned_state = jax.lax.scan(layer_fn_wrapped, x_in, (params, state))
       returned_kv_stacked = None
 
-      # Ensure metadata rank matches the stacked values
-      scanned_state = maxtext_utils_nnx.nnx_add_scan_axis(scanned_state, "layers", 0)
-
-      if scan_axis != 0:
-        new_params, new_rest = scanned_state.split(nnx.Param, ...)
-        new_params = maxtext_utils_nnx.nnx_sync_moveaxis(new_params, 0, scan_axis)
-        scanned_state = nnx.merge_state(new_params, new_rest)
+      # Move the scan axis to each variable's param_scan_axis and restore its name
+      # in the sharding metadata. jax.lax.scan emits it at position 0.
+      scanned_state = maxtext_utils_nnx.nnx_add_and_sync_scan_axis(scanned_state, "layers")
 
       returned_kv_stacked = None
 

src/maxtext/layers/quantizations.py

@@ -861,13 +861,20 @@ def maybe_quantize_model(model, config):
         dummy_tokens = jnp.ones(input_shape, dtype=jnp.int32)
         dummy_positions = jnp.ones(input_shape, dtype=jnp.int32)
         dummy_segment_ids = jnp.ones(input_shape, dtype=jnp.int32)
+        # The MTP block reads the decoder targets, so the qwix forward pass needs them.
+        # The Linen path supplies them from the is_initializing() guard in Transformer.
+        dummy_targets = {}
+        if config.mtp_num_layers > 0:
+          dummy_targets["decoder_target_tokens"] = jnp.ones(input_shape, dtype=jnp.int32)
+          dummy_targets["decoder_target_mask"] = jnp.ones(input_shape, dtype=jnp.int32)
         model = qwix.quantize_model(
             model,
             quantization_provider,
             dummy_tokens,
             dummy_positions,
             dummy_segment_ids,
             enable_dropout=False,
+            **dummy_targets,
         )
         # Qwix quantization runs a forward pass during tracing, which sows transient nnx.Intermediate variables
         # (e.g. max_logits from QK-Clip, MTP losses) into the model. Popping them here prevents structural mismatches

src/maxtext/trainers/pre_train/train.py

@@ -197,18 +197,22 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr
         decoder_target_tokens=data["targets"],
         decoder_target_mask=data["targets_segmentation"],
     )
+    # mtp_losses and mtp_acceptance subclass nnx.Intermediate, and nnx type filters match
+    # subclasses. Pop them before the generic Intermediate pop below, which would otherwise
+    # take them too and leave the MTP loss silently reading as 0.
+    mtp_losses_state, mtp_acceptance_state = None, None
+    if config.mtp_num_layers > 0:
+      mtp_losses_state = nnx.pop(model, mtp_losses)
+      mtp_acceptance_state = nnx.pop(model, mtp_acceptance)
+
     intermediates = nnx.pop(model, nnx.Intermediate)
     intermediate_outputs = intermediates.to_pure_dict()
 
-    # MTP sows mtp_losses/mtp_acceptance as custom Variable subclasses, not
-    # Intermediate, so the nnx.pop above misses them. Pop them here under their
-    # collection names so calculate_mtp_loss / calculate_mtp_acceptance_rate
-    # find them. Otherwise the MTP loss is silently zeroed. They are also
-    # excluded from the returned state below so they don't leak into
-    # out_shardings.
-    if config.mtp_num_layers > 0:
-      intermediate_outputs["mtp_losses"] = nnx.pop(model, mtp_losses).to_pure_dict()
-      intermediate_outputs["mtp_acceptance"] = nnx.pop(model, mtp_acceptance).to_pure_dict()
+    # Store them under the collection name so calculate_mtp_loss and
+    # calculate_mtp_acceptance_rate find them at the same path as the Linen collections.
+    if mtp_losses_state is not None and mtp_acceptance_state is not None:
+      intermediate_outputs["mtp_losses"] = mtp_losses_state.to_pure_dict()
+      intermediate_outputs["mtp_acceptance"] = mtp_acceptance_state.to_pure_dict()
 
     if (config.use_indexer and not config.indexer_sparse_training) and is_train:
       # In Dense Warm-up stage, we skip main model loss calculation for efficiency.

src/maxtext/utils/maxtext_utils_nnx.py

@@ -218,64 +218,64 @@ def nnx_update_sharding_meta(variable, transform_fn):
   return variable
 
 
-def nnx_sync_moveaxis(tree, from_axis, to_axis):
-  """Moves an axis in both values and sharding metadata of nnx.Variables."""
-  if from_axis == to_axis:
-    return tree
-
-  def _op(x):
-    is_var = isinstance(x, nnx.Variable)
-    val = x.get_value() if is_var else x
-    if not hasattr(val, "shape"):
-      return x
-
-    new_val = jnp.moveaxis(val, from_axis, to_axis)
-    if not is_var:
-      return new_val
-
-    def move_fn(l):
-      while len(l) < val.ndim:
-        l.append(None)
-      if len(l) > max(from_axis, to_axis):
-        l.insert(to_axis, l.pop(from_axis))
-      return l
-
-    return nnx_update_sharding_meta(x.replace(value=new_val), move_fn)
-
-  return jax.tree.map(_op, tree, is_leaf=lambda x: isinstance(x, nnx.Variable) or hasattr(x, "shape"))
-
-
 def nnx_remove_scan_axis(tree, name="layers"):
   """Removes the given scan axis from the PartitionSpec."""
 
   def _op(x):
     if not isinstance(x, nnx.Variable):
       return x
 
+    # Scanned stacks record their own axis name, such as "dense_layers" or "moe_layers",
+    # so prefer it over the caller's default. Otherwise the name never matches and the
+    # check below strips a real logical axis instead of the scan axis.
+    axis_name = x.get_metadata().get(nnx.PARTITION_NAME, name)
+
     def remove_fn(l):
-      if name in l:
-        l.remove(name)
-      while len(l) > x.get_value().ndim:
-        l.pop(0)
+      removed = axis_name in l
+      if removed:
+        l.remove(axis_name)
+      if len(l) > x.get_value().ndim:
+        if removed:
+          raise ValueError(
+              f"Sharding names {l} still exceed value rank {x.get_value().ndim} after removing scan axis "
+              f"{axis_name!r}; the partition metadata is inconsistent."
+          )
+        raise ValueError(
+            f"Scan axis {axis_name!r} not found in sharding names {l} for a rank-{x.get_value().ndim} value; "
+            "the partition metadata is inconsistent."
+        )
       return l
 
     return nnx_update_sharding_meta(x, remove_fn)
 
   return jax.tree.map(_op, tree, is_leaf=lambda x: isinstance(x, nnx.Variable))
 
 
-def nnx_add_scan_axis(tree, name="layers", pos=0):
-  """Adds the given scan axis to the PartitionSpec at the specified position."""
+def nnx_add_and_sync_scan_axis(tree, name="layers", pos=0):
+  """Restores the scan axis on each variable's value and sharding metadata.
+
+  jax.lax.scan stacks its outputs with the scan axis at position 0. For each
+  variable this moves that axis to the variable's own param_scan_axis (falling
+  back to pos when the metadata is absent) and inserts the matching axis name at
+  the same position, so the value and its sharding metadata stay aligned.
+  """
 
   def _op(x):
     if not isinstance(x, nnx.Variable):
       return x
 
+    axis_name = x.get_metadata().get(nnx.PARTITION_NAME, name)
+    target = x.get_metadata().get("param_scan_axis", pos)
+
+    val = x.get_value()
+    if target != 0 and hasattr(val, "ndim") and val.ndim > target:
+      x = x.replace(value=jnp.moveaxis(val, 0, target))
+
     def add_fn(l):
-      if name not in l:
+      if axis_name not in l:
         while len(l) < x.get_value().ndim - 1:
           l.append(None)
-        l.insert(pos, name)
+        l.insert(target, axis_name)
       else:
         while len(l) < x.get_value().ndim:
           l.append(None)

tests/unit/maxtext_utils_nnx_test.py

@@ -203,5 +203,68 @@ def test_nnx_ensure_scan_leading_axis_mixed(self):
     self.assertEqual(broadcast_state["raw_array"].shape, (10,))
 
 
+def _make_scanned_param(shape, out_sharding, partition_name):
+  """Build an nnx.Param mirroring a scanned decoder variable.
+
+  DeepSeek stacks name their scan axis via nnx.PARTITION_NAME ("dense_layers" or
+  "moe_layers"), so the variable carries that name in its metadata rather than
+  the caller literal "layers". Eager sharding is disabled because the sliced
+  value rank is intentionally one less than the metadata length.
+  """
+  with nnx.use_eager_sharding(False):
+    return nnx.Param(jnp.zeros(shape), out_sharding=out_sharding, **{nnx.PARTITION_NAME: partition_name})
+
+
+class TestScanAxisMetadata(unittest.TestCase):
+  """Lock in per-variable scan-axis resolution for DeepSeek-style stacks."""
+
+  def test_nnx_remove_scan_axis_preserves_real_leading_axis(self):
+    """nnx_remove_scan_axis must drop the stack's own scan axis, not "embed"."""
+    # Metadata carries the real fsdp axis "embed" plus the scan axis
+    # "dense_layers"; the sliced value is rank-2, one less than the 3 names.
+    param = _make_scanned_param((4, 8), ("embed", "dense_layers", "mlp"), "dense_layers")
+    state = nnx.State({"kernel": param})
+
+    result = maxtext_utils_nnx.nnx_remove_scan_axis(state, "layers")
+
+    out_sharding = result["kernel"].get_metadata().get("out_sharding")
+    # The scan axis is gone and the real leading logical axis survives. The
+    # unfixed fallback pops index 0 and strips "embed" to ("dense_layers", "mlp").
+    self.assertEqual(tuple(out_sharding), ("embed", "mlp"))
+
+  def test_nnx_remove_scan_axis_moe_layers(self):
+    """The same resolution holds for the "moe_layers" stack axis name."""
+    param = _make_scanned_param((4, 8), ("embed", "moe_layers", "mlp"), "moe_layers")
+    state = nnx.State({"kernel": param})
+
+    result = maxtext_utils_nnx.nnx_remove_scan_axis(state, "layers")
+
+    out_sharding = result["kernel"].get_metadata().get("out_sharding")
+    self.assertEqual(tuple(out_sharding), ("embed", "mlp"))
+
+  def test_nnx_remove_scan_axis_raises_on_inconsistent_metadata(self):
+    """A rank mismatch with no matching scan axis is an error, not a silent pop."""
+    # No name in the metadata matches the resolved scan axis, and the metadata
+    # is one longer than the value rank; the fixed code raises instead of
+    # blindly popping a real axis.
+    param = _make_scanned_param((4, 8), ("embed", "mlp", "vocab"), "dense_layers")
+    state = nnx.State({"kernel": param})
+
+    with self.assertRaises(ValueError):
+      maxtext_utils_nnx.nnx_remove_scan_axis(state, "layers")
+
+  def test_nnx_add_and_sync_scan_axis_uses_partition_name(self):
+    """nnx_add_and_sync_scan_axis must insert the stack's own axis name, not "layers"."""
+    # Stacked value is rank-3; metadata holds the two sliced logical axes.
+    param = _make_scanned_param((2, 4, 8), ("embed", "mlp"), "dense_layers")
+    state = nnx.State({"kernel": param})
+
+    result = maxtext_utils_nnx.nnx_add_and_sync_scan_axis(state, "layers", 0)
+
+    out_sharding = result["kernel"].get_metadata().get("out_sharding")
+    # The unfixed code inserts the literal "layers" here.
+    self.assertEqual(tuple(out_sharding), ("dense_layers", "embed", "mlp"))
+
+
 if __name__ == "__main__":
   unittest.main()