[Feature][Spec Decode] MTP with separate (possibly quantized) lm head for nemotron (#54574)

Signed-off-by: Yoav Miron <yomiron@nvidia.com>
Signed-off-by: YoavMiron <yomiron@nvidia.com>
Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com>

vllm/config/speculative.py

@@ -620,16 +620,15 @@ class SpeculativeConfig:
         )
         factors.append(uses_aux_hidden_states)
 
-        if uses_aux_hidden_states and self.draft_model_config is not None:
+        if self.draft_model_config is not None:
             factors.append(self.draft_model_config.compute_hash())
-
             # The specific layers used also affect the computation graph.
             layer_ids = getattr(
                 self.draft_model_config.hf_config,
                 "eagle_aux_hidden_state_layer_ids",
                 None,
             )
-            if layer_ids is not None:
+            if layer_ids is not None and uses_aux_hidden_states:
                 # Convert to tuple to make it hashable
                 factors.append(tuple(layer_ids))
 

vllm/model_executor/models/nemotron_h_mtp.py

@@ -22,16 +22,20 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader,
+    maybe_remap_kv_scale_name,
+)
 from vllm.model_executor.models.utils import (
     WeightsMapper,
+    get_draft_quant_config,
     make_empty_intermediate_tensors_factory,
     maybe_prefix,
 )
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.configs.nemotron_h import NemotronHConfig
 
-from .interfaces import SupportsPP
+from .interfaces import SupportsPP, SupportsQuant
 from .nemotron_h import (
     NemotronHAttentionDecoderLayer,
     NemotronHMoEDecoderLayer,
@@ -214,10 +218,22 @@ class NemotronHMTPMoEDecoderLayer(NemotronHMoEDecoderLayer):
 class NemotronHMultiTokenPredictor(nn.Module):
     """MTP predictor with NemotronH layers."""
 
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
         super().__init__()
 
-        config = vllm_config.model_config.hf_config.get_text_config()
+        speculative_config = vllm_config.speculative_config
+        assert speculative_config is not None
+        draft_model_config = speculative_config.draft_model_config
+        assert draft_model_config is not None
+        config = draft_model_config.hf_config.get_text_config()
+        if quant_config is None:
+            quant_config = get_draft_quant_config(vllm_config)
 
         self.config = config
         self.vocab_size = config.vocab_size
@@ -257,9 +273,9 @@ class NemotronHMultiTokenPredictor(nn.Module):
             common_kwargs = dict(
                 config=config,
                 layer_idx=self.mtp_start_layer_idx + i,
-                model_config=vllm_config.model_config,
+                model_config=draft_model_config,
                 cache_config=vllm_config.cache_config,
-                quant_config=vllm_config.quant_config,
+                quant_config=quant_config,
                 parallel_config=vllm_config.parallel_config,
                 prefix=layer_prefix,
                 has_start_projections=is_start_of_step,
@@ -310,14 +326,9 @@ class NemotronHMultiTokenPredictor(nn.Module):
         return hidden_states
 
 
-class NemotronHMTP(nn.Module, SupportsPP):
+class NemotronHMTP(nn.Module, SupportsPP, SupportsQuant):
     """NemotronH MTP model."""
 
-    # Quant configs name modules in checkpoint space ("language_model.mtp.layers.0*"),
-    # but this draft is built under "mtp" (maybe_prefix below). SupportsQuant only
-    # re-roots exclude_modules when the model defines a mapper, so without one the
-    # exclusions never match and the MTP experts are wrongly quantized. Mirrors the
-    # prefix handling load_weights() already does by hand.
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={"language_model.": ""},
         orig_to_new_substr={"embeddings": "embed_tokens"},
@@ -333,11 +344,13 @@ class NemotronHMTP(nn.Module, SupportsPP):
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
-        config = vllm_config.model_config.hf_config.get_text_config()
+        speculative_config = vllm_config.speculative_config
+        assert speculative_config is not None
+        draft_model_config = speculative_config.draft_model_config
+        assert draft_model_config is not None
+        config = draft_model_config.hf_config.get_text_config()
         self.vllm_config = vllm_config
         self.config = config
-        self.quant_config = vllm_config.quant_config
-
         # Needed for load_weights mapping
         self.mtp_start_layer_idx = config.num_hidden_layers
 
@@ -350,13 +363,17 @@ class NemotronHMTP(nn.Module, SupportsPP):
 
         # MTP predictor
         self.model = NemotronHMultiTokenPredictor(
-            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "mtp")
+            vllm_config=vllm_config,
+            quant_config=self.quant_config,
+            prefix=maybe_prefix(prefix, "mtp"),
         )
 
         # LM head for generating logits
+        self.has_own_lm_head = False
         self.lm_head = ParallelLMHead(
             self.config.vocab_size,
             self.config.hidden_size,
+            quant_config=self.quant_config,
             prefix=maybe_prefix(prefix, "lm_head"),
         )
 
@@ -366,6 +383,11 @@ class NemotronHMTP(nn.Module, SupportsPP):
             self.model.make_empty_intermediate_tensors
         )
 
+    @staticmethod
+    def _find_quant_config(*args, **kwargs) -> QuantizationConfig | None:
+        vllm_config = kwargs.get("vllm_config")
+        return get_draft_quant_config(vllm_config)
+
     def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.model.get_input_embeddings(input_ids)
 
@@ -428,9 +450,16 @@ class NemotronHMTP(nn.Module, SupportsPP):
             # MTP weights are nested in "language_model."
             # in Multimodal Nemotron-H checkpoints.
             name = name.removeprefix("language_model.")
-
-            # Only process MTP weights - skip all non-MTP weights
-            if not name.startswith("mtp.") and "embeddings" not in name:
+            is_lm_head_weight = name.startswith("lm_head.")
+            if is_lm_head_weight:
+                self.has_own_lm_head = True
+            # Only process MTP and LM head weights -
+            # skip all non-MTP and non-LM head weights
+            if (
+                not name.startswith("mtp.")
+                and "embeddings" not in name
+                and not is_lm_head_weight
+            ):
                 continue
             # Skip rotary embeddings (computed, not loaded)
             if "rotary_emb.inv_freq" in name:
@@ -443,6 +472,11 @@ class NemotronHMTP(nn.Module, SupportsPP):
                 if name.startswith("backbone."):
                     name = name.replace("backbone.", "model.")
 
+            if "scale" in name or "zero_point" in name:
+                name = maybe_remap_kv_scale_name(name, params_dict)
+                if name is None:
+                    continue
+
             # Handle stacked parameters (qkv_proj) for attention layers
             is_stacked = False
             for param_name, weight_name, shard_id in stacked_params_mapping:
@@ -516,4 +550,9 @@ class NemotronHMTP(nn.Module, SupportsPP):
             weight_loader(param, loaded_weight)
             loaded_params.add(name)
 
+        if not self.has_own_lm_head:
+            loaded_params.update(
+                name for name in params_dict if name.startswith("lm_head.")
+            )
+
         return loaded_params