diff --git a/src/maxtext/checkpoint_conversion/standalone_scripts/convert_deepseek_family_ckpt.py b/src/maxtext/checkpoint_conversion/standalone_scripts/convert_deepseek_family_ckpt.py index 6a750bace4..caa24a2a96 100644 --- a/src/maxtext/checkpoint_conversion/standalone_scripts/convert_deepseek_family_ckpt.py +++ b/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"] diff --git a/src/maxtext/checkpoint_conversion/utils/utils.py b/src/maxtext/checkpoint_conversion/utils/utils.py index 17209a037b..8777214cdf 100644 --- a/src/maxtext/checkpoint_conversion/utils/utils.py +++ b/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.") diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index a9bacf692f..a6c573b08b 100644 --- a/src/maxtext/layers/decoders.py +++ b/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): diff --git a/src/maxtext/layers/multi_token_prediction.py b/src/maxtext/layers/multi_token_prediction.py index fce3dfb940..befa2b2730 100644 --- a/src/maxtext/layers/multi_token_prediction.py +++ b/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", diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 1e06020e72..617745c1d3 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/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}: diff --git a/tests/unit/multi_token_prediction_test.py b/tests/unit/multi_token_prediction_test.py index 93c8ef6f81..1b2fae8fdf 100644 --- a/tests/unit/multi_token_prediction_test.py +++ b/tests/unit/multi_token_prediction_test.py @@ -126,6 +126,17 @@ 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.""" @@ -133,14 +144,16 @@ class _MockDecoderForMTP: 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,