diff --git a/tests/v1/e2e/spec_decode/eagle/test_eagle3_pp.py b/tests/v1/e2e/spec_decode/eagle/test_eagle3_pp.py index ddd7dbb910b4..77ee85f08555 100644 --- a/tests/v1/e2e/spec_decode/eagle/test_eagle3_pp.py +++ b/tests/v1/e2e/spec_decode/eagle/test_eagle3_pp.py @@ -14,12 +14,15 @@ MODEL = "meta-llama/Llama-3.2-1B-Instruct" DRAFT = "nm-testing/Llama3_2_1B_speculator.eagle3" -PROMPTS = [ - "The capital of France is", - "2 + 2 equals", - "In one word, the color of the sky is", - "Q: If a train travels 60 miles in 1.5 hours, what is its average speed?\nA:", +# Enough tokens that one greedy path changing under compilation can't move the mean +TOPICS = ["the ocean", "a robot", "Rome", "volcanoes"] +TEMPLATES = [ + "Write a short story about {}.", + "Explain {} to a child.", + "List five facts about {}:", + "Q: What is the history of {}?\nA:", ] +PROMPTS = [template.format(topic) for topic in TOPICS for template in TEMPLATES] ACCEPTANCE_TOLERANCE = 0.95 @@ -29,11 +32,13 @@ def _run( model: str = MODEL, draft: str = DRAFT, max_model_len: int = 512, - max_tokens: int = 32, + max_tokens: int = 64, chat_prompts: list[Messages] | None = None, + model_impl: str = "auto", ) -> float: llm = LLM( model=model, + model_impl=model_impl, tensor_parallel_size=1, pipeline_parallel_size=pp_size, max_model_len=max_model_len, @@ -64,10 +69,11 @@ def _run( @multi_gpu_test(num_gpus=4) -def test_eagle3_pipeline_parallel_acceptance(): - baseline = _run(1) +@pytest.mark.parametrize("model_impl", ["vllm", "transformers"]) +def test_eagle3_pipeline_parallel_acceptance(model_impl: str): + baseline = _run(1, model_impl=model_impl) for pp_size in (2, 4): - parallel = _run(pp_size) + parallel = _run(pp_size, model_impl=model_impl) assert parallel >= baseline * ACCEPTANCE_TOLERANCE, ( f"PP={pp_size} acceptance regressed: {parallel:.3f} < " f"{baseline:.3f} * {ACCEPTANCE_TOLERANCE}" diff --git a/tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py b/tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py index db8272887fcb..a18673235bbc 100644 --- a/tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py +++ b/tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py @@ -358,3 +358,40 @@ def test_extract_hidden_states_tp2(): len(layer_ids), hidden_size, ) + + +@pytest.mark.parametrize("model_impl", ["vllm", "transformers"]) +@create_new_process_for_each_test() +def test_extract_layer_zero_is_input_embeddings(tmp_path, model_impl): + """Layer 0 is the input to the first decoder layer, i.e. the input embeddings.""" + from transformers import AutoModelForCausalLM + + model = "meta-llama/Llama-3.2-1B-Instruct" + llm = LLM( + model=model, + model_impl=model_impl, + speculative_config={ + "method": "extract_hidden_states", + "num_speculative_tokens": 1, + "draft_model_config": { + "hf_config": {"eagle_aux_hidden_state_layer_ids": [0, 1]} + }, + }, + kv_transfer_config={ + "kv_connector": "ExampleHiddenStatesConnector", + "kv_role": "kv_producer", + "kv_connector_extra_config": {"shared_storage_path": str(tmp_path)}, + }, + max_model_len=256, + gpu_memory_utilization=0.4, + ) + hidden_size = llm.llm_engine.model_config.get_hidden_size() + sampling_params = SamplingParams(max_tokens=1, temperature=0.0) + output = llm.generate(["Hello world"], sampling_params)[0] + expected_shape = (len(output.prompt_token_ids), 2, hidden_size) + token_ids, hidden_states = get_and_check_output(output, expected_shape) + + hf_model = AutoModelForCausalLM.from_pretrained(model, dtype=hidden_states.dtype) + with torch.no_grad(): + embeddings = hf_model.get_input_embeddings()(token_ids) + torch.testing.assert_close(hidden_states[:, 0, :], embeddings) diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 49ecc6ba0f00..259c1a26725e 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -1667,7 +1667,11 @@ def _maybe_add_hidden_state( layer_idx >= self._aux_first_local_layer and layer_idx in self.aux_hidden_state_layers ): - value = hidden_states + residual if residual is not None else hidden_states + if residual is None: + # The first layer updates its residual (the embeddings) in place + value = hidden_states.clone() + else: + value = hidden_states + residual aux_hidden_states.append(value) return aux_hidden_states @@ -1721,6 +1725,9 @@ def _get_decoder_holder(model: object) -> object: parent_ref = model.get_language_model() elif hasattr(model, "language_model"): parent_ref = model.language_model + # The Transformers modeling backend's `.model` is the HF model, not the decoder + if isinstance(parent_ref, EagleModelMixin): + return parent_ref # A multimodal model that builds its decoder inside `_mark_language_model` # has get_language_model() return the decoder itself, with no `.model`. return getattr(parent_ref, "model", parent_ref) diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py index 222bc0be0a5c..484dd8e45883 100644 --- a/vllm/model_executor/models/transformers/base.py +++ b/vllm/model_executor/models/transformers/base.py @@ -51,6 +51,7 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.fused_moe import MoERunner from vllm.model_executor.models.interfaces import ( + EagleModelMixin, SupportsEagle, SupportsEagle3, SupportsLoRA, @@ -84,6 +85,7 @@ WeightsMapper, make_empty_intermediate_tensors_factory, maybe_prefix, + spec_decode_needs_target_embed, ) from vllm.sequence import IntermediateTensors @@ -114,8 +116,10 @@ class Base( SupportsPP, SupportsEagle, SupportsEagle3, + EagleModelMixin, ): hf_to_vllm_mapper: WeightsMapper + supports_aux_hidden_states_over_pp = True def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): nn.Module.__init__(self) @@ -188,6 +192,7 @@ def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): # Create weight name to module qualname mapper self._create_hf_to_vllm_mapper() # Remove layers not on this pipeline parallel rank + self.start_layer, self.end_layer = self._pp_indices() self.pipeline_parallel() # Skip the device sync of rope updates where `inv_freq` can be fixed maybe_disable_dynamic_rope_update(self.model.get_decoder(), vllm_config) @@ -432,8 +437,10 @@ def pipeline_parallel(self): # Layers before module list for name in names[:module_list_idx]: - if self.pp_group.is_first_rank or ( - self._get_tie_word_embeddings() and self.pp_group.is_last_rank + if ( + self.pp_group.is_first_rank + or (self._get_tie_word_embeddings() and self.pp_group.is_last_rank) + or spec_decode_needs_target_embed(self.vllm_config) ): continue # attrsetter in case the module is nested (e.g. "text_model.embed_tokens") @@ -551,7 +558,8 @@ def _recursive_replace(module: nn.Module, prefix: str): and len(module) == self.text_config.num_hidden_layers ): # Populate Eagle3 attrs - self._target_class = type(child_module) + if not isinstance(child_module, PPMissingLayer): + self._target_class = type(child_module) layer_name = qual_name.removeprefix("model.") self._layer_names[int(child_name)] = layer_name # MTP weights should not be loaded into the base model @@ -790,6 +798,7 @@ def forward( assert intermediate_tensors is not None input_ids = None inputs_embeds = intermediate_tensors["hidden_states"] + remote_aux = self.collect_remote_aux_hidden_states(intermediate_tensors) # Add batch dimension before entering Transformers model if input_ids is not None and input_ids.ndim == 1: @@ -818,13 +827,24 @@ def forward( # Remove batch dimension after exiting Transformers model hidden_states = outputs[0][0, ...] + aux_hidden_states = [] if self._output_aux_hidden_states_kwargs: - aux_hidden_states = [x[0][0, ...] for x in outputs[1:]] + # Layers on other PP ranks capture nothing + aux_hidden_states = [x[0][0, ...] for x in outputs[1:] if x] + # The input embedding isn't called when the caller passes `inputs_embeds` + caller_embeds = inputs_embeds if self.pp_group.is_first_rank else None + if 0 in self.aux_hidden_state_layers and caller_embeds is not None: + aux_hidden_states.insert(0, caller_embeds[0, ...]) if not self.pp_group.is_last_rank: - return IntermediateTensors({"hidden_states": hidden_states}) + return IntermediateTensors( + { + "hidden_states": hidden_states, + **self.pack_local_aux_hidden_states(aux_hidden_states), + } + ) - if self._output_aux_hidden_states_kwargs and len(aux_hidden_states) > 0: + if aux_hidden_states := remote_aux + aux_hidden_states: return hidden_states, aux_hidden_states return hidden_states @@ -856,10 +876,17 @@ def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: if self.model._can_record_outputs is None: self.model._can_record_outputs = {} - target_class = self._target_class - for layer in layers: - # layer - 1 because we want the input to the layer - layer_name = self._layer_names[layer - 1] + self._set_aux_hidden_state_layers(layers) + for layer in self.aux_hidden_state_layers: + # The input to a layer is the output of the module before it + if layer == 0: + embeddings = self.model.get_input_embeddings() + target_class = type(embeddings) + modules = self.model.named_modules() + layer_name = next(n for n, m in modules if m is embeddings) + else: + target_class = self._target_class + layer_name = self._layer_names[layer - 1] layer_key = f"aux_hidden_state_{layer}" aux_hidden_state_i = OutputRecorder(target_class, layer_name=layer_name) self.model._can_record_outputs[layer_key] = aux_hidden_state_i diff --git a/vllm/model_executor/models/transformers/multimodal.py b/vllm/model_executor/models/transformers/multimodal.py index 1b0213c21292..ed72366185c0 100644 --- a/vllm/model_executor/models/transformers/multimodal.py +++ b/vllm/model_executor/models/transformers/multimodal.py @@ -1083,6 +1083,9 @@ def get_language_model(self) -> torch.nn.Module: bases = [b for b in bases if b is not MultiModalMixin] class LanguageModel(*bases): # type: ignore[misc] + # `reserve_aux_intermediate_tensor_slots` needs `config.hidden_size` + supports_aux_hidden_states_over_pp = False + def __init__(self, multimodal_model): # Don't call super().__init__() to avoid re-initialization self.__dict__.update(multimodal_model.__dict__) diff --git a/vllm/v1/worker/gpu/spec_decode/extract_hidden_states_config.py b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states_config.py index 159827e644c0..0fd9629fcafe 100644 --- a/vllm/v1/worker/gpu/spec_decode/extract_hidden_states_config.py +++ b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states_config.py @@ -37,6 +37,8 @@ def configure(cls, spec: "SpeculativeConfig") -> None: hf_config = {} spec.draft_model_config = copy.copy(spec.target_model_config) + # `ExtractHiddenStatesModel` is native, ignore the target's `model_impl` + spec.draft_model_config.model_impl = "auto" spec.draft_model_config.hf_config = ExtractHiddenStatesConfig( spec.draft_model_config.hf_config, **hf_config )