From 98f9713399bd602ff954a83e6e6abcb4cf8b8864 Mon Sep 17 00:00:00 2001 From: Joe Runde Date: Tue, 17 Sep 2024 17:17:08 -0600 Subject: [PATCH] [Bugfix] Fix TP > 1 for new granite (#8544) Signed-off-by: Joe Runde --- vllm/model_executor/models/granite.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/granite.py b/vllm/model_executor/models/granite.py index b0325e8b..5f365bbc 100644 --- a/vllm/model_executor/models/granite.py +++ b/vllm/model_executor/models/granite.py @@ -428,7 +428,8 @@ class GraniteForCausalLM(nn.Module, SupportsLoRA): sampling_metadata: SamplingMetadata) -> Optional[torch.Tensor]: logits = self.logits_processor(self.lm_head, hidden_states, sampling_metadata) - logits /= self.config.logits_scaling + if logits is not None: + logits /= self.config.logits_scaling return logits def sample(