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

Signed-off-by: Joe Runde <[email protected]> Signed-off-by: Alvant <[email protected]>
compressa-ai · Oct 26, 2024 · 128425b · 128425b
1 parent d5f7c8b
commit 128425b
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(