Merge pull request #3130 from AI-Hypercomputer:bvandermoon-mtp-params

PiperOrigin-RevId: 871486617

Author: Google-ML-Automation

src/MaxText/layers/multi_token_prediction.py

@@ -128,6 +128,40 @@ def __init__(
         model_mode=MODEL_MODE_TRAIN,
     )
 
+
+  @property
+  def embedding_norm(self):
+    return getattr(self, f"mtp_{self.layer_number}_embedding_norm")
+
+  @embedding_norm.setter
+  def embedding_norm(self, module):
+    setattr(self, f"mtp_{self.layer_number}_embedding_norm", module)
+
+  @property
+  def hidden_state_norm(self):
+    return getattr(self, f"mtp_{self.layer_number}_hidden_state_norm")
+
+  @hidden_state_norm.setter
+  def hidden_state_norm(self, module):
+    setattr(self, f"mtp_{self.layer_number}_hidden_state_norm", module)
+
+  @property
+  def projection_layer(self):
+    return getattr(self, f"mtp_{self.layer_number}_projection")
+
+  @projection_layer.setter
+  def projection_layer(self, module):
+    setattr(self, f"mtp_{self.layer_number}_projection", module)
+
+  @property
+  def transformer_layer(self):
+    return getattr(self, f"mtp_{self.layer_number}_transformer_layer")
+
+  @transformer_layer.setter
+  def transformer_layer(self, module):
+    setattr(self, f"mtp_{self.layer_number}_transformer_layer", module)
+
+
   def __call__(
       self,
       prev_hidden_state: jnp.ndarray,
@@ -192,13 +226,6 @@ def __init__(
     self.decoder = decoder
     self.rngs = rngs if rngs is not None else nnx.Rngs(0)
 
-    # NNX Variables are exposed as Linen mutable collections by ToLinen wrapper.
-    self.losses = mtp_losses(jnp.zeros((config.mtp_num_layers,), dtype=jnp.float32))
-    self.weights = mtp_losses(jnp.zeros((config.mtp_num_layers,), dtype=jnp.float32))
-    # Float32 used to avoid gradient errors; converted to int32 in acceptance rate calculation.
-    self.mtp_preds = mtp_acceptance(jnp.zeros((1,), dtype=jnp.float32))
-    self.mtp_mask = mtp_acceptance(jnp.zeros((1,), dtype=jnp.float32))
-
     # 1-indexed to match paper convention.
     for k in range(1, config.mtp_num_layers + 1):
       layer = MultiTokenPredictionLayer(
@@ -278,11 +305,13 @@ def __call__(
         mtp_masks_list.append(rolled_target_mask)
 
     if mtp_losses_list:
-      self.losses.value = jnp.stack(mtp_losses_list)
-      self.weights.value = jnp.stack(mtp_weights_list)
+      # Not part of checkpoints, don't declare in __init__
+      self.losses = mtp_losses(jnp.stack(mtp_losses_list))
+      self.weights = mtp_losses(jnp.stack(mtp_weights_list))
     if mtp_preds_list:
-      self.mtp_preds.value = jnp.stack(mtp_preds_list)
-      self.mtp_mask.value = jnp.stack(mtp_masks_list)
+      # Not part of checkpoints, don't declare in __init__
+      self.mtp_preds = mtp_acceptance(jnp.stack(mtp_preds_list))
+      self.mtp_mask = mtp_acceptance(jnp.stack(mtp_masks_list))
 
     return {}
 

tests/unit/multi_token_prediction_test.py

@@ -231,22 +231,6 @@ def setUp(self):
         rngs=self.rngs,
     )
 
-  def test_no_sow_during_init(self):
-    """Verifies losses/weights are initialized with zeros (NNX behavior)."""
-    # NNX pre-initializes Variables with zeros to avoid checkpointing errors.
-    # Unlike Linen which sows during forward pass, NNX creates Variables in __init__.
-    initial_state = nnx.state(self.test_model)
-    self.assertTrue(hasattr(initial_state.mtp_block, "losses"))
-    self.assertTrue(hasattr(initial_state.mtp_block, "weights"))
-
-    # Verify they're initialized with zeros of correct shape.
-    losses_val = initial_state.mtp_block.losses.value
-    weights_val = initial_state.mtp_block.weights.value
-    self.assertEqual(losses_val.shape, (self.cfg.mtp_num_layers,))
-    self.assertEqual(weights_val.shape, (self.cfg.mtp_num_layers,))
-    self.assertTrue(jnp.all(losses_val == 0.0))
-    self.assertTrue(jnp.all(weights_val == 0.0))
-
   def test_sow_functionality(self):
     """Verifies that the block correctly sows losses and weights."""
     _ = self.test_model(