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] Fix device ordinal for multi-node spec decode (#13269)
Signed-off-by: Shangming Cai <caishangming@linux.alibaba.com>
shangmingc committed
Feb 19, 2025 at 22:13 UTC
5ae9f26a5ae2aa370cbee48924aa7ab885a45925
1 file changed
+2
-1
vllm/spec_decode/spec_decode_worker.py
+2
-1
@@ -10,6 +10,7 @@ import torch.nn as nn
10
11
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.logger import init_logger
16
from vllm.model_executor.layers.rejection_sampler import RejectionSampler
@@ -365,7 +366,7 @@ class SpecDecodeWorker(LoraNotSupportedWorkerBase):
366
target_lm_head_weight)
367
368
self._metrics.init_tensors(self.rank, device_type=self.device)
368
- self.spec_decode_sampler.init_tensors(self.rank,
369
+ self.spec_decode_sampler.init_tensors(get_tp_group().local_rank,
370
device_type=self.device)
371
372
scorer_cls: Type[SpeculativeScorer]