Align KTO with DPO: Use _metrics attribute

trl/experimental/kto/kto_trainer.py

@@ -339,9 +339,6 @@ def __init__(
         else:
             self.use_dpo_data_collator = False
 
-        # metric
-        self._stored_metrics = defaultdict(lambda: defaultdict(list))
-
         # Training arguments
         self.beta = args.beta
         self.precompute_ref_log_probs = args.precompute_ref_log_probs
@@ -417,6 +414,9 @@ def __init__(
             if self.ref_model is not None:
                 disable_dropout_in_model(self.ref_model)
 
+        # Initialize the metrics
+        self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)}
+
         # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the
         # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set
         # self.model_accepts_loss_kwargs to False to enable scaling.
@@ -1139,7 +1139,7 @@ def compute_loss(
 
     def store_metrics(self, metrics: dict[str, float], train_eval: Literal["train", "eval"] = "train") -> None:
         for key, value in metrics.items():
-            self._stored_metrics[train_eval][key].append(value)
+            self._metrics[train_eval][key].append(value)
 
     def _get_train_sampler(self, dataset: Dataset | None = None) -> torch.utils.data.Sampler | None:
         if dataset is None:
@@ -1199,23 +1199,22 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
         prefix = "eval_" if train_eval == "eval" else ""
         # accumulate average metrics from sums and lengths
         for split in ["chosen", "rejected"]:
-            if f"count/{split}" in self._stored_metrics[train_eval]:
-                count_sum = torch.Tensor(self._stored_metrics[train_eval][f"count/{split}"]).sum().item()
+            if f"count/{split}" in self._metrics[train_eval]:
+                count_sum = torch.Tensor(self._metrics[train_eval][f"count/{split}"]).sum().item()
                 for metric in ["rewards", "logps", "logits"]:
                     logs[f"{prefix}{metric}/{split}"] = (
-                        torch.Tensor(self._stored_metrics[train_eval][f"{metric}/{split}_sum"]).sum().item()
-                        / count_sum
+                        torch.Tensor(self._metrics[train_eval][f"{metric}/{split}_sum"]).sum().item() / count_sum
                     )
                     # delete obsolete metric
-                    del self._stored_metrics[train_eval][f"{metric}/{split}_sum"]
-                del self._stored_metrics[train_eval][f"count/{split}"]
+                    del self._metrics[train_eval][f"{metric}/{split}_sum"]
+                del self._metrics[train_eval][f"count/{split}"]
         # calculate reward margin
         if f"{prefix}rewards/chosen" in logs and f"{prefix}rewards/rejected" in logs:
             logs[f"{prefix}rewards/margins"] = logs[f"{prefix}rewards/chosen"] - logs[f"{prefix}rewards/rejected"]
         # Add averaged stored metrics to logs
-        for key, metrics in self._stored_metrics[train_eval].items():
+        for key, metrics in self._metrics[train_eval].items():
             logs[f"{prefix}{key}"] = torch.Tensor(metrics).mean().item()
-        del self._stored_metrics[train_eval]
+        self._metrics[train_eval].clear()
         return super().log(logs, start_time)
 
     # Ensure the model card is saved along with the checkpoint