From f5cf4e956b409e69294ff1b21891f123fa2f43c2 Mon Sep 17 00:00:00 2001 From: youkaichao Date: Tue, 14 Jan 2025 15:01:07 +0800 Subject: [PATCH] fix neuron Signed-off-by: youkaichao --- vllm/worker/neuron_worker.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/worker/neuron_worker.py b/vllm/worker/neuron_worker.py index c5d00ec30f756..2c95363fa0b2a 100644 --- a/vllm/worker/neuron_worker.py +++ b/vllm/worker/neuron_worker.py @@ -26,6 +26,7 @@ def __init__( local_rank: int, rank: int, distributed_init_method: str, + is_driver_worker: bool = True, ) -> None: WorkerBase.__init__(self, vllm_config=vllm_config) self.local_rank = local_rank @@ -38,7 +39,7 @@ def __init__( self.model_runner: NeuronModelRunner = NeuronModelRunner( vllm_config=vllm_config) - self.is_driver_worker = True + self.is_driver_worker = is_driver_worker def execute_model( self,