A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 22:15 UTC
[bugfix] spec decode worker get tp group only when initialized (#13578)
Simon Mo committed
Feb 19, 2025 at 20:46 UTC
8c755c3b6d37954895a9952eb6bf7691c4d25a50
1 file changed
+7
-2
vllm/spec_decode/spec_decode_worker.py
+7
-2
@@ -12,6 +12,7 @@ from vllm.config import ParallelConfig, SpeculativeConfig, VllmConfig
12
from vllm.distributed.communication_op import (broadcast_tensor_dict,
13
get_tp_group,
14
tensor_model_parallel_gather)
15
+from vllm.distributed.parallel_state import model_parallel_is_initialized
16
from vllm.logger import init_logger
17
from vllm.model_executor.layers.rejection_sampler import RejectionSampler
18
from vllm.model_executor.layers.sampler import SamplerOutput
@@ -366,8 +367,12 @@ class SpecDecodeWorker(LoraNotSupportedWorkerBase):
367
target_lm_head_weight)
368
369
self._metrics.init_tensors(self.rank, device_type=self.device)
369
- self.spec_decode_sampler.init_tensors(get_tp_group().local_rank,
370
- device_type=self.device)
370
+ if model_parallel_is_initialized():
371
+ self.spec_decode_sampler.init_tensors(get_tp_group().local_rank,
372
+ device_type=self.device)
373
+ else:
374
+ self.spec_decode_sampler.init_tensors(self.rank,
375
+ device_type=self.device)
376
377
scorer_cls: Type[SpeculativeScorer]
378
if self.disable_mqa_scorer: