[V1][Spec Decode] Make Eagle compatible with pipeline parallelism

vllm/config.py

@@ -2071,6 +2071,9 @@ class SpeculativeConfig:
         - draft_tensor_parallel_size (Optional[int]): The degree of the tensor
             parallelism for the draft model. Can only be 1 or the same as the
             target model's tensor parallel size.
+        - draft_pipeline_parallel_size (Optional[int]): The degree of the
+            pipeline parallelism for the draft model. Can only be 1 or the same
+            as the target model's tensor parallel size.
         - disable_logprobs (bool): If set to True, token log probabilities are
             not returned during speculative decoding. If set to False, token
             log probabilities are returned according to the log probability
@@ -2126,6 +2129,7 @@ class SpeculativeConfig:
     method: Optional[str] = None
     acceptance_method: str = "rejection_sampler"
     draft_tensor_parallel_size: Optional[int] = None
+    draft_pipeline_parallel_size: Optional[int] = None
     disable_logprobs: bool = True
 
     model: Optional[str] = None
@@ -2339,6 +2343,13 @@ def __post_init__(self):
                         self.draft_model_config.hf_config
                 )
 
+                self.draft_pipeline_parallel_size = \
+                    SpeculativeConfig._verify_and_get_draft_pp(
+                        self.target_parallel_config,
+                        self.draft_pipeline_parallel_size,
+                        self.draft_model_config.hf_config
+                    )
+
                 self.draft_model_config.max_model_len = (
                     SpeculativeConfig._maybe_override_draft_max_model_len(
                         self.max_model_len,
@@ -2349,7 +2360,8 @@ def __post_init__(self):
                 self.draft_parallel_config = (
                     SpeculativeConfig.create_draft_parallel_config(
                         self.target_parallel_config,
-                        self.draft_tensor_parallel_size))
+                        self.draft_tensor_parallel_size,
+                        self.draft_pipeline_parallel_size))
 
         if self.acceptance_method == "typical_acceptance_sampler":
             if self.posterior_threshold is None:
@@ -2423,18 +2435,38 @@ def _verify_and_get_draft_tp(
                 f"other value than 1 or target model tensor_parallel_size")
         return speculative_draft_tensor_parallel_size
 
+    @staticmethod
+    def _verify_and_get_draft_pp(
+            target_parallel_config: ParallelConfig,
+            speculative_draft_pipeline_parallel_size: Optional[int],
+            draft_hf_config: PretrainedConfig) -> int:
+        """
+        Verifies and adjusts the tensor parallel size for a draft model
+        specified using speculative_draft_pipeline_parallel_size.
+        """
+        # If speculative_draft_pipeline_parallel_size is unset then set it
+        # appropriately else verify that it is set correctly.
+        if speculative_draft_pipeline_parallel_size is None:
+            speculative_draft_pipeline_parallel_size = 1
+        elif speculative_draft_pipeline_parallel_size not in (
+                1, target_parallel_config.pipeline_parallel_size):
+            raise ValueError(
+                f"{speculative_draft_pipeline_parallel_size=} cannot be "
+                f"other value than 1 or target model pipeline_parallel_size")
+        return speculative_draft_pipeline_parallel_size
+
     @staticmethod
     def create_draft_parallel_config(
         target_parallel_config: ParallelConfig,
         speculative_draft_tensor_parallel_size: int,
+        speculative_draft_pipeline_parallel_size: int,
     ) -> ParallelConfig:
         """Create a parallel config for use by the draft worker.
 
         This is mostly a copy of the target parallel config, except the tp_size.
         """
         draft_parallel_config = ParallelConfig(
-            pipeline_parallel_size=target_parallel_config.
-            pipeline_parallel_size,
+            pipeline_parallel_size=speculative_draft_pipeline_parallel_size,
             tensor_parallel_size=speculative_draft_tensor_parallel_size,
             distributed_executor_backend=target_parallel_config.
             distributed_executor_backend,

vllm/model_executor/models/llama_eagle.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Iterable, Set, Tuple
+from typing import Iterable, Optional, Set, Tuple
 
 import torch
 import torch.nn as nn
@@ -14,6 +14,7 @@
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.llama import (LlamaDecoderLayer,
                                               LlamaForCausalLM)
+from vllm.sequence import IntermediateTensors
 
 from .utils import AutoWeightsLoader, maybe_prefix
 
@@ -70,6 +71,7 @@ def forward(
         input_ids: torch.Tensor,
         positions: torch.Tensor,
         hidden_states: torch.Tensor,
+        intermediate_tensors: Optional[IntermediateTensors] = None,
     ) -> torch.Tensor:
         input_embeds = self.embed_tokens(input_ids)
         hidden_states = self.fc(
@@ -132,6 +134,7 @@ def forward(
         input_ids: torch.Tensor,
         positions: torch.Tensor,
         hidden_states: torch.Tensor,
+        intermediate_tensors: Optional[IntermediateTensors] = None,
     ) -> torch.Tensor:
         return self.model(input_ids, positions, hidden_states)
 

vllm/v1/spec_decode/eagle.py

@@ -21,6 +21,7 @@ def __init__(
         device: torch.device,
     ):
         self.vllm_config = vllm_config
+        self.device = device
         self.num_speculative_tokens = (
             vllm_config.speculative_config.num_speculative_tokens)
         self.block_size = vllm_config.cache_config.block_size
@@ -187,12 +188,12 @@ def prepare_inputs(
     def load_model(self, target_model: nn.Module) -> None:
         loader = get_model_loader(self.vllm_config.load_config)
         target_layer_num = self.vllm_config.model_config.get_num_layers(
-            self.vllm_config.parallel_config)
+            self.vllm_config.parallel_config
+        ) * self.vllm_config.parallel_config.pipeline_parallel_size
 
         draft_model_config = \
             self.vllm_config.speculative_config.draft_model_config
-        # FIXME(lily): This does not handle with distributed inference.
-        target_device = self.vllm_config.device_config.device
+        target_device = self.device
         # We need to set the vllm_config here to register attention
         # layers in the forward context.
         with set_default_torch_dtype(

vllm/v1/worker/gpu_model_runner.py

@@ -161,16 +161,21 @@ def __init__(
         self.use_spec_decode = False
         if self.speculative_config:
             self.use_spec_decode = True
-            if get_pp_group().is_last_rank:
-                if self.speculative_config.method == "ngram":
-                    self.drafter = NgramProposer(self.vllm_config)
-                elif self.speculative_config.method == "eagle":
-                    self.drafter = EagleProposer(self.vllm_config,
-                                                 self.device)  # type: ignore
-                else:
-                    raise ValueError("Unknown speculative decoding method: "
-                                     f"{self.speculative_config.method}")
-                self.rejection_sampler = RejectionSampler()
+            if self.speculative_config.draft_pipeline_parallel_size == 1:
+                if get_pp_group().is_last_rank:
+                    if self.speculative_config.method == "ngram":
+                        self.drafter = NgramProposer(self.vllm_config)
+                    elif self.speculative_config.method == "eagle":
+                        self.drafter = EagleProposer(
+                            self.vllm_config, self.device)  # type: ignore
+                    else:
+                        raise ValueError(
+                            "Unknown speculative decoding method: "
+                            f"{self.speculative_config.method}")
+                    self.rejection_sampler = RejectionSampler()
+            else:
+                raise ValueError(
+                    "draft_pipeline_parallel_size > 1 is not supported")
 
         # Request states.
         self.requests: dict[str, CachedRequestState] = {}