diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 39d1f0ab4b..d4085d7d52 100644 --- a/vllm/config/speculative.py +++ b/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)) diff --git a/vllm/model_executor/models/nemotron_h_mtp.py b/vllm/model_executor/models/nemotron_h_mtp.py index 6db6868bda..96f1b6685e 100644 --- a/vllm/model_executor/models/nemotron_h_mtp.py +++ b/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