[NNX] Clean up transient intermediates at setup and unify metrics sowing as nnx.Intermediate

This change cleans up transient variables sowed during model creation/tracing from the persistent sharding state and checkpoint layout.

1. Clean Quantization Setup:
- Updated maybe_quantize_model in quantizations.py to pop sowed nnx.Intermediate variables in-place immediately after Qwix tracing.
- This cleans the initial state returned by the model factory fn, allowing us to revert the setup-time nnx.Not(nnx.Intermediate) filters in maxtext_utils.py and train_compile.py.
- Optimized create_nnx_abstract_model in model_creation_utils.py to reuse abs_model.mesh and avoid calling eval_shape twice.

2. Metrics Sowing Uniformity:
- Changed all self.sow("intermediates", ...) calls to self.sow(nnx.Intermediate, ...) in all pure NNX models (gemma.py, llama2.py, etc.).
- Sowing string-based "intermediates" was dynamically creating custom Variable types that did not inherit from nnx.Intermediate, leading to parameter/checkpoint bloat.
- train_step in train.py now strips all intermediates via nnx.Not(nnx.Intermediate) before returning.

3. Testing:
- Added MaybeQuantizeModelTest in quantizations_test.py to assert that sowed intermediates are popped from the model state and that abstract model state contains no intermediates.

Author: xibinliu

src/maxtext/layers/multi_token_prediction.py

@@ -37,11 +37,11 @@
 # Custom Variable types for MTP intermediate outputs
 # These will be automatically converted to Linen mutable collections by ToLinen wrapper
 # The class names become collection names directly (no case conversion)
-class mtp_losses(nnx.Variable):  # pylint: disable=invalid-name
+class mtp_losses(nnx.Intermediate):  # pylint: disable=invalid-name
   """Variable type for storing MTP loss components -> 'mtp_losses' collection."""
 
 
-class mtp_acceptance(nnx.Variable):  # pylint: disable=invalid-name
+class mtp_acceptance(nnx.Intermediate):  # pylint: disable=invalid-name
   """Variable type for storing MTP acceptance predictions -> 'mtp_acceptance' collection."""
 
 

src/maxtext/layers/quantizations.py

@@ -862,6 +862,10 @@ def maybe_quantize_model(model, config):
             dummy_segment_ids,
             enable_dropout=False,
         )
+        # 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
+        # between the initial setup GraphDef/state_mesh_shardings and the stripped states during train steps.
+        nnx.pop(model, nnx.Intermediate)
       else:
         model = qwix.quantize_model(model, quantization_provider)
   return model

src/maxtext/models/gemma.py

@@ -171,10 +171,10 @@ def __call__(
     )
 
     if self.config.record_internal_nn_metrics:
-      self.sow("intermediates", "activation_mean", jnp.mean(layer_output))
-      self.sow("intermediates", "activation_stdev", jnp.std(layer_output))
+      self.sow(nnx.Intermediate, "activation_mean", jnp.mean(layer_output))
+      self.sow(nnx.Intermediate, "activation_stdev", jnp.std(layer_output))
       self.sow(
-          "intermediates",
+          nnx.Intermediate,
           "activation_fraction_zero",
           jnp.sum(layer_output == 0) / jnp.size(layer_output),
       )

src/maxtext/models/gemma2.py

