Add separate final normalization for Multi-Token Prediction in DeepSeek models.

src/maxtext/checkpoint_conversion/standalone_scripts/convert_deepseek_family_ckpt.py

@@ -178,7 +178,6 @@ def dequantize_pack_quantized_int4(
 # The MTP block in MaxText will reuse the main embedding and output head.
 MTP_KEYS_TO_SKIP = [
     "model.layers.61.embed_tokens.weight",
-    "model.layers.61.shared_head.norm.weight",
     "model.layers.61.shared_head.head.weight",
 ]
 
@@ -267,6 +266,7 @@ def hf_to_maxtext_mapping(layer_idx, num_experts, first_num_dense_layers, num_ma
             f"model.layers.{layer_idx}.enorm.weight": "mtp_block.mtp_layer_1.mtp_1_embedding_norm.scale",
             f"model.layers.{layer_idx}.hnorm.weight": "mtp_block.mtp_layer_1.mtp_1_hidden_state_norm.scale",
             f"model.layers.{layer_idx}.eh_proj.weight": "mtp_block.mtp_layer_1.mtp_1_projection.kernel",
+            f"model.layers.{layer_idx}.shared_head.norm.weight": "mtp_block.mtp_layer_1.mtp_1_final_norm.scale",
         }
     )
     for expert_idx in range(num_experts):
@@ -649,6 +649,7 @@ def _normalize(raw_key):
             "mtp_1_embedding_norm": {"scale": None},
             "mtp_1_hidden_state_norm": {"scale": None},
             "mtp_1_projection": {"kernel": None},
