fix dpo metrics

Former-commit-id: 57029280da825a39fbf5a05097921b861f126669

src/llamafactory/train/dpo/trainer.py

@@ -250,20 +250,18 @@ class CustomDPOTrainer(DPOTrainer):
         if self.ftx_gamma > 1e-6:
             losses += self.ftx_gamma * sft_loss
 
-        reward_accuracies = (chosen_rewards > rejected_rewards).float()
-
         prefix = "eval_" if train_eval == "eval" else ""
-        metrics[f"{prefix}rewards/chosen"] = chosen_rewards.mean().cpu()
-        metrics[f"{prefix}rewards/rejected"] = rejected_rewards.mean().cpu()
-        metrics[f"{prefix}rewards/accuracies"] = reward_accuracies.mean().cpu()
-        metrics[f"{prefix}rewards/margins"] = (chosen_rewards - rejected_rewards).mean().cpu()
-        metrics[f"{prefix}logps/rejected"] = policy_rejected_logps.detach().mean().cpu()
-        metrics[f"{prefix}logps/chosen"] = policy_chosen_logps.detach().mean().cpu()
-        metrics[f"{prefix}logits/rejected"] = policy_rejected_logits.detach().mean().cpu()
-        metrics[f"{prefix}logits/chosen"] = policy_chosen_logits.detach().mean().cpu()
+        metrics[f"{prefix}rewards/chosen"] = chosen_rewards.mean().item()
+        metrics[f"{prefix}rewards/rejected"] = rejected_rewards.mean().item()
+        metrics[f"{prefix}rewards/accuracies"] = (chosen_rewards > rejected_rewards).float().mean().item()
+        metrics[f"{prefix}rewards/margins"] = (chosen_rewards - rejected_rewards).mean().item()
+        metrics[f"{prefix}logps/rejected"] = policy_chosen_logps.mean().item()
+        metrics[f"{prefix}logps/chosen"] = policy_rejected_logps.mean().item()
+        metrics[f"{prefix}logits/rejected"] = policy_chosen_logits.mean().item()
+        metrics[f"{prefix}logits/chosen"] = policy_rejected_logits.mean().item()
         if self.loss_type == "orpo":
-            metrics[f"{prefix}sft_loss"] = sft_loss.detach().mean().cpu()
-            metrics[f"{prefix}odds_ratio_loss"] = ((losses - sft_loss) / self.beta).detach().mean().cpu()
+            metrics[f"{prefix}sft_loss"] = sft_loss.mean().item()
+            metrics[f"{prefix}odds_ratio_loss"] = ((losses - sft_loss) / self.beta).mean().item()
 
         return losses.mean(), metrics
 
@@ -275,6 +273,36 @@ class CustomDPOTrainer(DPOTrainer):
         """
         loss = super().compute_loss(model, inputs, return_outputs)
         if is_transformers_version_equal_to_4_46() and kwargs.pop("num_items_in_batch", False):
-            loss /= self.args.gradient_accumulation_steps
+            if return_outputs:
+                return (loss[0] / self.args.gradient_accumulation_steps, *loss[1:])
+            else:
+                return loss / self.args.gradient_accumulation_steps
 
         return loss
+
+    @override
+    def log(self, logs: Dict[str, float]) -> None:
+        r"""
+        Log `logs` on the various objects watching training, including stored metrics.
+        """
+        # logs either has "loss" or "eval_loss"
+        train_eval = "train" if "loss" in logs else "eval"
+        # Add averaged stored metrics to logs
+        key_list, metric_list = [], []
+        for key, metrics in self._stored_metrics[train_eval].items():
+            key_list.append(key)
+            metric_list.append(torch.tensor(metrics, dtype=torch.float).to(self.accelerator.device).mean().item())
+
+        del self._stored_metrics[train_eval]
+        if len(metric_list) < 10:  # pad to for all reduce
+            for i in range(10 - len(metric_list)):
+                key_list.append(f"dummy_{i}")
+                metric_list.append(0.0)
+
+        metric_list = torch.tensor(metric_list, dtype=torch.float).to(self.accelerator.device)
+        metric_list = self.accelerator.reduce(metric_list, "mean").tolist()
+        for key, metric in zip(key_list, metric_list):  # add remaining items
+            if not key.startswith("dummy_"):
+                logs[key] = metric
+
+        return Trainer.log(self, logs)