NNX: fix MTP eval under pure_nnx (sharding rule + sown-var handling)

Two NNX-only bugs blocked MTP eval (mtp_num_layers>0 + mtp_eval_target_module)
under pure_nnx=True. Both exist on main and surface only on the NNX path (main
defaults to Linen).

1. logical_axis_rules: the MTP projection kernel is tagged with logical axis
   'concat_embed' (and the fused gated-MLP wi uses 'num_activations'), but neither
   was in logical_axis_rules. get_abstract_model resolves logical->physical via
   flax's from_sharding_rules, which returns unmapped names verbatim, so the literal
   name reached NamedSharding and was rejected. Linen replicates unmapped axes by
   default. Add both as replicated ([])  (base.yml).

2. loss_fn: mtp_losses/mtp_acceptance are nnx.Variable subclasses, not
   nnx.Intermediate, so nnx.pop(model, nnx.Intermediate) missed them. They leaked
   into the returned train state and broke out_shardings at compile, and were
   absent from intermediate_outputs so calculate_mtp_loss read None and the MTP loss
   was silently dropped. Pop them into intermediate_outputs under their collection
   names and exclude them from the returned state (train.py).

Validated on CPU (gpt3-52k, mtp_num_layers=1 mtp_eval_target_module=1): trains,
mtp_loss is non-zero in the total loss, and mtp_acceptance_rate is reported in eval.

Author: ecnal-cienet

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: