Fix logprob does not handle temperature (THUDM#557)

Author: fzyzcjy

slime/backends/fsdp_utils/actor.py

@@ -230,7 +230,7 @@ def compute_log_prob(
                             model_args["pixel_values"] = batch["pixel_values"]
                         logits = self.model(**model_args).logits
                     batch[f"{store_prefix}log_probs"] = gather_log_probs_packed(
-                        logits, batch["tokens"], self.args.rollout_temperature
+                        logits, batch["tokens"], temperature=self.args.rollout_temperature
                     )
             return rollout_data
 
@@ -386,7 +386,9 @@ def train(self, rollout_id: int, rollout_data_ref: Box) -> None:
                 ).logits
 
             # Handle packed sequences
-            log_probs = gather_log_probs_packed(logits, packed_batch["tokens"], packed_batch["cu_seqlens"])
+            log_probs = gather_log_probs_packed(
+                logits, packed_batch["tokens"], packed_batch["cu_seqlens"], temperature=self.args.rollout_temperature
+            )
             packed_batch["cur_log_probs"] = log_probs
             unpacked_batches = unpack_sequences(packed_batch)
 
@@ -655,7 +657,10 @@ def gather_log_probs(logits: torch.Tensor, input_ids: torch.Tensor, rollout_temp
 
 
 def gather_log_probs_packed(
-    logits: torch.Tensor, input_ids: torch.Tensor, cu_seqlens: torch.Tensor | float | None = None
+    logits: torch.Tensor,
+    input_ids: torch.Tensor,
+    cu_seqlens: torch.Tensor | float | None = None,
+    temperature: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Gather next-token log probabilities for packed sequences.
 
@@ -674,6 +679,9 @@ def gather_log_probs_packed(
         logits = logits.squeeze(0)
         input_ids = input_ids.squeeze(0)
 
+    if temperature is not None:
+        logits = logits.div(temperature)
+
     # Shift for next-token prediction: logits[:-1] predicts input_ids[1:]
     log_probs = torch.log_softmax(logits[:-1], dim=-1)
     targets = input_ids[1:].to(device=log_probs.device)