Merge pull request #6 from tomasruizt/feature/correct-tensor-parallelism-on-draft-model

Feature/correct tensor parallelism on draft model

Author: tomasruizt

tests/v1/e2e/test_spec_decode.py

@@ -11,9 +11,12 @@
 from vllm import LLM, SamplingParams
 from vllm.assets.base import VLLM_S3_BUCKET_URL
 from vllm.assets.image import VLM_IMAGES_DIR
+from vllm.config.vllm import VllmConfig
 from vllm.distributed import cleanup_dist_env_and_memory
+from vllm.engine.arg_utils import EngineArgs
 from vllm.outputs import RequestOutput
 from vllm.platforms import current_platform
+from vllm.v1.spec_decode.draft_model import create_vllm_config_for_draft_model
 from vllm.v1.spec_decode.metrics import compute_acceptance_len, compute_acceptance_rate
 
 MTP_SIMILARITY_RATE = 0.8
@@ -359,7 +362,7 @@ def test_mtp_correctness(
 
 @dataclass
 class ArgsTest:
-    model: str
+    target_model: str
     draft_model: str
     sampling_config: SamplingParams
     num_speculative_tokens: int
@@ -376,7 +379,7 @@ class ArgsTest:
 cases = [
     # Same model for draft and target, greedy sampling.
     ArgsTest(
-        model="Qwen/Qwen3-0.6B",
+        target_model="Qwen/Qwen3-0.6B",
         draft_model="Qwen/Qwen3-0.6B",
         sampling_config=greedy_sampling(),
         num_speculative_tokens=3,  # K
@@ -386,7 +389,7 @@ class ArgsTest:
     ),
     # Smaller draft model, stochastic sampling.
     ArgsTest(
-        model="Qwen/Qwen3-1.7B",
+        target_model="Qwen/Qwen3-1.7B",
         draft_model="Qwen/Qwen3-0.6B",
         sampling_config=stochastic_sampling(),
         num_speculative_tokens=3,
@@ -416,31 +419,80 @@ def test_draft_model_correctness(args: ArgsTest, enforce_eager: bool):
 def test_draft_model_quantization(models: tuple[str, str], enforce_eager: bool):
     tgt_model, draft_model = models
     sd_case = ArgsTest(
-        model=tgt_model,
+        target_model=tgt_model,
         draft_model=draft_model,
-        sampling_config=greedy_sampling(),
-        num_speculative_tokens=3,
-        expected_acceptance_len=2.95 + 1,
-        expected_acceptance_rate=0.95,
-        expected_same_output_fraction=0.95,
+        **some_high_acceptance_metrics(),
     )
     assert_draft_model_correctness(sd_case, enforce_eager)
 
 
+def test_draft_model_tensor_parallelism():
+    """Ensure spec decode works when running with TP > 1."""
+    sd_case = ArgsTest(
+        target_model="Qwen/Qwen3-1.7B",
+        target_tensor_parallel_size=2,
+        draft_model="Qwen/Qwen3-0.6B",
+        draft_tensor_parallel_size=2,
+        **some_high_acceptance_metrics(),
+    )
+    assert_draft_model_correctness(sd_case, enforce_eager=False)
+
+
+def test_draft_model_engine_args_tensor_parallelism():
+    """Ensure the vllm_config for the draft model is created correctly,
+    and independently of the target model (quantization, TP, etc.)"""
+
+    engine_args = EngineArgs(
+        model="Qwen/Qwen3-1.7B-FP8",  # <<< tgt quantized
+        tensor_parallel_size=4,
+        speculative_config={
+            "model": "Qwen/Qwen3-0.6B",  # <<< draft not quantized
+            "method": "draft_model",
+            "num_speculative_tokens": 3,
+            "draft_tensor_parallel_size": 1,  # <<< valid arg name
+        },
+    )
+    tgt_vllm_config: VllmConfig = engine_args.create_engine_config()
+    assert tgt_vllm_config.parallel_config.tensor_parallel_size == 4
+    assert tgt_vllm_config.quant_config.get_name() == "fp8"
+
+    draft_vllm_config: VllmConfig = create_vllm_config_for_draft_model(tgt_vllm_config)
+    assert draft_vllm_config.parallel_config.tensor_parallel_size == 1
+    assert draft_vllm_config.quant_config is None
+
+
+def test_draft_model_engine_args_rejects_invalid_tp_argname():
+    """The user should pass "draft_tensor_parallel_size" rather than
+    "tensor_parallel_size". We enforce this with validation."""
+
+    engine_args = EngineArgs(
+        model="Qwen/Qwen3-1.7B",
+        tensor_parallel_size=1,
+        speculative_config={
+            "model": "Qwen/Qwen3-0.6B",
+            "method": "draft_model",
+            "num_speculative_tokens": 3,
+            "tensor_parallel_size": 1,  # <<< invalid arg name
+        },
+    )
+    with pytest.raises(ValueError):
+        engine_args.create_engine_config()
+
+
 def assert_draft_model_correctness(args: ArgsTest, enforce_eager: bool):
     """Compare the outputs using and not using speculative decoding.
     In the greedy decoding case, the outputs must match EXACTLY."""
     test_prompts = get_test_prompts(mm_enabled=False, quiet=True)
 
     spec_llm = LLM(
-        model=args.model,
+        model=args.target_model,
         speculative_config={
             "model": args.draft_model,
             "method": "draft_model",
             "num_speculative_tokens": args.num_speculative_tokens,
             "max_model_len": args.max_model_len,
             "enforce_eager": enforce_eager,
-            "tensor_parallel_size": args.draft_tensor_parallel_size,
+            "draft_tensor_parallel_size": args.draft_tensor_parallel_size,
             "disable_padded_drafter_batch": True,
             "max_num_seqs": 100,  # limit cudagraph capture runtime
         },
@@ -462,7 +514,7 @@ def assert_draft_model_correctness(args: ArgsTest, enforce_eager: bool):
     assert acceptance_len >= args.expected_acceptance_len
 
     ref_llm = LLM(
-        model=args.model,
+        model=args.target_model,
         max_model_len=args.max_model_len,
         gpu_memory_utilization=args.gpu_memory_utilization,
         tensor_parallel_size=args.target_tensor_parallel_size,
@@ -480,7 +532,7 @@ def assert_draft_model_correctness(args: ArgsTest, enforce_eager: bool):
     assert match_fraction >= args.expected_same_output_fraction
 
     print(
-        f"spec-decode: target={args.model}, draft={args.draft_model}, "
+        f"spec-decode: target={args.target_model}, draft={args.draft_model}, "
         f"temperature={args.sampling_config.temperature:.2f}, "
         f"acceptance_rate={acceptance_rate:.2f}, "
         f"acceptance_len={acceptance_len:.2f}, "
@@ -501,3 +553,13 @@ def compute_exact_matches(
             print(f"ref_output: {ref_output.outputs[0].text}")
             print(f"spec_output: {spec_output.outputs[0].text}")
     return matches / len(ref_outputs)
+
+
+def some_high_acceptance_metrics() -> dict:
+    return {
+        "sampling_config": greedy_sampling(),
+        "num_speculative_tokens": 3,
+        "expected_acceptance_len": 2.95 + 1,
+        "expected_acceptance_rate": 0.95,
+        "expected_same_output_fraction": 0.95,
+    }

vllm/benchmarks/lib/ready_checker.py

@@ -8,8 +8,12 @@
 import aiohttp
 from tqdm.asyncio import tqdm
 
+from vllm.logger import init_logger
+
 from .endpoint_request_func import RequestFunc, RequestFuncInput, RequestFuncOutput
 
+logger = init_logger(__name__)
+
 
 async def wait_for_endpoint(
     request_func: RequestFunc,
@@ -61,6 +65,8 @@ async def wait_for_endpoint(
                 if output.success:
                     pbar.close()
                     return output
+                else:
+                    logger.warning("Endpoint is not ready. Error='%s'", output.error)
             except aiohttp.ClientConnectorError:
                 pass
 

vllm/config/parallel.py

@@ -3,6 +3,7 @@
 
 import hashlib
 import os
+from dataclasses import replace
 from typing import TYPE_CHECKING, Any, Literal
 
 import torch
@@ -564,3 +565,6 @@ def _verify_args(self) -> Self:
             )
 
         return self
+
+    def replace(self, **kwargs) -> Self:
+        return replace(self, **kwargs)

vllm/config/speculative.py

@@ -79,6 +79,10 @@ class SpeculativeConfig:
     draft_tensor_parallel_size: int | None = None
     """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."""
+    tensor_parallel_size: int | None = None
+    """Users should pass "draft_tensor_parallel_size". This parameters is only 
+    to reject it if passed."""
+
     disable_logprobs: bool = True
     """If set to True, token log probabilities are not returned during
     speculative decoding. If set to False, token log probabilities are returned
@@ -537,6 +541,12 @@ def create_draft_parallel_config(
 
     @model_validator(mode="after")
     def _verify_args(self) -> Self:
+        if self.tensor_parallel_size is not None:
+            raise ValueError(
+                "'tensor_parallel_size' is not a valid argument in the "
+                "speculative_config. Please pass 'draft_tensor_parallel_size' instead."
+            )
+
         if self.num_speculative_tokens is None:
             raise ValueError(
                 "num_speculative_tokens must be provided with "

vllm/v1/spec_decode/draft_model.py

@@ -6,7 +6,9 @@
 import torch
 
 from vllm.attention.layer import Attention
-from vllm.config import ModelConfig, VllmConfig, get_layers_from_vllm_config
+from vllm.config import VllmConfig, get_layers_from_vllm_config
+from vllm.config.speculative import SpeculativeConfig
+from vllm.logger import init_logger
 from vllm.model_executor.model_loader import get_model
 from vllm.v1.attention.backends.utils import (
     CommonAttentionMetadata,
@@ -16,6 +18,8 @@
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.spec_decode.eagle import PADDING_SLOT_ID, SpecDecodeBaseProposer
 
+logger = init_logger(__name__)
+
 
 class DraftModelProposer(SpecDecodeBaseProposer):
     def __init__(
@@ -34,6 +38,7 @@ def __init__(
         self._raise_if_mrope()
         self._raise_if_padded_drafter_batch()
         self._raise_if_vocab_size_mismatch()
+        self._raise_if_draft_tp_mismatch()
 
     def propose(
         self,
@@ -98,12 +103,29 @@ def _raise_if_padded_drafter_batch(self):
             raise NotImplementedError(
                 "Speculative Decoding with draft models does not support "
                 "padded drafter batch yet. Please pass --disable-padded-drafter-batch "
-                "in the speculative config."
+                "in the speculative_config."
             )
 
     def _raise_if_vocab_size_mismatch(self):
         self.vllm_config.speculative_config.verify_equal_vocab_size_if_draft_model()
 
+    def _raise_if_draft_tp_mismatch(self):
+        # Note(Tomas Ruiz) If we run the target model with TP > 1 and
+        # the draft model with TP = 1, then the different TP ranks collide.
+        # Specifically when all ranks compile the draft model on rank 0
+        # (because TP=1), then the torch compile cache is overwritten and corrupted.
+        # We need a mechanism like this: https://github.com/vllm-project/vllm/pull/5414
+        # To prevent this error, we assert that both TP sizes must be the same.
+        spec_cfg: SpeculativeConfig = self.vllm_config.speculative_config
+        tgt_tp = spec_cfg.target_parallel_config.tensor_parallel_size
+        draft_tp = spec_cfg.draft_parallel_config.tensor_parallel_size
+        if draft_tp != tgt_tp:
+            raise ValueError(
+                f"Currently, 'draft_tensor_parallel_size' and 'tensor_parallel_size' "
+                f"must be the same. Got {draft_tp} and {tgt_tp}. "
+                "Please pass 'draft_tensor_parallel_size' in the speculative_config."
+            )
+
     def set_input_ids_first_pass(
         self,
         target_token_ids: torch.Tensor,
@@ -115,15 +137,6 @@ def set_input_ids_first_pass(
 
     def load_model(self, target_model: Any) -> None:
         """Takes target_model to satisfy the type checker."""
-        draft_model_config: ModelConfig = (
-            self.vllm_config.speculative_config.draft_model_config
-        )
-        # Recompute quant_config, which is configured for the target model
-        # But the draft model might not be quantized.
-        vllm_config_draft: VllmConfig = self.vllm_config.replace(
-            quant_config=None,
-            model_config=draft_model_config,
-        )
 
         # This must be computed before loading the draft model
         # because that mutates the forward_context of the vllm_config
@@ -133,12 +146,17 @@ def load_model(self, target_model: Any) -> None:
 
         from vllm.compilation.backends import set_model_tag
 
+        draft_vllm_config: VllmConfig = create_vllm_config_for_draft_model(
+            target_model_vllm_config=self.vllm_config
+        )
+        logger.info(
+            "Starting to load draft model %s. TP=%d, rank=%d",
+            draft_vllm_config.model_config.model,
+            draft_vllm_config.parallel_config.tensor_parallel_size,
+            draft_vllm_config.parallel_config.rank,
+        )
         with set_model_tag("draft_model"):
-            self.model = get_model(
-                vllm_config=vllm_config_draft,
-                model_config=draft_model_config,
-                prefix="draft_model",
-            )
+            self.model = get_model(vllm_config=draft_vllm_config, prefix="draft_model")
 
         # This must be computed after loading the draft model
         # because that mutates the forward_context of the vllm_config
@@ -149,6 +167,27 @@ def load_model(self, target_model: Any) -> None:
         self.attn_layer_names = list(draft_attn_layer_names)
 
 
+def create_vllm_config_for_draft_model(
+    target_model_vllm_config: VllmConfig,
+) -> VllmConfig:
+    """The vllm_config is configured for the target model, e.g.
+    its quant_config and parallel_config. But the draft model is potentially
+    quantized differently, and has potentially different tensor_parallel_size.
+    This function creates a new vllm_config configured for the draft model.
+    The vllm_config is useful when loading the draft model with get_model().
+    """
+    old = target_model_vllm_config
+    new_parallel_config = old.speculative_config.draft_parallel_config.replace(
+        rank=old.parallel_config.rank
+    )
+    new: VllmConfig = old.replace(
+        quant_config=None,  # quant_config is recomputed in __init__()
+        model_config=old.speculative_config.draft_model_config,
+        parallel_config=new_parallel_config,
+    )
+    return new
+
+
 @dataclass
 class DraftModelInputs:
     token_ids: torch.Tensor