[Speculative Decoding] Support draft model on different tensor-parallel size than target model

vllm/config.py

@@ -783,6 +783,7 @@ def maybe_create_spec_config(
         target_parallel_config: ParallelConfig,
         target_dtype: str,
         speculative_model: Optional[str],
+        speculative_tensor_parallel_size: Optional[int],
         num_speculative_tokens: Optional[int],
         speculative_max_model_len: Optional[int],
         enable_chunked_prefill: bool,
@@ -878,7 +879,6 @@ def maybe_create_spec_config(
             # config, in future, we may try refactor it out, and set
             # draft related config as None here.
             draft_model_config = target_model_config
-            draft_parallel_config = target_parallel_config
         else:
             ngram_prompt_lookup_max = 0
             ngram_prompt_lookup_min = 0
@@ -907,9 +907,9 @@ def maybe_create_spec_config(
                     target_model_config.max_model_len,
                 ))
 
-            draft_parallel_config = (
-                SpeculativeConfig.create_draft_parallel_config(
-                    target_parallel_config))
+        draft_parallel_config = (
+            SpeculativeConfig.create_draft_parallel_config(
+                target_parallel_config, speculative_tensor_parallel_size))
 
         return SpeculativeConfig(
             draft_model_config,
@@ -957,16 +957,27 @@ def _maybe_override_draft_max_model_len(
 
     @staticmethod
     def create_draft_parallel_config(
-            target_parallel_config: ParallelConfig) -> ParallelConfig:
+            target_parallel_config: ParallelConfig,
+            speculative_tensor_parallel_size: Optional[int]) -> ParallelConfig:
         """Create a parallel config for use by the draft worker.
 
-        This is mostly a copy of the target parallel config. In the future the
-        draft worker can have a different parallel strategy, e.g. TP=1.
+        This is mostly a copy of the target parallel config, except the tp_size.
         """
+
+        speculative_tensor_parallel_size = (
+            speculative_tensor_parallel_size
+            or target_parallel_config.tensor_parallel_size)
+
+        if speculative_tensor_parallel_size > \
+            target_parallel_config.tensor_parallel_size:
+            raise ValueError(
+                f"{speculative_tensor_parallel_size=} cannot be "
+                f"larger than {target_parallel_config.tensor_parallel_size}")
+
         draft_parallel_config = ParallelConfig(
             pipeline_parallel_size=target_parallel_config.
             pipeline_parallel_size,
-            tensor_parallel_size=target_parallel_config.tensor_parallel_size,
+            tensor_parallel_size=speculative_tensor_parallel_size,
             distributed_executor_backend=target_parallel_config.
             distributed_executor_backend,
             max_parallel_loading_workers=target_parallel_config.

vllm/distributed/parallel_state.py

@@ -550,6 +550,10 @@ def init_distributed_environment(
     global _WORLD
     if _WORLD is None:
         ranks = list(range(torch.distributed.get_world_size()))
+        if world_size != -1:
+            assert world_size == len(ranks), (
+                "given world_size does not match with world_size of torch")
+
         _WORLD = GroupCoordinator(
             group_ranks=[ranks],
             local_rank=local_rank,
@@ -558,7 +562,7 @@ def init_distributed_environment(
             use_custom_allreduce=False,
         )
     else:
-        assert _WORLD.world_size == torch.distributed.get_world_size(), (
+        assert _WORLD.world_size == world_size, (
             "world group already initialized with a different world size")
 
 
@@ -673,6 +677,31 @@ def model_parallel_is_initialized():
     return (_TP is not None and _PP is not None)
 
 
+OVERRIDE_TP_STATE = False
+
+
+@contextlib.contextmanager
+def patch_tensor_parallel_group(world_group, tp_group):
+    """Patch the tp group temporarily until this function ends."""
+    global OVERRIDE_TP_STATE
+    if OVERRIDE_TP_STATE:
+        return
+
+    OVERRIDE_TP_STATE = True
+    old_world_group = get_world_group()
+    old_tp_group = get_tp_group()
+    global _WORLD, _TP
+    _WORLD = world_group
+    _TP = tp_group
+    try:
+        yield
+    finally:
+        # restore the original state
+        OVERRIDE_TP_STATE = False
+        _WORLD = old_world_group
+        _TP = old_tp_group
+
+
 def get_tensor_model_parallel_world_size():
     """Return world size for the tensor model parallel group."""
     return get_tp_group().world_size

vllm/engine/arg_utils.py

@@ -93,6 +93,7 @@ class EngineArgs:
     guided_decoding_backend: str = 'outlines'
     # Speculative decoding configuration.
     speculative_model: Optional[str] = None
+    speculative_tensor_parallel_size: Optional[int] = None
     num_speculative_tokens: Optional[int] = None
     speculative_max_model_len: Optional[int] = None
     speculative_disable_by_batch_size: Optional[int] = None
@@ -534,6 +535,13 @@ def add_cli_args(
             default=EngineArgs.num_speculative_tokens,
             help='The number of speculative tokens to sample from '
             'the draft model in speculative decoding.')
+        parser.add_argument(
+            '--speculative-tensor-parallel-size',
+            '-spec-tp',
+            type=int,
+            default=EngineArgs.speculative_tensor_parallel_size,
+            help='Number of tensor parallel replicas for '
+            'the draft model in speculative decoding.')
 
         parser.add_argument(
             '--speculative-max-model-len',
@@ -676,6 +684,8 @@ def create_engine_config(self, ) -> EngineConfig:
             target_parallel_config=parallel_config,
             target_dtype=self.dtype,
             speculative_model=self.speculative_model,
+            speculative_tensor_parallel_size = \
+                self.speculative_tensor_parallel_size,
             num_speculative_tokens=self.num_speculative_tokens,
             speculative_disable_by_batch_size=self.
             speculative_disable_by_batch_size,

vllm/spec_decode/multi_step_worker.py

@@ -6,7 +6,8 @@
 
 from vllm.sequence import (ExecuteModelRequest, SamplerOutput,
                            SequenceGroupMetadata)
-from vllm.spec_decode.interfaces import SpeculativeProposals
+from vllm.spec_decode.interfaces import (SpeculativeProposals,
+                                         SpeculativeProposer)
 from vllm.spec_decode.proposer_worker_base import ProposerWorkerBase
 from vllm.spec_decode.top1_proposer import Top1Proposer
 from vllm.worker.worker import Worker
@@ -28,7 +29,7 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
         # Lazy initialization list.
-        self._proposer: Top1Proposer
+        self._proposer: SpeculativeProposer
 
     def init_device(self):
         super().init_device()
@@ -71,9 +72,9 @@ def sampler_output(
                                      sample_len)
 
         # Run model sample_len times.
-        model_outputs = []
+        model_outputs: List[SamplerOutput] = []
         for _ in range(sample_len):
-            model_output = super().execute_model(
+            model_output: List[SamplerOutput] = super().execute_model(
                 execute_model_req=copied_execute_model_req)
             assert (len(model_output) == 1
                     ), "composing multistep workers not supported"

vllm/spec_decode/smaller_tp_proposer_worker.py

@@ -0,0 +1,213 @@
+from typing import List, Optional, Set, Tuple, Union
+
+import torch
+import torch.distributed
+
+from vllm.config import ParallelConfig
+from vllm.distributed.parallel_state import (_ENABLE_CUSTOM_ALL_REDUCE,
+                                             GroupCoordinator, get_tp_group,
+                                             get_world_group,
+                                             patch_tensor_parallel_group)
+from vllm.logger import init_logger
+from vllm.lora.request import LoRARequest
+from vllm.sequence import ExecuteModelRequest, SamplerOutput
+from vllm.spec_decode.interfaces import SpeculativeProposals
+from vllm.spec_decode.proposer_worker_base import ProposerWorkerBase
+from vllm.worker.worker import Worker
+
+logger = init_logger(__name__)
+
+
+class SmallerTpProposerWorker(ProposerWorkerBase):
+    """Class which allows a speculative draft model to run with smaller tensor
+    parallel degree than target model.
+    This reduces the communication overhead of small draft models.
+
+    This is implemented by changing vLLM's tensor parallel group to a group of
+    size temporarily during forward passes of draft models.
+    """
+
+    @classmethod
+    def maybe_wrap_worker(cls, worker, draft_parallel_config: ParallelConfig,
+                          target_parallel_config: ParallelConfig, rank: int):
+        """Wrap the worker in a SmallerTpProposerWorker if necessary.
+        """
+        draft_tp = draft_parallel_config.tensor_parallel_size
+        target_tp = target_parallel_config.tensor_parallel_size
+
+        if draft_tp == target_tp:
+            return worker
+
+        if draft_tp > target_tp:
+            raise ValueError(
+                f"{cls} only supports draft_tp smaller than target_tp."
+                f"{draft_tp=} {target_tp=}")
+
+        # gpu ranks that will generate draft tokens together
+        ranks = list(range(draft_tp))
+
+        if rank in ranks:
+            logger.info("Wrapping {%s} in {%s}", type(worker), cls)
+            return cls(worker, ranks)
+        else:
+            # for workers not participating in the draft generation
+            logger.info("Returning dummy worker")
+            return DummyProposerWorker(worker)
+
+    def __init__(self, worker: Union[Worker, ProposerWorkerBase],
+                 ranks: List[int]):
+        self._worker = worker
+        self._ranks = ranks
+        self._world_group = None
+        self._tp_group = None
+
+    def _patch_tensor_parallel_group(self):
+        return patch_tensor_parallel_group(self._world_group, self._tp_group)
+
+    def init_device(self):
+        """Initialize the device.
+
+        This also creates an additional tensor-parallel process group containing
+        only a subset of the whole ranks.
+        """
+        local_rank = get_world_group().local_rank
+        world_backend = torch.distributed.get_backend(
+            get_world_group().device_group)
+        tp_backend = torch.distributed.get_backend(get_tp_group().device_group)
+
+        self._world_group = GroupCoordinator(
+            group_ranks=[self._ranks],
+            local_rank=local_rank,
+            torch_distributed_backend=world_backend,
+            use_pynccl=False,
+            use_custom_allreduce=False,
+        )
+        self._tp_group = GroupCoordinator(
+            group_ranks=[self._ranks],
+            local_rank=local_rank,
+            torch_distributed_backend=tp_backend,
+            use_pynccl=True,
+            use_custom_allreduce=_ENABLE_CUSTOM_ALL_REDUCE,
+        )
+
+        with self._patch_tensor_parallel_group():
+            self._worker.init_device()
+
+    def set_include_gpu_probs_tensor(self):
+        self._worker.set_include_gpu_probs_tensor()
+
+    def load_model(self):
+        with self._patch_tensor_parallel_group():
+            self._worker.load_model()
+
+    def determine_num_available_blocks(self):
+        with self._patch_tensor_parallel_group():
+            return self._worker.determine_num_available_blocks()
+
+    def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks: int):
+        with self._patch_tensor_parallel_group():
+            self._worker.initialize_cache(num_gpu_blocks, num_cpu_blocks)
+
+    @torch.inference_mode()
+    def sampler_output(
+        self,
+        execute_model_req: ExecuteModelRequest,
+        sample_len: int,
+    ) -> Tuple[Optional[List[SamplerOutput]], bool]:
+        # do not call _parch_tensor_parallel_group, because
+        # it's always called after tp_group has already been overridden
+        return self._worker.sampler_output(execute_model_req, sample_len)
+
+    def get_spec_proposals(
+        self,
+        execute_model_req: ExecuteModelRequest,
+    ) -> SpeculativeProposals:
+        with self._patch_tensor_parallel_group():
+            return self._worker.get_spec_proposals(execute_model_req)
+
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        execute_model_req: Optional[ExecuteModelRequest] = None
+    ) -> List[SamplerOutput]:
+        with self._patch_tensor_parallel_group():
+            return self._worker.execute_model(execute_model_req)
+
+    def get_cache_block_size_bytes(self) -> int:
+        return self._worker.get_cache_block_size_bytes()
+
+    def add_lora(self, lora_request: LoRARequest) -> bool:
+        raise NotImplementedError
+
+    def remove_lora(self, lora_id: int) -> bool:
+        raise NotImplementedError
+
+    def list_loras(self) -> Set[int]:
+        raise NotImplementedError
+
+    @property
+    def max_model_len(self) -> int:
+        return self._worker.max_model_len
+
+    @property
+    def vocab_size(self) -> int:
+        return self._worker.vocab_size
+
+
+class DummyProposerWorker(ProposerWorkerBase):
+    """Dummy proposer worker that do nothing.
+    It's for workers that do not participate in draft generation.
+    """
+
+    def __init__(
+        self,
+        worker: Union[Worker, ProposerWorkerBase],
+    ):
+        self._worker = worker
+
+    def init_device(self):
+        pass
+
+    def load_model(self):
+        pass
+
+    def determine_num_available_blocks(self):
+        pass
+
+    def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks: int):
+        pass
+
+    def sampler_output(
+        self,
+        execute_model_req: ExecuteModelRequest,
+        sample_len: int,
+    ) -> Tuple[List[SamplerOutput], bool]:
+        return [], True
+
+    def get_spec_proposals(
+        self,
+        execute_model_req: ExecuteModelRequest,
+    ) -> SpeculativeProposals:
+        return SpeculativeProposals(None, None, None)
+
+    def execute_model(
+        self,
+        execute_model_req: Optional[ExecuteModelRequest] = None
+    ) -> List[SamplerOutput]:
+        return []
+
+    def get_cache_block_size_bytes(self) -> int:
+        return 0
+
+    def add_lora(self, lora_request: LoRARequest) -> bool:
+        return False
+
+    def remove_lora(self, lora_id: int) -> bool:
+        return False
+
+    def list_loras(self) -> Set[int]:
+        return set()
+
+    @property
+    def vocab_size(self) -> int:
+        return self._worker.vocab_size