Merge pull request #4525 from AI-Hypercomputer:fix/nnx-mtp-and-scan-axis-sharding

PiperOrigin-RevId: 951564354

Author: Google-ML-Automation

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()