@@ -307,10 +307,10 @@ def __call__(
     layer_output = nn.with_logical_constraint(layer_output, self.activation_axis_names)
 
     if self.config.record_internal_nn_metrics:
-      self.sow("intermediates", "activation_mean", jnp.mean(layer_output))
-      self.sow("intermediates", "activation_stdev", jnp.std(layer_output))
+      self.sow(nnx.Intermediate, "activation_mean", jnp.mean(layer_output))
+      self.sow(nnx.Intermediate, "activation_stdev", jnp.std(layer_output))
       self.sow(
-          "intermediates",
+          nnx.Intermediate,
           "activation_fraction_zero",
           jnp.sum(layer_output == 0) / jnp.size(layer_output),
       )

src/maxtext/models/gemma3.py

@@ -242,10 +242,10 @@ def __call__(
     layer_output = nn.with_logical_constraint(layer_output, self.activation_axis_names)
 
     if cfg.record_internal_nn_metrics:
-      self.sow("intermediates", "activation_mean", jnp.mean(layer_output))
-      self.sow("intermediates", "activation_stdev", jnp.std(layer_output))
+      self.sow(nnx.Intermediate, "activation_mean", jnp.mean(layer_output))
+      self.sow(nnx.Intermediate, "activation_stdev", jnp.std(layer_output))
       self.sow(
-          "intermediates",
+          nnx.Intermediate,
           "activation_fraction_zero",
           jnp.sum(layer_output == 0) / jnp.size(layer_output),
       )

src/maxtext/models/gemma4.py

@@ -365,7 +365,7 @@ def __call__(
     if getattr(self.config, "num_experts", 1) > 1:
       mlp_lnx, load_balance_loss, _ = self.mlp(attn_output, original_inputs=attention_lnx)
       if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None:
-        self.sow("intermediates", "moe_lb_loss", load_balance_loss)
+        self.sow(nnx.Intermediate, "moe_lb_loss", load_balance_loss)
     else:
       mlp_lnx = self.mlp(attn_output, deterministic=deterministic)
 
@@ -381,10 +381,10 @@ def __call__(
     layer_output = nn.with_logical_constraint(layer_output, self.activation_axis_names)
 
     if cfg.record_internal_nn_metrics:
-      self.sow("intermediates", "activation_mean", jnp.mean(layer_output))
-      self.sow("intermediates", "activation_stdev", jnp.std(layer_output))
+      self.sow(nnx.Intermediate, "activation_mean", jnp.mean(layer_output))
+      self.sow(nnx.Intermediate, "activation_stdev", jnp.std(layer_output))
       self.sow(
-          "intermediates",
+          nnx.Intermediate,
           "activation_fraction_zero",
           jnp.sum(layer_output == 0) / jnp.size(layer_output),
       )

src/maxtext/models/gpt3.py

@@ -510,10 +510,10 @@ def __call__(
     layer_output = nn.with_logical_constraint(layer_output, self.activation_axis_names)
 
     if self.config.record_internal_nn_metrics:
-      self.sow("intermediates", "activation_mean", jnp.mean(layer_output))
-      self.sow("intermediates", "activation_stdev", jnp.std(layer_output))
+      self.sow(nnx.Intermediate, "activation_mean", jnp.mean(layer_output))
+      self.sow(nnx.Intermediate, "activation_stdev", jnp.std(layer_output))
       self.sow(
-          "intermediates",
+          nnx.Intermediate,
           "activation_fraction_zero",
           jnp.sum(layer_output == 0) / jnp.size(layer_output),
       )

src/maxtext/models/gpt_oss.py

@@ -198,13 +198,13 @@ def __call__(
     )
 
     if cfg.load_balance_loss_weight > 0.0 and load_balance_loss is not None:
-      self.sow("intermediates", "moe_lb_loss", load_balance_loss)
+      self.sow(nnx.Intermediate, "moe_lb_loss", load_balance_loss)
 
     if cfg.record_internal_nn_metrics:
-      self.sow("intermediates", "activation_mean", jnp.mean(layer_output))
-      self.sow("intermediates", "activation_stdev", jnp.std(layer_output))
+      self.sow(nnx.Intermediate, "activation_mean", jnp.mean(layer_output))
+      self.sow(nnx.Intermediate, "activation_stdev", jnp.std(layer_output))
       self.sow(
-          "intermediates",
+          nnx.Intermediate,
           "activation_fraction_zero",
           jnp.sum(layer_output == 0) / jnp.size(layer_output),
       )

src/maxtext/models/llama2.py

@@ -203,10 +203,10 @@ def __call__(
     layer_output = self._maybe_shard_with_logical(layer_output, self.activation_axis_names)
 
     if cfg.record_internal_nn_metrics:
-      self.sow("intermediates", "activation_mean", jnp.mean(layer_output))
-      self.sow("intermediates", "activation_stdev", jnp.std(layer_output))
+      self.sow(nnx.Intermediate, "activation_mean", jnp.mean(layer_output))
+      self.sow(nnx.Intermediate, "activation_stdev", jnp.std(layer_output))
       self.sow(
-          "intermediates",
+          nnx.Intermediate,
           "activation_fraction_zero",
           jnp.sum(layer_output == 0) / jnp.size(layer_output),
       )

src/maxtext/models/llama4.py

@@ -497,13 +497,13 @@ def __call__(
     layer_output = nn.with_logical_constraint(layer_output, self.activation_axis_names)
 
     if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None:
-      self.sow("intermediates", "moe_lb_loss", load_balance_loss)
+      self.sow(nnx.Intermediate, "moe_lb_loss", load_balance_loss)
 
     if cfg.record_internal_nn_metrics:
-      self.sow("intermediates", "activation_mean", jnp.mean(layer_output))
-      self.sow("intermediates", "activation_stdev", jnp.std(layer_output))
+      self.sow(nnx.Intermediate, "activation_mean", jnp.mean(layer_output))
+      self.sow(nnx.Intermediate, "activation_stdev", jnp.std(layer_output))
       self.sow(
-          "intermediates",
+          nnx.Intermediate,
           "activation_fraction_zero",
           jnp.sum(layer_output == 0) / jnp.size(layer_output),
       )