fix small bug

Author: YeAnbang

applications/ColossalChat/coati/distributed/grpo_consumer.py

@@ -294,7 +294,7 @@ def step(self, step_idx: int, pbar: Any, **kwargs) -> Optional[float]:
 
                         if self.booster.plugin.stage_manager.is_last_stage():
                             reference_action_log_probs = memory_efficient_logprob(
-                                reference_model_outputs["outputs"]["logits"],
+                                reference_model_outputs["outputs"]["logits"] / self.generate_config["temperature"],
                                 input_ids_forward_micro_batch,
                                 num_action,
                                 shard_config=self.plugin.shard_config,
@@ -321,7 +321,7 @@ def step(self, step_idx: int, pbar: Any, **kwargs) -> Optional[float]:
                     def _criterion(outputs, inputs):
                         action_logits = outputs.logits
                         action_log_probs = memory_efficient_logprob(
-                            action_logits,
+                            action_logits / self.generate_config["temperature"],
                             inputs["input_ids"],
                             num_action,
                             shard_config=self.plugin.shard_config,
@@ -388,7 +388,7 @@ def _criterion(outputs, inputs):
                             reference_model_logits / self.generate_config["temperature"],
                             input_ids_forward_micro_batch,
                             num_action,
-                            self.plugin.shard_config,
+                            shard_config=self.plugin.shard_config,
                         )
                         per_token_kl = (
                             torch.exp(reference_action_log_probs - action_log_probs)