+            "mtp_1_final_norm": {"scale": None},
             "mtp_1_transformer_layer": {
                 "pre_self_attention_layer_norm": {"scale": None},
                 "post_self_attention_layer_norm": {"scale": None},
@@ -685,6 +686,9 @@ def _normalize(raw_key):
     jax_weights["mtp_block"]["mtp_layer_1"]["mtp_1_projection"]["kernel"] = (
         chkpt_vars["mtp_block.mtp_layer_1.mtp_1_projection.kernel"].to(torch.float16).numpy().transpose()
     )
+    jax_weights["mtp_block"]["mtp_layer_1"]["mtp_1_final_norm"]["scale"] = (
+        chkpt_vars["mtp_block.mtp_layer_1.mtp_1_final_norm.scale"].to(torch.float16).numpy()
+    )
 
     # MTP internal transformer layer - Attention and Norms
     mtp_transformer_layer = jax_weights["mtp_block"]["mtp_layer_1"]["mtp_1_transformer_layer"]

src/maxtext/checkpoint_conversion/utils/utils.py

@@ -1281,6 +1281,7 @@ def save_weights_to_checkpoint(
       save_interval_steps,
       use_ocdbt=use_ocdbt,
       use_zarr3=use_zarr3,
+      enable_continuous_checkpointing=True,
   )
   if checkpoint_manager is None:
     raise RuntimeError("Failed to create Orbax checkpoint manager.")

src/maxtext/layers/decoders.py

@@ -794,23 +794,24 @@ def _apply_embedding(
     return y
 
   @nn.compact
-  def apply_output_head(self, shared_embedding: nn.Module | nnx.Module, y, deterministic, model_mode):
+  def apply_output_head(self, shared_embedding: nn.Module | nnx.Module, y, deterministic, model_mode, normalize_y=True):
     """Applies final normalization and projects hidden states to logits."""
 
     cfg = self.config
-    if cfg.shard_mode == ShardMode.EXPLICIT:
-      norm_out_sharding = create_sharding(self.mesh, ("activation_batch", "activation_length", "activation_embed"))
-    else:
-      norm_out_sharding = None
+    if normalize_y:
+      if cfg.shard_mode == ShardMode.EXPLICIT:
+        norm_out_sharding = create_sharding(self.mesh, ("activation_batch", "activation_length", "activation_embed"))
+      else:
+        norm_out_sharding = None
 
-    y = self.get_norm_layer(num_features=y.shape[-1])(
-        dtype=cfg.dtype,
-        weight_dtype=cfg.weight_dtype,
-        name="decoder_norm",
-        epsilon=cfg.normalization_layer_epsilon,
-        kernel_axes=("norm",),
-        parameter_memory_host_offload=cfg.parameter_memory_host_offload,
-    )(y, out_sharding=norm_out_sharding)
+      y = self.get_norm_layer(num_features=y.shape[-1])(
+          dtype=cfg.dtype,
+          weight_dtype=cfg.weight_dtype,
+          name="decoder_norm",
+          epsilon=cfg.normalization_layer_epsilon,
+          kernel_axes=("norm",),
+          parameter_memory_host_offload=cfg.parameter_memory_host_offload,
+      )(y, out_sharding=norm_out_sharding)
     y = nn.Dropout(rate=cfg.dropout_rate, broadcast_dims=(-2,))(y, deterministic=deterministic)
 
     if model_mode in (MODEL_MODE_PREFILL, MODEL_MODE_AUTOREGRESSIVE):

src/maxtext/layers/multi_token_prediction.py

@@ -208,6 +208,15 @@ def __init__(
         rngs=rngs,
     )
 
+    self.final_norm = RMSNorm(
+        num_features=cfg.emb_dim,
+        epsilon=cfg.normalization_layer_epsilon,
+        dtype=cfg.dtype,
+        weight_dtype=cfg.weight_dtype,
+        kernel_axes=("norm",),
+        rngs=rngs,
+    )
+
   @property
   def embedding_norm(self):
     return getattr(self, f"mtp_{self.layer_number}_embedding_norm")
@@ -240,6 +249,14 @@ def transformer_layer(self):
   def transformer_layer(self, module):
     setattr(self, f"mtp_{self.layer_number}_transformer_layer", module)
 
+  @property
+  def final_norm(self):
+    return getattr(self, f"mtp_{self.layer_number}_final_norm")
+
+  @final_norm.setter
+  def final_norm(self, module):
+    setattr(self, f"mtp_{self.layer_number}_final_norm", module)
+
   def __call__(
       self,
       prev_hidden_state: jnp.ndarray,
@@ -462,7 +479,19 @@ def __call__(
           model_mode=self.decoder.model_mode,
       )
 
-      mtp_logits = self.decoder.apply_output_head(shared_embedding, mtp_hidden_state, deterministic, model_mode)
+      # Apply separate normalization to the MTP hidden state before projecting to logits.
+      normed_mtp_hidden_state = mtp_layer.final_norm(mtp_hidden_state)
+      normed_mtp_hidden_state = sharding.maybe_shard_with_logical(
+          normed_mtp_hidden_state,
+          ("activation_batch", "activation_length", "activation_embed"),
+          self.mesh,
+          cfg.shard_mode,
+          sharding.get_logical_axis_rules(),
+      )
+
+      mtp_logits = self.decoder.apply_output_head(
+          shared_embedding, normed_mtp_hidden_state, deterministic, model_mode, normalize_y=False
+      )
 
       logits_logical_axes = (
           "activation_embed_and_logits_batch",

src/maxtext/layers/nnx_decoders.py

@@ -1531,19 +1531,20 @@ def _apply_embedding(
 
     return y
 
-  def apply_output_head(self, shared_embedding, y, deterministic, model_mode):
+  def apply_output_head(self, shared_embedding, y, deterministic, model_mode, normalize_y=True):
     """Applies final normalization and projects hidden states to logits."""
 
     cfg = self.config
-    if cfg.shard_mode == ShardMode.EXPLICIT:
-      norm_out_sharding = create_sharding(
-          self.mesh,
-          ("activation_batch", "activation_length", "activation_embed"),
-      )
-    else:
-      norm_out_sharding = None
+    if normalize_y:
+      if cfg.shard_mode == ShardMode.EXPLICIT:
+        norm_out_sharding = create_sharding(
+            self.mesh,
+            ("activation_batch", "activation_length", "activation_embed"),
+        )
+      else:
+        norm_out_sharding = None
 
-    y = self.decoder_norm(y, out_sharding=norm_out_sharding)
+      y = self.decoder_norm(y, out_sharding=norm_out_sharding)
     y = self.dropout(y, deterministic=deterministic)  # NNX call
 
     if model_mode in {MODEL_MODE_PREFILL, MODEL_MODE_AUTOREGRESSIVE}:

tests/unit/multi_token_prediction_test.py

@@ -126,21 +126,34 @@ def test_multi_token_prediction_layer_output(self):
     max_logging.log(f"  Config Batch: {self.batch_size}, SeqLen: {self.seq_len}, EmbedDim: {self.embed_dim}")
     max_logging.log(f"  Output shape: {output_hidden_state.shape}")
 
+  def test_multi_token_prediction_layer_final_norm(self):
+    """Tests that final_norm is instantiated, accessible via property, and operational."""
+    self.assertTrue(hasattr(self.mtp_layer, "final_norm"))
+    self.assertTrue(hasattr(self.mtp_layer, f"mtp_{TEST_LAYER_NUM}_final_norm"))
+    self.assertIs(self.mtp_layer.final_norm, getattr(self.mtp_layer, f"mtp_{TEST_LAYER_NUM}_final_norm"))
+
+    norm_output = self.mtp_layer.final_norm(self.prev_hidden_state)
+    self.assertEqual(norm_output.shape, self.prev_hidden_state.shape)
+    self.assertEqual(norm_output.dtype, self.cfg.dtype)
+    self.assertFalse(jnp.isnan(norm_output).any())
+
 
 class _MockDecoderForMTP:
   """A mock decoder that simulates the behavior needed by MTPBlock."""
 
   def __init__(self, config: Config):
     self.config = config
     self.model_mode = MODEL_MODE_TRAIN
+    self.last_normalize_y = None
 
   def _apply_embedding(self, _shared_embedding, input_ids, _position_ids, _deterministic, model_mode):
     """Returns a zero tensor with the correct embedding shape."""
     batch_size, seq_len = input_ids.shape
     return jnp.zeros((batch_size, seq_len, self.config.base_emb_dim), dtype=self.config.dtype)
 
-  def apply_output_head(self, _shared_embedding, hidden_state, _deterministic, model_mode):
+  def apply_output_head(self, _shared_embedding, hidden_state, _deterministic, model_mode, normalize_y=True):
     """Returns a zero tensor with the correct logit shape."""
+    self.last_normalize_y = normalize_y
     batch_size, seq_len, _ = hidden_state.shape
     return jnp.zeros((batch_size, seq_len, self.config.vocab_size), dtype=self.config.dtype)
 
@@ -268,6 +281,26 @@ def test_sow_functionality(self):
     self.assertEqual(len(losses_val), self.cfg.mtp_num_layers)
     self.assertEqual(len(weights_val), self.cfg.mtp_num_layers)
 
+  def test_final_norm_in_mtp_block_forward(self):
+    """Verifies that MTPBlock executes final_norm and passes normalize_y=False to apply_output_head."""
+    _ = self.test_model(
+        main_hidden_state=self.main_hidden_state,
+        input_ids=self.input_ids,
+        target_ids=self.target_ids,
+        target_mask=self.target_mask,
+        position_ids=self.position_ids,
+        decoder_segment_ids=self.decoder_segment_ids,
+        model_mode=MODEL_MODE_TRAIN,
+        deterministic=True,
+    )
+    self.assertFalse(self.test_model.decoder.last_normalize_y)
+    state = nnx.state(self.test_model)
+    for k in range(1, self.cfg.mtp_num_layers + 1):
+      mtp_layer_state = getattr(state.mtp_block, f"mtp_layer_{k}")
+      self.assertTrue(hasattr(mtp_layer_state, f"mtp_{k}_final_norm"))
+      final_norm_scale = getattr(mtp_layer_state, f"mtp_{k}_final_norm").scale.value
+      self.assertEqual(final_norm_scale.shape, (self.cfg.base_emb_dim,))
+
   def _forward(self, model):
     return model(
         main_hidden_state=self.main_hidden_state,