fix: granite logit scale in logit computation (#1054)

PygmalionAI · Dec 27, 2024 · 11f49b5 · 11f49b5
1 parent 1264e0b
commit 11f49b5
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/aphrodite/modeling/models/granite.py b/aphrodite/modeling/models/granite.py
@@ -439,11 +439,12 @@ def forward(
         return model_output
 
     def compute_logits(
-        self, hidden_states: torch.Tensor, sampling_metadata: SamplingMetadata
-    ) -> Optional[torch.Tensor]:
-        logits = self.logits_processor(
-            self.lm_head, hidden_states, sampling_metadata
-        )
+            self, hidden_states: torch.Tensor,
+            sampling_metadata: SamplingMetadata) -> Optional[torch.Tensor]:
+        logits = self.logits_processor(self.lm_head, hidden_states,
+                                       sampling_metadata)
+        if logits is not None:
+            logits /= self.config.logits_scaling
         return logits
 
     def sample(