[Spec Decode] Make speculative decoding compatible with pipeline parallelism

Signed-off-by: Xin Yang <xyangx@amazon.com>

Author: xyang16

tests/spec_decode/e2e/test_integration_dist_pp2.py

@@ -0,0 +1,124 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests which cover integration of the speculative decoding framework with
+pipeline parallelism.
+"""
+
+from typing import Optional
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+from .conftest import run_equality_correctness_test_tp
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2,
+                    reason="Need at least 2 GPUs to run the test.")
+@pytest.mark.parametrize(
+    "common_llm_kwargs",
+    [[
+        # Skip cuda graph recording for fast test.
+        "--enforce-eager",
+        "--pipeline-parallel-size",
+        "2",
+
+        # precision
+        "--dtype",
+        "bfloat16",
+    ]])
+@pytest.mark.parametrize("per_test_common_llm_kwargs", [[]])
+@pytest.mark.parametrize("baseline_llm_kwargs", [[]])
+@pytest.mark.parametrize("model, test_llm_kwargs",
+                         [("JackFram/llama-68m", [
+                             "--speculative-model",
+                             "JackFram/llama-68m",
+                             "--num-speculative-tokens",
+                             "5",
+                             "--speculative-draft-pipeline-parallel-size",
+                             "1",
+                         ]),
+                          ("ibm-granite/granite-3b-code-instruct", [
+                              "--speculative-model",
+                              "ibm-granite/granite-3b-code-instruct",
+                              "--num-speculative-tokens",
+                              "5",
+                              "--speculative-draft-pipeline-parallel-size",
+                              "1",
+                          ])])
+@pytest.mark.parametrize("batch_size", [2])
+@pytest.mark.parametrize("seed", [1])
+def test_draft_model_pp_lt_target_model_pp2(model, common_llm_kwargs,
+                                            per_test_common_llm_kwargs,
+                                            baseline_llm_kwargs,
+                                            test_llm_kwargs, batch_size: int,
+                                            seed: int):
+    """Verify spec decode works well with smaller pp for draft models.
+    """
+    if current_platform.is_rocm():
+        pytest.skip("hip is not well-supported yet")
+    run_equality_correctness_test_tp(model,
+                                     common_llm_kwargs,
+                                     per_test_common_llm_kwargs,
+                                     baseline_llm_kwargs,
+                                     test_llm_kwargs,
+                                     batch_size,
+                                     max_output_len=32,
+                                     seed=seed,
+                                     temperature=0.0)
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2,
+                    reason="Need at least 2 GPUs to run the test.")
+@pytest.mark.parametrize(
+    "common_llm_kwargs",
+    [[
+        # Skip cuda graph recording for fast test.
+        "--enforce-eager",
+        "--pipeline-parallel-size",
+        "2",
+
+        # precision
+        "--dtype",
+        "bfloat16",
+    ]])
+@pytest.mark.parametrize(
+    "per_test_common_llm_kwargs",
+    [["--enable-chunked-prefill", "False"],
+     [
+         "--enable-chunked-prefill", "True", "--max-num-batched-tokens", "4",
+         "--max-num-seqs", "4"
+     ]])
+@pytest.mark.parametrize("baseline_llm_kwargs", [[]])
+@pytest.mark.parametrize("model, test_llm_kwargs", [("JackFram/llama-68m", [
+    "--speculative-model",
+    "JackFram/llama-68m",
+    "--num-speculative-tokens",
+    "3",
+    "--speculative-draft-pipeline-parallel-size",
+    "1",
+])])
+@pytest.mark.parametrize("logprobs", [None, 2])
+@pytest.mark.parametrize("batch_size", [2])
+@pytest.mark.parametrize("seed", [1])
+def test_spec_decode_chunked_prefill_pp2(model, common_llm_kwargs,
+                                         per_test_common_llm_kwargs,
+                                         baseline_llm_kwargs, test_llm_kwargs,
+                                         logprobs: Optional[int],
+                                         batch_size: int, seed: int):
+    """Verify spec decode works well with same and different PP size for
+    the draft model with chunked prefill.
+    """
+    if logprobs:
+        test_llm_kwargs.extend(
+            ["--disable-logprobs-during-spec-decoding", "False"])
+    run_equality_correctness_test_tp(model,
+                                     common_llm_kwargs,
+                                     per_test_common_llm_kwargs,
+                                     baseline_llm_kwargs,
+                                     test_llm_kwargs,
+                                     batch_size,
+                                     max_output_len=32,
+                                     seed=seed,
+                                     temperature=0.0,
+                                     logprobs=logprobs)

tests/spec_decode/test_multi_step_worker.py

@@ -12,7 +12,7 @@
 from vllm.model_executor.utils import set_random_seed
 from vllm.sequence import (ExecuteModelRequest, HiddenStates, Logprob,
                            get_all_seq_ids)
-from vllm.spec_decode.draft_model_runner import TP1DraftModelRunner
+from vllm.spec_decode.draft_model_runner import TP1PP1DraftModelRunner
 from vllm.spec_decode.multi_step_worker import MultiStepWorker
 from vllm.spec_decode.top1_proposer import Top1Proposer
 from vllm.worker.worker import Worker
@@ -91,7 +91,7 @@ def test_same_output_for_single_step():
         block_size,
         num_gpu_blocks,
         seed,
-        model_runner_cls=TP1DraftModelRunner,
+        model_runner_cls=TP1PP1DraftModelRunner,
     )
     worker = create_worker(
         Worker,
@@ -304,7 +304,7 @@ def test_multi_step_with_batch_expansion_correct_output():
         block_size,
         num_gpu_blocks,
         seed,
-        model_runner_cls=TP1DraftModelRunner,
+        model_runner_cls=TP1PP1DraftModelRunner,
     )
     multi_step_worker.set_include_gpu_probs_tensor()
     worker = create_worker(
@@ -399,7 +399,7 @@ def test_multi_step_with_batch_expansion_incorrect_output():
         block_size,
         num_gpu_blocks,
         seed,
-        model_runner_cls=TP1DraftModelRunner,
+        model_runner_cls=TP1PP1DraftModelRunner,
     )
     multi_step_worker.set_include_gpu_probs_tensor()
     worker = create_worker(
@@ -502,13 +502,14 @@ def test_multi_step_correct_kvcache(num_steps, attn_backend):
 
     with global_force_attn_backend_context_manager(attn_backend):
         dtype = 'float16' if attn_backend == _Backend.FLASH_ATTN else 'float32'
-        multi_step_worker = create_worker(MultiStepWorker,
-                                          model_name,
-                                          block_size,
-                                          num_gpu_blocks,
-                                          seed,
-                                          model_runner_cls=TP1DraftModelRunner,
-                                          dtype=dtype)
+        multi_step_worker = create_worker(
+            MultiStepWorker,
+            model_name,
+            block_size,
+            num_gpu_blocks,
+            seed,
+            model_runner_cls=TP1PP1DraftModelRunner,
+            dtype=dtype)
         multi_step_worker.set_include_gpu_probs_tensor()
         worker = create_worker(Worker,
                                model_name,
@@ -771,7 +772,7 @@ def test_use_draft_model_runner_advance_step():
         block_size,
         num_gpu_blocks,
         seed,
-        model_runner_cls=TP1DraftModelRunner,
+        model_runner_cls=TP1PP1DraftModelRunner,
     )
 
     # Mock "_gpu_advance_step" to raise an exception when called.

tests/spec_decode/test_spec_decode_worker.py

@@ -12,7 +12,7 @@
 from vllm.model_executor.utils import set_random_seed
 from vllm.sequence import ExecuteModelRequest, SequenceOutput
 from vllm.spec_decode.batch_expansion import BatchExpansionTop1Scorer
-from vllm.spec_decode.draft_model_runner import TP1DraftModelRunner
+from vllm.spec_decode.draft_model_runner import TP1PP1DraftModelRunner
 from vllm.spec_decode.interfaces import SpeculativeProposals
 from vllm.spec_decode.metrics import (AsyncMetricsCollector,
                                       SpecDecodeWorkerMetrics)
@@ -929,7 +929,7 @@ def test_correctly_load_weight_for_eagle():
         block_size,
         num_gpu_blocks,
         seed,
-        model_runner_cls=TP1DraftModelRunner,
+        model_runner_cls=TP1PP1DraftModelRunner,
     )
 
     spec_decode_sampler = mock_spec_decode_sampler("rejection_sampler")

vllm/attention/layer.py

@@ -152,8 +152,8 @@ def __init__(
         # by bind_kv_cache
         # this variable will not be accessed if use_direct_call is True
         self.kv_cache = [
-            torch.tensor([]) for _ in range(get_current_vllm_config(
-            ).parallel_config.pipeline_parallel_size)
+            torch.tensor([]) for _ in range(
+                get_current_vllm_config().parallel_config.virtual_engine_size)
         ]
 
         self.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32)

vllm/config.py

@@ -1398,6 +1398,7 @@ class ParallelConfig:
     data_parallel_master_ip: str = "127.0.0.1"
     data_parallel_master_port: int = 29500  # Port of the data parallel master.
     enable_expert_parallel: bool = False  # Use EP instead of TP for MoE layers.
+    virtual_engine_size: int = 1  # Number of virtual engine.
 
     # Maximum number of multiple batches
     # when load model sequentially. To avoid RAM OOM when using tensor
@@ -1960,6 +1961,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
@@ -2173,6 +2175,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,
@@ -2183,7 +2192,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:
@@ -2257,19 +2267,41 @@ 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 = \
+                target_parallel_config.pipeline_parallel_size
+        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,
+            virtual_engine_size=target_parallel_config.virtual_engine_size,
             distributed_executor_backend=target_parallel_config.
             distributed_executor_backend,
             max_parallel_loading_workers=target_parallel_config.

vllm/distributed/parallel_state.py

@@ -1011,11 +1011,13 @@ def model_parallel_is_initialized():
 
 
 _TP_STATE_PATCHED = False
+_PP_STATE_PATCHED = False
 
 
 @contextmanager
-def patch_tensor_parallel_group(tp_group: GroupCoordinator):
-    """Patch the tp group temporarily until this function ends.
+def patch_model_parallel_group(tp_group: GroupCoordinator,
+                               pp_group: GroupCoordinator):
+    """Patch the tp and pp group temporarily until this function ends.
 
     This method is for draft workers of speculative decoding to run draft model
     with different tp degree from that of target model workers.
@@ -1026,16 +1028,26 @@ def patch_tensor_parallel_group(tp_group: GroupCoordinator):
     global _TP_STATE_PATCHED
     assert not _TP_STATE_PATCHED, "Should not call when it's already patched"
 
+    global _PP_STATE_PATCHED
+    assert not _PP_STATE_PATCHED, "Should not call when it's already patched"
+
     _TP_STATE_PATCHED = True
     old_tp_group = get_tp_group()
     global _TP
     _TP = tp_group
+
+    _PP_STATE_PATCHED = True
+    old_pp_group = get_pp_group()
+    global _PP
+    _PP = pp_group
     try:
         yield
     finally:
         # restore the original state
         _TP_STATE_PATCHED = False
         _TP = old_tp_group
+        _PP_STATE_PATCHED = False
+        _PP = old_pp_group
 
 
 def get_tensor_model_parallel_world_size():

vllm/engine/arg_utils.py

@@ -187,6 +187,7 @@ class EngineArgs:
     speculative_model: Optional[str] = None
     speculative_model_quantization: Optional[str] = None
     speculative_draft_tensor_parallel_size: Optional[int] = None
+    speculative_draft_pipeline_parallel_size: Optional[int] = None
     num_speculative_tokens: Optional[int] = None
     speculative_disable_mqa_scorer: Optional[bool] = False
     speculative_max_model_len: Optional[int] = None
@@ -843,6 +844,13 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
             default=EngineArgs.speculative_draft_tensor_parallel_size,
             help='Number of tensor parallel replicas for '
             'the draft model in speculative decoding.')
+        parser.add_argument(
+            '--speculative-draft-pipeline-parallel-size',
+            '-spec-draft-pp',
+            type=int,
+            default=EngineArgs.speculative_draft_pipeline_parallel_size,
+            help='Number of pipeline parallel replicas for '
+            'the draft model in speculative decoding.')
 
         parser.add_argument(
             '--speculative-max-model-len',
@@ -1257,6 +1265,8 @@ def create_speculative_config(
                 "max_model_len": self.speculative_max_model_len,
                 "draft_tensor_parallel_size":
                 self.speculative_draft_tensor_parallel_size,
+                "draft_pipeline_parallel_size":
+                self.speculative_draft_pipeline_parallel_size,
                 "num_speculative_tokens": self.num_speculative_tokens,
                 "disable_mqa_scorer": self.speculative_disable_mqa_scorer,
                 "disable_by_batch_size":
@@ -1369,6 +1379,7 @@ def create_engine_config(
         parallel_config = ParallelConfig(
             pipeline_parallel_size=self.pipeline_parallel_size,
             tensor_parallel_size=self.tensor_parallel_size,
+            virtual_engine_size=self.pipeline_parallel_size,
             data_parallel_size=self.data_parallel_size,
             enable_expert_parallel=self.enable_expert_parallel,
             max_parallel_loading_workers=self.max_parallel_loading_workers,