[Bugfix] Fix TP > 1 for new granite (vllm-project#8544)

Signed-off-by: Joe Runde <Joseph.Runde@ibm.com>
aibrix · Sep 19, 2024 · 9d9c360 · 9d9c360
1 parent 35ede69
commit 9d9c360
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/vllm/model_executor/models/granite.py b/vllm/model_executor/models/granite.py
@@ -428,7 +428,8 @@ def compute_logits(
             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(