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

Signed-off-by: Joe Runde <Joseph.Runde@ibm.com> Signed-off-by: Amit Garg <mitgarg17495@gmail.com>
garg-amit · Oct 28, 2024 · ec70e8f · ec70e8f
1 parent ca794e0
commit ec70e8f
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(