Merge pull request #4133 from AI-Hypercomputer:feat/nnx-e2e-fixes

PiperOrigin-RevId: 931227563

Author: Google-ML-Automation

src/maxtext/trainers/pre_train/train.py

@@ -45,7 +45,7 @@
 # Placeholder: internal
 
 # pylint: disable=too-many-positional-arguments
-from maxtext.layers.multi_token_prediction import calculate_mtp_acceptance_rate, calculate_mtp_loss
+from maxtext.layers.multi_token_prediction import calculate_mtp_acceptance_rate, calculate_mtp_loss, mtp_acceptance, mtp_losses
 from maxtext.common import checkpointing, profiler
 from maxtext.common.goodput import (
     GoodputEvent,
@@ -197,6 +197,16 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr
     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()
+
     if config.num_vocab_tiling > 1:
       hidden_state_key = ("decoder", "hidden_states")
       hidden_states = maxtext_utils.get_nested_value(intermediate_outputs, hidden_state_key)[0]
@@ -532,9 +542,10 @@ def move(path, value):
 
   if isinstance(model, nn.Module):
     return new_state, metrics
-  # Drop Intermediates (e.g. sowed max_logits for QK-Clip) before returning;
-  # they're absent from state_mesh_shardings and would cause a leaf-count mismatch.
-  return nnx.state(new_state, nnx.Not(nnx.Intermediate)), metrics
+  # Drop Intermediates (e.g. sowed max_logits for QK-Clip) and the MTP sown
+  # vars (mtp_losses/mtp_acceptance) before returning. They're absent from
+  # state_mesh_shardings and would cause a leaf-count / structure mismatch.
+  return nnx.state(new_state, nnx.Not(nnx.Any(nnx.Intermediate, mtp_losses, mtp_acceptance))), metrics
 
 
 def eval_step(model, config, state, data, dropout_rng=None):

src/maxtext/utils/maxtext_utils.py

@@ -1670,7 +1670,23 @@ def _make_named_sharding(v):
       local_rules = metadata.get("sharding_rules", ())
       if context_rules or local_rules:
         rules = composite_rules(context_rules, local_rules)
-        pspec = PartitionSpec(*from_sharding_rules(out_sharding, rules))
+        raw_sharding = from_sharding_rules(out_sharding, rules)
+        mesh_axis_names = mesh.axis_names if mesh is not None else ()
+
+        # from_sharding_rules leaves a logical name with no matching rule unchanged, so a
+        # name missing from logical_axis_rules (e.g. concat_embed on the MTP kernel)
+        # reaches NamedSharding and is rejected as an unknown mesh axis. Map any such
+        # leftover name to None (replicated), matching Linen, whose logical_to_mesh_axes
+        # replicates unmatched names.
+        def _sanitize(x):
+          if isinstance(x, list):
+            x = tuple(x)
+          if x is None or (isinstance(x, str) and x in mesh_axis_names) or isinstance(x, tuple):
+            return x
+          return None
+
+        sanitized_sharding = [_sanitize(x) for x in raw_sharding]
+        pspec = PartitionSpec(*sanitized_sharding)
       else:
         pspec = PartitionSpec(*out_sharding)
     return v.replace(NamedSharding(mesh, pspec))

src/maxtext/utils/train_utils.py

@@ -278,8 +278,12 @@ def create_train_state_fn():
     )
     if config.pure_nnx:
       with nn_partitioning.axis_rules(config.logical_axis_rules):
-        # train_state is instance of TrainStateNNX
-        state_graphdef, _ = nnx.get_abstract_model(init_state_fn, mesh)
+        # We only need the graphdef here; it's merged with state below. Avoid
+        # nnx.get_abstract_model: it eagerly builds a NamedSharding for every variable
+        # under jax.set_mesh(mesh) and rejects any logical name missing from
+        # logical_axis_rules (e.g. concat_embed on the MTP kernel). Tracing shapes
+        # without a mesh skips sharding resolution, so it avoids the crash.
+        state_graphdef = nnx.graphdef(nnx.eval_shape(init_state_fn))
         _, state_params, _ = nnx.split(state.model, nnx.Param, ...)
         _, state_mesh_shardings_params, _ = nnx.split(state_mesh_shardings.model, nnx.Param, ...)
     else: