From b89c094ed14448fa3f8a539dd4e9744a96c5b9ad Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?=
 <quentin.gallouedec@huggingface.co>
Date: Mon, 24 Aug 2026 16:35:32 +0000
Subject: [PATCH 1/6] Accumulate SFT metrics as running sums, reduce once per
 logging window

---
 trl/trainer/distillation_trainer.py | 26 ++++++++---
 trl/trainer/sft_trainer.py          | 70 ++++++++++++++++-------------
 2 files changed, 58 insertions(+), 38 deletions(-)

diff --git a/trl/trainer/distillation_trainer.py b/trl/trainer/distillation_trainer.py
index 74899c00978..ffa1bbb37fb 100644
--- a/trl/trainer/distillation_trainer.py
+++ b/trl/trainer/distillation_trainer.py
@@ -748,6 +748,9 @@ def __init__(
 
         # Metrics & Logging
         self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)}
+        # Per-mode running sums of on-device metric accumulators. `compute_loss` only adds local tensors here (no
+        # collective, no host sync); `log()` reduces them across ranks in a single collective and resets them.
+        self._metric_sums = {"train": defaultdict(int), "eval": defaultdict(int)}
         self._total_train_tokens = 0
         self._current_train_step_time = 0.0
         self.log_completions = args.log_completions
@@ -1374,15 +1377,13 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
             model, unwrapped_student, self._compute_loss, unwrapped_student, inputs, num_items_in_batch
         )
 
-        # Log the mean per-token student entropy (in nats). The reduction runs here, after `_forward_redirection`
-        # returns, so the `gather_for_metrics` collective does not run inside the DDP/FSDP-wrapped forward (a hang/
-        # ordering risk). The Liger path produces no entropy, so it logs none. Mirrors `SFTTrainer.compute_loss`.
+        # Log the mean per-token student entropy (in nats). Accumulated as per-rank running sums; the cross-rank
+        # reduction happens in `log()`, once per logging window, so no collective runs inside the DDP/FSDP-wrapped
+        # forward. The Liger path produces no entropy, so it logs none. Mirrors `SFTTrainer.compute_loss`.
         if entropy_sum is not None:
             mode = "train" if self.model.training else "eval"
-            num_valid_tokens = self.accelerator.gather_for_metrics(num_valid_tokens).sum()
-            entropy_sum = self.accelerator.gather_for_metrics(entropy_sum).sum()
-            entropy = (entropy_sum / num_valid_tokens).item() if num_valid_tokens > 0 else 0.0
-            self._metrics[mode]["entropy"].append(entropy)
+            self._metric_sums[mode]["entropy_sum"] += entropy_sum.detach()
+            self._metric_sums[mode]["total_tokens"] += num_valid_tokens.detach()
 
         return (loss, None) if return_outputs else loss
 
@@ -1539,6 +1540,16 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
             valid = [v for v in val if not math.isnan(v)]
             metrics[key] = sum(valid) / len(valid) if valid else None
 
+        sums = self._metric_sums[mode]
+        # Entropy is accumulated in `compute_loss` as per-rank running sums. Reduce it across ranks here, in a
+        # single collective per logging window. The logged value is token-weighted over the window: a ratio of
+        # global sums, not a mean of per-step ratios. Mirrors `SFTTrainer.log`.
+        if sums:
+            values = torch.stack([value.double() for value in sums.values()])
+            totals = dict(zip(sums.keys(), self.accelerator.reduce(values, reduction="sum").tolist(), strict=True))
+            total_tokens = totals["total_tokens"]
+            metrics["entropy"] = totals["entropy_sum"] / total_tokens if total_tokens > 0 else 0.0
+
         # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs`
         # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format.
         if mode == "eval":
@@ -1547,6 +1558,7 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
         logs.update(metrics)
         super().log(logs, start_time)
         self._metrics[mode].clear()
+        self._metric_sums[mode].clear()
 
         if self.accelerator.is_main_process and self.log_completions:
             if is_rich_available():
diff --git a/trl/trainer/sft_trainer.py b/trl/trainer/sft_trainer.py
index 2c1c2950ed1..0fb4a4e2213 100644
--- a/trl/trainer/sft_trainer.py
+++ b/trl/trainer/sft_trainer.py
@@ -1401,7 +1401,9 @@ def __init__(
             text_config.router_aux_loss_coef = self.args.router_aux_loss_coef
 
         # Initialize the metrics
-        self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)}
+        # Per-mode running sums of on-device metric accumulators. `compute_loss` only adds local tensors here (no
+        # collective, no host sync); `log()` reduces them across ranks in a single collective and resets them.
+        self._metric_sums = {"train": defaultdict(int), "eval": defaultdict(int)}
         self._total_train_tokens = 0
 
         # Add tags to the model
@@ -1790,10 +1792,8 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
             # PEFT (PromptTuning, P-Tuning) prepends `-100`-padded virtual tokens before delegating into the patched
             # forward, so the valid-token count over the padded labels can differ from the un-padded `labels[..., 1:]`
             # count by up to one per sequence; using the patched output keeps numerator and denominator aligned.
-            num_valid = self.accelerator.gather_for_metrics(outputs.num_valid_tokens).sum()
-            entropy_sum = self.accelerator.gather_for_metrics(outputs.entropy_sum).sum()
-            entropy = (entropy_sum / num_valid).item() if num_valid > 0 else 0.0
-            self._metrics[mode]["entropy"].append(entropy)
+            self._metric_sums[mode]["entropy_sum"] += outputs.entropy_sum.detach()
+            self._metric_sums[mode]["total_tokens"] += outputs.num_valid_tokens.detach()
         elif not self.args.use_liger_kernel:  # liger doesn't return logits
             with torch.no_grad():
                 if "shift_labels" in inputs:
@@ -1819,39 +1819,27 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
                 total_tokens = mask.sum()
                 correct_predictions = (predictions == shift_labels) & mask
                 correct_tokens = correct_predictions.sum()
-
-                # Gather counts across ranks and weight-average
-                entropy_sum = self.accelerator.gather_for_metrics(entropy_sum).sum()
-                total_tokens = self.accelerator.gather_for_metrics(total_tokens).sum()
-                correct_tokens = self.accelerator.gather_for_metrics(correct_tokens)
-                entropy = (entropy_sum / total_tokens).item() if total_tokens > 0 else 0.0
-
-                total_sum = total_tokens.sum()
-                accuracy = (correct_tokens.sum() / total_sum).item() if total_sum > 0 else 0.0
-            self._metrics[mode]["entropy"].append(entropy)
-            self._metrics[mode]["mean_token_accuracy"].append(accuracy)
+            self._metric_sums[mode]["entropy_sum"] += entropy_sum
+            self._metric_sums[mode]["total_tokens"] += total_tokens
+            self._metric_sums[mode]["correct_tokens"] += correct_tokens
 
         if mode == "train":
             # When using padding-free, the attention_mask is not present in the inputs, instead we have cu_seq_lens_q,
             # cu_seq_lens_k, and max_length_k, max_length_q and position_ids.
             if "attention_mask" in inputs:
-                num_tokens_in_batch = self.accelerator.gather_for_metrics(inputs["attention_mask"].sum()).sum().item()
+                num_tokens_in_batch = inputs["attention_mask"].sum()
             elif "position_ids" in inputs:
-                local_num_tokens = torch.tensor(inputs["position_ids"].size(1), device=inputs["position_ids"].device)
-                num_tokens_in_batch = self.accelerator.gather_for_metrics(local_num_tokens).sum().item()
+                num_tokens_in_batch = torch.tensor(inputs["position_ids"].size(1), device=inputs["position_ids"].device)
             else:
                 raise ValueError("Expected 'attention_mask' or 'position_ids' in inputs.")
-            self._total_train_tokens += num_tokens_in_batch
-        self._metrics[mode]["num_tokens"] = [self._total_train_tokens]
+            self._metric_sums[mode]["num_tokens_in_batch"] += num_tokens_in_batch
 
         if self.args.loss_type == "chunked_nll":
-            correct = self.accelerator.gather_for_metrics(outputs.num_correct_tokens).sum()
-            accuracy = (correct / num_valid).item() if num_valid > 0 else 0.0
-            self._metrics[mode]["mean_token_accuracy"].append(accuracy)
+            self._metric_sums[mode]["correct_tokens"] += outputs.num_correct_tokens.detach()
         elif self.args.use_liger_kernel:
             if hasattr(outputs, "token_accuracy") and outputs.token_accuracy is not None:
-                token_accuracy = self.accelerator.gather_for_metrics(outputs.token_accuracy).mean().item()
-                self._metrics[mode]["mean_token_accuracy"].append(token_accuracy)
+                self._metric_sums[mode]["token_accuracy_sum"] += outputs.token_accuracy.detach()
+                self._metric_sums[mode]["token_accuracy_count"] += torch.ones_like(outputs.token_accuracy)
             else:
                 warnings.warn(
                     "liger-kernel did not return token_accuracy when requested. The mean_token_accuracy metric will "
@@ -1860,9 +1848,8 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
                 )
         # Log auxiliary loss if enabled (applies to both Liger and non-Liger)
         if self.aux_loss_enabled:
-            aux_loss = outputs.aux_loss
-            aux_loss = self.accelerator.gather_for_metrics(aux_loss).mean().item()
-            self._metrics[mode]["aux_loss"].append(aux_loss)
+            self._metric_sums[mode]["aux_loss_sum"] += outputs.aux_loss.detach()
+            self._metric_sums[mode]["aux_loss_count"] += torch.ones_like(outputs.aux_loss)
 
         return (loss, outputs) if return_outputs else loss
 
@@ -1878,7 +1865,28 @@ def training_step(self, *args, **kwargs):
 
     def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
         mode = "train" if self.model.training else "eval"
-        metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()}  # average the metrics
+        sums = self._metric_sums[mode]
+
+        # Metrics are accumulated in `compute_loss` as per-rank running sums. Reduce them across ranks here, in a
+        # single collective per logging window. Token-level metrics (entropy, accuracy) are token-weighted over the
+        # window: a ratio of global sums, not a mean of per-step ratios.
+        metrics = {}
+        if sums:
+            values = torch.stack([value.double() for value in sums.values()])
+            totals = dict(zip(sums.keys(), self.accelerator.reduce(values, reduction="sum").tolist(), strict=True))
+            if "entropy_sum" in totals:
+                total_tokens = totals["total_tokens"]
+                metrics["entropy"] = totals["entropy_sum"] / total_tokens if total_tokens > 0 else 0.0
+                metrics["mean_token_accuracy"] = totals["correct_tokens"] / total_tokens if total_tokens > 0 else 0.0
+            if "token_accuracy_sum" in totals:
+                metrics["mean_token_accuracy"] = totals["token_accuracy_sum"] / totals["token_accuracy_count"]
+            if "aux_loss_sum" in totals:
+                metrics["aux_loss"] = totals["aux_loss_sum"] / totals["aux_loss_count"]
+            # `num_tokens` advances only when a train-mode log folds in the pending sums, so an eval log between two
+            # train logs can lag by up to one logging window.
+            if mode == "train" and "num_tokens_in_batch" in totals:
+                self._total_train_tokens += int(totals["num_tokens_in_batch"])
+            metrics["num_tokens"] = self._total_train_tokens
 
         # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs`
         # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format.
@@ -1887,7 +1895,7 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
 
         logs.update(metrics)
         super().log(logs, start_time)
-        self._metrics[mode].clear()
+        self._metric_sums[mode].clear()
 
     # Ensure the model card is saved along with the checkpoint
     def _save_checkpoint(self, model, trial):

From 989a5cb14f3211c5c2076b7c009542fb7de6e80a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?=
 <quentin.gallouedec@huggingface.co>
Date: Mon, 24 Aug 2026 17:19:49 +0000
Subject: [PATCH 2/6] Generalize the log-time reduction to <name>_sum /
 <name>_count pairs

---
 trl/trainer/distillation_trainer.py | 21 +++++++++------
 trl/trainer/sft_trainer.py          | 41 +++++++++++++++--------------
 2 files changed, 34 insertions(+), 28 deletions(-)

diff --git a/trl/trainer/distillation_trainer.py b/trl/trainer/distillation_trainer.py
index ffa1bbb37fb..f07dcf6d0fc 100644
--- a/trl/trainer/distillation_trainer.py
+++ b/trl/trainer/distillation_trainer.py
@@ -1383,7 +1383,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
         if entropy_sum is not None:
             mode = "train" if self.model.training else "eval"
             self._metric_sums[mode]["entropy_sum"] += entropy_sum.detach()
-            self._metric_sums[mode]["total_tokens"] += num_valid_tokens.detach()
+            self._metric_sums[mode]["entropy_count"] += num_valid_tokens.detach()
 
         return (loss, None) if return_outputs else loss
 
@@ -1540,15 +1540,20 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
             valid = [v for v in val if not math.isnan(v)]
             metrics[key] = sum(valid) / len(valid) if valid else None
 
+        # Metrics are accumulated in `compute_loss` as per-rank running sums. Aggregate them across ranks here, in
+        # a single collective per logging window, then compute each `<name>` metric as `<name>_sum / <name>_count`,
+        # i.e. weighted by whatever the count counts (tokens, batches). Keys are sorted so that every rank stacks
+        # them in the same order.
         sums = self._metric_sums[mode]
-        # Entropy is accumulated in `compute_loss` as per-rank running sums. Reduce it across ranks here, in a
-        # single collective per logging window. The logged value is token-weighted over the window: a ratio of
-        # global sums, not a mean of per-step ratios. Mirrors `SFTTrainer.log`.
         if sums:
-            values = torch.stack([value.double() for value in sums.values()])
-            totals = dict(zip(sums.keys(), self.accelerator.reduce(values, reduction="sum").tolist(), strict=True))
-            total_tokens = totals["total_tokens"]
-            metrics["entropy"] = totals["entropy_sum"] / total_tokens if total_tokens > 0 else 0.0
+            keys = sorted(sums)
+            values = torch.stack([sums[key].double() for key in keys])
+            totals = dict(zip(keys, self.accelerator.reduce(values, reduction="sum").tolist(), strict=True))
+            for key in keys:
+                if key.endswith("_sum"):
+                    name = key.removesuffix("_sum")
+                    count = totals[name + "_count"]
+                    metrics[name] = totals[key] / count if count > 0 else 0.0
 
         # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs`
         # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format.
diff --git a/trl/trainer/sft_trainer.py b/trl/trainer/sft_trainer.py
index 0fb4a4e2213..e2ddb96600c 100644
--- a/trl/trainer/sft_trainer.py
+++ b/trl/trainer/sft_trainer.py
@@ -1793,7 +1793,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
             # forward, so the valid-token count over the padded labels can differ from the un-padded `labels[..., 1:]`
             # count by up to one per sequence; using the patched output keeps numerator and denominator aligned.
             self._metric_sums[mode]["entropy_sum"] += outputs.entropy_sum.detach()
-            self._metric_sums[mode]["total_tokens"] += outputs.num_valid_tokens.detach()
+            self._metric_sums[mode]["entropy_count"] += outputs.num_valid_tokens.detach()
         elif not self.args.use_liger_kernel:  # liger doesn't return logits
             with torch.no_grad():
                 if "shift_labels" in inputs:
@@ -1820,8 +1820,9 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
                 correct_predictions = (predictions == shift_labels) & mask
                 correct_tokens = correct_predictions.sum()
             self._metric_sums[mode]["entropy_sum"] += entropy_sum
-            self._metric_sums[mode]["total_tokens"] += total_tokens
-            self._metric_sums[mode]["correct_tokens"] += correct_tokens
+            self._metric_sums[mode]["entropy_count"] += total_tokens
+            self._metric_sums[mode]["mean_token_accuracy_sum"] += correct_tokens
+            self._metric_sums[mode]["mean_token_accuracy_count"] += total_tokens
 
         if mode == "train":
             # When using padding-free, the attention_mask is not present in the inputs, instead we have cu_seq_lens_q,
@@ -1835,11 +1836,12 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
             self._metric_sums[mode]["num_tokens_in_batch"] += num_tokens_in_batch
 
         if self.args.loss_type == "chunked_nll":
-            self._metric_sums[mode]["correct_tokens"] += outputs.num_correct_tokens.detach()
+            self._metric_sums[mode]["mean_token_accuracy_sum"] += outputs.num_correct_tokens.detach()
+            self._metric_sums[mode]["mean_token_accuracy_count"] += outputs.num_valid_tokens.detach()
         elif self.args.use_liger_kernel:
             if hasattr(outputs, "token_accuracy") and outputs.token_accuracy is not None:
-                self._metric_sums[mode]["token_accuracy_sum"] += outputs.token_accuracy.detach()
-                self._metric_sums[mode]["token_accuracy_count"] += torch.ones_like(outputs.token_accuracy)
+                self._metric_sums[mode]["mean_token_accuracy_sum"] += outputs.token_accuracy.detach()
+                self._metric_sums[mode]["mean_token_accuracy_count"] += torch.ones_like(outputs.token_accuracy)
             else:
                 warnings.warn(
                     "liger-kernel did not return token_accuracy when requested. The mean_token_accuracy metric will "
@@ -1865,23 +1867,22 @@ def training_step(self, *args, **kwargs):
 
     def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
         mode = "train" if self.model.training else "eval"
-        sums = self._metric_sums[mode]
 
-        # Metrics are accumulated in `compute_loss` as per-rank running sums. Reduce them across ranks here, in a
-        # single collective per logging window. Token-level metrics (entropy, accuracy) are token-weighted over the
-        # window: a ratio of global sums, not a mean of per-step ratios.
+        # Metrics are accumulated in `compute_loss` as per-rank running sums. Aggregate them across ranks here, in
+        # a single collective per logging window, then compute each `<name>` metric as `<name>_sum / <name>_count`,
+        # i.e. weighted by whatever the count counts (tokens, batches). Keys are sorted so that every rank stacks
+        # them in the same order.
         metrics = {}
+        sums = self._metric_sums[mode]
         if sums:
-            values = torch.stack([value.double() for value in sums.values()])
-            totals = dict(zip(sums.keys(), self.accelerator.reduce(values, reduction="sum").tolist(), strict=True))
-            if "entropy_sum" in totals:
-                total_tokens = totals["total_tokens"]
-                metrics["entropy"] = totals["entropy_sum"] / total_tokens if total_tokens > 0 else 0.0
-                metrics["mean_token_accuracy"] = totals["correct_tokens"] / total_tokens if total_tokens > 0 else 0.0
-            if "token_accuracy_sum" in totals:
-                metrics["mean_token_accuracy"] = totals["token_accuracy_sum"] / totals["token_accuracy_count"]
-            if "aux_loss_sum" in totals:
-                metrics["aux_loss"] = totals["aux_loss_sum"] / totals["aux_loss_count"]
+            keys = sorted(sums)
+            values = torch.stack([sums[key].double() for key in keys])
+            totals = dict(zip(keys, self.accelerator.reduce(values, reduction="sum").tolist(), strict=True))
+            for key in keys:
+                if key.endswith("_sum"):
+                    name = key.removesuffix("_sum")
+                    count = totals[name + "_count"]
+                    metrics[name] = totals[key] / count if count > 0 else 0.0
             # `num_tokens` advances only when a train-mode log folds in the pending sums, so an eval log between two
             # train logs can lag by up to one logging window.
             if mode == "train" and "num_tokens_in_batch" in totals:

From 715c7c0c32c0ee41cf74bd12b4b18bcdaf2ff741 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?=
 <quentin.gallouedec@huggingface.co>
Date: Mon, 24 Aug 2026 17:34:48 +0000
Subject: [PATCH 3/6] Aggregate GRPO and RLOO metrics at log time

---
 tests/test_grpo_trainer.py  |  14 +++-
 trl/trainer/grpo_trainer.py | 143 +++++++++++++++++++++++-------------
 trl/trainer/rloo_trainer.py |  84 ++++++++++++++-------
 3 files changed, 159 insertions(+), 82 deletions(-)

diff --git a/tests/test_grpo_trainer.py b/tests/test_grpo_trainer.py
index 834833b3dc3..372903289a1 100644
--- a/tests/test_grpo_trainer.py
+++ b/tests/test_grpo_trainer.py
@@ -1632,15 +1632,21 @@ def generate_with_one_unscorable_token(prompts):
 
         trainer._generate = generate_with_one_unscorable_token
 
-        # Snapshot the divergence metric as soon as it is produced. Reading `_metrics` after training is useless,
-        # because the dict is cleared on every log.
+        # Snapshot the divergence metric as soon as it is produced. Reading the accumulators after training is
+        # useless, because they are cleared on every log.
         original_score = trainer._generate_and_score_completions
         recorded_metrics = []
 
         def record_metrics(inputs):
             outputs = original_score(inputs)
-            for key in ["sampling/sampling_logp_difference/mean", "sampling/sampling_logp_difference/max"]:
-                recorded_metrics.extend((key, value) for value in trainer._metrics["train"][key])
+            sums = trainer._metric_sums["train"]
+            mean = (
+                sums["sampling/sampling_logp_difference/mean_sum"]
+                / sums["sampling/sampling_logp_difference/mean_count"]
+            )
+            recorded_metrics.append(("sampling/sampling_logp_difference/mean", mean.item()))
+            max_delta = trainer._metric_maxs["train"]["sampling/sampling_logp_difference/max"]
+            recorded_metrics.append(("sampling/sampling_logp_difference/max", max_delta.item()))
             return outputs
 
         trainer._generate_and_score_completions = record_metrics
diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index de126792d2f..3b9ee438086 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -84,8 +84,6 @@
     get_config_model_id,
     identity,
     maybe_gather_lm_head_ctx,
-    nanmax,
-    nanmin,
     nanstd,
     pad,
     print_prompt_completions_sample,
@@ -1058,6 +1056,12 @@ def cast_outputs_to_original_dtype(module, args, output):
 
         # Initialize the metrics
         self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)}
+        # Per-mode running sums/extrema of on-device metric accumulators. They only receive local tensors (no
+        # collective, no host sync); `log()` aggregates them across ranks in one collective per kind and resets
+        # them. `_metrics` above keeps plain floats for values that are already identical on every rank.
+        self._metric_sums = {"train": defaultdict(int), "eval": defaultdict(int)}
+        self._metric_mins = {"train": {}, "eval": {}}
+        self._metric_maxs = {"train": {}, "eval": {}}
         self._total_train_tokens = 0
         self._current_train_step_time = 0.0
         self.log_completions = args.log_completions
@@ -2865,56 +2869,52 @@ def _generate_and_score_completions(
             self._logs["extra"][column].extend(gather_object(self._pending_extra_logs[column]))
         self._pending_extra_logs.clear()
 
-        # Flush user-logged metrics (from log_metric), averaging across processes.
-        # Keys must be sorted so that all ranks call accelerator.gather in the same order, otherwise values
-        # get mis-attributed across metrics (dict insertion order may differ between processes).
-        for name in sorted(self._pending_metrics):
+        # Flush user-logged metrics (from log_metric), accumulated locally and averaged across processes at `log()`
+        # time. Every rank must log the same metric names, otherwise the log-time aggregation mismatches ranks
+        # (the same requirement the previous per-step gather had).
+        for name in self._pending_metrics:
             values = self._pending_metrics[name]
             local_mean = sum(values) / len(values)
-            global_mean = self.accelerator.gather(torch.tensor(local_mean, device=device)).mean().item()
-            self._metrics[mode][name].append(global_mean)
+            self._metric_sums[mode][f"{name}_sum"] += torch.tensor(local_mean, device=device)
+            self._metric_sums[mode][f"{name}_count"] += torch.ones((), device=device)
         self._pending_metrics.clear()
 
         if images is not None and self.log_multimodal:
             self._logs["images"].extend(gather_object(images))
 
         if self.use_vllm and self.vllm_importance_sampling_correction:
+            sums = self._metric_sums[mode]
+            mins, maxs = self._metric_mins[mode], self._metric_maxs[mode]
             delta = torch.abs(old_per_token_logps - sampling_per_token_logps)
             mask = completion_mask.bool() if tool_mask is None else (completion_mask * tool_mask).bool()
             # Tokens vLLM could not score carry NaN, so exclude them rather than let them turn the reported
             # divergence into NaN. Counting them as zero instead would understate the divergence.
-            delta = delta[mask & ~torch.isnan(delta)]
-            mean_delta = torch.mean(delta) if delta.numel() > 0 else torch.tensor(0.0, device=device)
-            max_delta = torch.max(delta) if delta.numel() > 0 else torch.tensor(0.0, device=device)
-            self._metrics[mode]["sampling/sampling_logp_difference/mean"].append(
-                self.accelerator.gather(mean_delta).mean().item()
-            )
-            self._metrics[mode]["sampling/sampling_logp_difference/max"].append(
-                self.accelerator.gather(max_delta).max().item()
-            )
+            delta_valid = mask & ~torch.isnan(delta)
+            sums["sampling/sampling_logp_difference/mean_sum"] += torch.where(delta_valid, delta, 0.0).sum()
+            sums["sampling/sampling_logp_difference/mean_count"] += delta_valid.sum()
+            max_delta = torch.where(delta_valid, delta, -torch.inf).max()
+            key = "sampling/sampling_logp_difference/max"
+            maxs[key] = torch.maximum(maxs.setdefault(key, max_delta), max_delta)
             if sequence_level_is:
                 flat_is_ratio = vllm_importance_sampling_ratio.flatten()
             else:
                 flat_is_ratio = vllm_importance_sampling_ratio[mask]
 
-            min_importance_sampling_ratio = (
-                torch.min(flat_is_ratio) if flat_is_ratio.numel() > 0 else torch.tensor(0.0, device=device)
-            )
-            mean_importance_sampling_ratio = (
-                torch.mean(flat_is_ratio) if flat_is_ratio.numel() > 0 else torch.tensor(0.0, device=device)
-            )
-            max_importance_sampling_ratio = (
-                torch.max(flat_is_ratio) if flat_is_ratio.numel() > 0 else torch.tensor(0.0, device=device)
-            )
-            self._metrics[mode]["sampling/importance_sampling_ratio/min"].append(
-                nanmin(self.accelerator.gather(min_importance_sampling_ratio)).item()
-            )
-            self._metrics[mode]["sampling/importance_sampling_ratio/mean"].append(
-                self.accelerator.gather(mean_importance_sampling_ratio).nanmean().item()
-            )
-            self._metrics[mode]["sampling/importance_sampling_ratio/max"].append(
-                nanmax(self.accelerator.gather(max_importance_sampling_ratio)).item()
-            )
+            is_ratio_valid = ~torch.isnan(flat_is_ratio)
+            sums["sampling/importance_sampling_ratio/mean_sum"] += torch.where(
+                is_ratio_valid, flat_is_ratio, 0.0
+            ).sum()
+            sums["sampling/importance_sampling_ratio/mean_count"] += is_ratio_valid.sum()
+            if flat_is_ratio.numel() > 0:
+                min_is_ratio = torch.where(is_ratio_valid, flat_is_ratio, torch.inf).min()
+                max_is_ratio = torch.where(is_ratio_valid, flat_is_ratio, -torch.inf).max()
+            else:
+                min_is_ratio = torch.tensor(torch.inf, device=device)
+                max_is_ratio = torch.tensor(-torch.inf, device=device)
+            key = "sampling/importance_sampling_ratio/min"
+            mins[key] = torch.minimum(mins.setdefault(key, min_is_ratio), min_is_ratio)
+            key = "sampling/importance_sampling_ratio/max"
+            maxs[key] = torch.maximum(maxs.setdefault(key, max_is_ratio), max_is_ratio)
 
         output = {
             "prompt_ids": prompt_ids,
@@ -3006,8 +3006,10 @@ def compute_liger_loss(self, unwrapped_model, inputs):
 
         mode = "train" if self.model.training else "eval"
         if self.beta != 0.0:
-            self._metrics[mode]["kl"].append(self.accelerator.gather(mean_kl).mean().item())
-        self._metrics[mode]["clip_ratio"].append(self.accelerator.gather(clip_ratio).mean().item())
+            self._metric_sums[mode]["kl_sum"] += mean_kl.detach()
+            self._metric_sums[mode]["kl_count"] += torch.ones_like(mean_kl)
+        self._metric_sums[mode]["clip_ratio_sum"] += clip_ratio.detach()
+        self._metric_sums[mode]["clip_ratio_count"] += torch.ones_like(clip_ratio)
         # DAPO/CISPO/VESPO normalize by num_items_in_batch / num_processes (applied internally by
         # the Liger loss), then need a `current_gradient_accumulation_steps / steps_per_generation`
         # rescale to land on the per-window token-mean — matching the non-Liger path
@@ -3299,7 +3301,9 @@ def _compute_loss(self, model, inputs):
 
             loss = loss - apply_coef * entropy_loss
 
-            self._metrics[mode]["policy_loss"].append(self.accelerator.gather(policy_loss).nanmean().item())
+            policy_loss_valid = ~torch.isnan(policy_loss)
+            self._metric_sums[mode]["policy_loss_sum"] += torch.where(policy_loss_valid, policy_loss, 0.0).detach()
+            self._metric_sums[mode]["policy_loss_count"] += policy_loss_valid.float()
 
             # Adaptive update. Gated on train mode so evaluation cannot mutate the entropy controller state.
             if self.use_adaptive_entropy and mode == "train":
@@ -3337,7 +3341,8 @@ def _compute_loss(self, model, inputs):
         if self.aux_loss_enabled:
             normalizer = self.current_gradient_accumulation_steps if mode == "train" else 1.0
             loss = loss + self.router_aux_loss_coef * aux_loss / normalizer
-            self._metrics[mode]["aux_loss"].append(self.accelerator.gather_for_metrics(aux_loss).mean().item())
+            self._metric_sums[mode]["aux_loss_sum"] += aux_loss.detach()
+            self._metric_sums[mode]["aux_loss_count"] += torch.ones_like(aux_loss)
 
         # Log the metrics
         def masked_seq_mean(x):
@@ -3345,36 +3350,40 @@ def masked_seq_mean(x):
                 return x.squeeze(1)
             return (x * mask).sum(-1) / mask.sum(-1)
 
-        def global_masked_mean(x):
+        def accumulate_masked_mean(name, x):
+            x = x.detach()
             if x.shape[1] == 1:  # when importance_sampling_level == "sequence": one value per sequence
                 local_sum, local_count = x.sum(), torch.tensor(float(x.shape[0]), device=x.device)
             else:
                 local_sum, local_count = (x * mask).sum(), mask.sum().float()
-            totals = self.accelerator.reduce(torch.stack([local_sum, local_count]), reduction="sum")
-            return (totals[0] / totals[1].clamp(min=1.0)).item()
+            self._metric_sums[mode][f"{name}_sum"] += local_sum
+            self._metric_sums[mode][f"{name}_count"] += local_count
 
         if self.beta != 0.0:
-            self._metrics[mode]["kl"].append(global_masked_mean(per_token_kl))
+            accumulate_masked_mean("kl", per_token_kl)
 
-        self._metrics[mode]["entropy"].append(global_masked_mean(entropies))
+        accumulate_masked_mean("entropy", entropies)
 
         if self.loss_type in ["grpo", "bnpo", "dr_grpo", "dapo", "luspo"]:
             # Compute the clipped probability ratios
             is_low_clipped = (coef_1 < 1 - self.epsilon_low) & (advantages < 0)
             is_high_clipped = (coef_1 > 1 + self.epsilon_high) & (advantages > 0)
             is_region_clipped = is_low_clipped | is_high_clipped
-            self._metrics[mode]["clip_ratio/low_mean"].append(global_masked_mean(is_low_clipped.float()))
-            self._metrics[mode]["clip_ratio/high_mean"].append(global_masked_mean(is_high_clipped.float()))
-            self._metrics[mode]["clip_ratio/region_mean"].append(global_masked_mean(is_region_clipped.float()))
-            gathered_low_clip = self.accelerator.gather(masked_seq_mean(is_low_clipped.float()))
-            self._metrics[mode]["clip_ratio/low_min"].append(nanmin(gathered_low_clip).item())
-            gathered_high_clip = self.accelerator.gather(masked_seq_mean(is_high_clipped.float()))
-            self._metrics[mode]["clip_ratio/high_max"].append(nanmax(gathered_high_clip).item())
+            accumulate_masked_mean("clip_ratio/low_mean", is_low_clipped.float())
+            accumulate_masked_mean("clip_ratio/high_mean", is_high_clipped.float())
+            accumulate_masked_mean("clip_ratio/region_mean", is_region_clipped.float())
+            mins, maxs = self._metric_mins[mode], self._metric_maxs[mode]
+            low_clip_seq = masked_seq_mean(is_low_clipped.float())
+            low_min = torch.where(low_clip_seq.isnan(), torch.inf, low_clip_seq).min()
+            mins["clip_ratio/low_min"] = torch.minimum(mins.setdefault("clip_ratio/low_min", low_min), low_min)
+            high_clip_seq = masked_seq_mean(is_high_clipped.float())
+            high_max = torch.where(high_clip_seq.isnan(), -torch.inf, high_clip_seq).max()
+            maxs["clip_ratio/high_max"] = torch.maximum(maxs.setdefault("clip_ratio/high_max", high_max), high_max)
         elif self.loss_type == "cispo":
             is_cispo_clipped = (coef_1 > self.epsilon_high) & (advantages > 0)
-            self._metrics[mode]["cispo_clip_ratio"].append(global_masked_mean(is_cispo_clipped.float()))
+            accumulate_masked_mean("cispo_clip_ratio", is_cispo_clipped.float())
         elif self.loss_type == "vespo":
-            self._metrics[mode]["vespo/phi_seq_mean"].append(global_masked_mean(phi_seq))
+            accumulate_masked_mean("vespo/phi_seq_mean", phi_seq)
 
         return loss
 
@@ -3400,6 +3409,31 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
             valid = [v for v in val if not math.isnan(v)]
             metrics[key] = sum(valid) / len(valid) if valid else None
 
+        # Metrics accumulated as per-rank running sums are aggregated across ranks here, in a single collective per
+        # logging window: each `<name>` metric is computed as `<name>_sum / <name>_count`, i.e. weighted by
+        # whatever the count counts (tokens, sequences, batches). Keys are sorted so that every rank stacks them in
+        # the same order.
+        sums = self._metric_sums[mode]
+        if sums:
+            keys = sorted(sums)
+            values = torch.stack([sums[key].double() for key in keys])
+            totals = dict(zip(keys, self.accelerator.reduce(values, reduction="sum").tolist(), strict=True))
+            for key in keys:
+                if key.endswith("_sum"):
+                    name = key.removesuffix("_sum")
+                    count = totals[name + "_count"]
+                    metrics[name] = totals[key] / count if count > 0 else 0.0
+        # Running extrema take the min/max across ranks, one collective each (min(x) is computed as -max(-x)). A
+        # rank with no data contributed +/-inf sentinels; a window with no data at all logs None, like the
+        # NaN-filtered metrics above.
+        for extrema, sign in ((self._metric_mins[mode], -1.0), (self._metric_maxs[mode], 1.0)):
+            if extrema:
+                keys = sorted(extrema)
+                values = torch.stack([sign * extrema[key].double() for key in keys])
+                reduced = sign * self.accelerator.reduce(values, reduction="max")
+                for key, value in zip(keys, reduced.tolist(), strict=True):
+                    metrics[key] = value if math.isfinite(value) else None
+
         # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs`
         # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format.
         if mode == "eval":
@@ -3408,6 +3442,9 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
         logs.update(metrics)
         super().log(logs, start_time)
         self._metrics[mode].clear()
+        self._metric_sums[mode].clear()
+        self._metric_mins[mode].clear()
+        self._metric_maxs[mode].clear()
 
         if self.accelerator.is_main_process and self.log_completions:
             if is_rich_available():
diff --git a/trl/trainer/rloo_trainer.py b/trl/trainer/rloo_trainer.py
index 3c9de79d254..6f280b0cf8f 100644
--- a/trl/trainer/rloo_trainer.py
+++ b/trl/trainer/rloo_trainer.py
@@ -67,8 +67,6 @@
     get_callable_name,
     get_config_model_id,
     identity,
-    nanmax,
-    nanmin,
     nanstd,
     pad,
     print_prompt_completions_sample,
@@ -626,6 +624,12 @@ def __init__(
 
         # Initialize the metrics
         self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)}
+        # Per-mode running sums/extrema of on-device metric accumulators. They only receive local tensors (no
+        # collective, no host sync); `log()` aggregates them across ranks in one collective per kind and resets
+        # them. `_metrics` above keeps plain floats for values that are already identical on every rank.
+        self._metric_sums = {"train": defaultdict(int), "eval": defaultdict(int)}
+        self._metric_mins = {"train": {}, "eval": {}}
+        self._metric_maxs = {"train": {}, "eval": {}}
         self._total_train_tokens = 0
         self._current_train_step_time = 0.0
         self.log_completions = args.log_completions
@@ -1617,10 +1621,8 @@ def _generate_and_score_completions(
 
         # Calculate and log the mean KL divergence between current and reference model
         if self.beta != 0.0:
-            kl_stats = self.accelerator.reduce(
-                torch.stack([(per_token_kl * completion_mask).sum(), completion_mask.sum().float()]), reduction="sum"
-            )
-            self._metrics[mode]["kl"].append((kl_stats[0] / kl_stats[1].clamp(min=1.0)).item())
+            self._metric_sums[mode]["kl_sum"] += (per_token_kl * completion_mask).sum().detach()
+            self._metric_sums[mode]["kl_count"] += completion_mask.sum().float()
 
         # Calculate mean reward per function, but only for samples where the function was applied (non-NaN values)
         for i, reward_func_name in enumerate(self.reward_func_names):
@@ -1648,14 +1650,14 @@ def _generate_and_score_completions(
             self._logs["extra"][column].extend(gather_object(self._pending_extra_logs[column]))
         self._pending_extra_logs.clear()
 
-        # Flush user-logged metrics (from log_metric), averaging across processes.
-        # Keys must be sorted so that all ranks call accelerator.gather in the same order, otherwise values
-        # get mis-attributed across metrics (dict insertion order may differ between processes).
-        for name in sorted(self._pending_metrics):
+        # Flush user-logged metrics (from log_metric), accumulated locally and averaged across processes at `log()`
+        # time. Every rank must log the same metric names, otherwise the log-time aggregation mismatches ranks
+        # (the same requirement the previous per-step gather had).
+        for name in self._pending_metrics:
             values = self._pending_metrics[name]
             local_mean = sum(values) / len(values)
-            global_mean = self.accelerator.gather(torch.tensor(local_mean, device=device)).mean().item()
-            self._metrics[mode][name].append(global_mean)
+            self._metric_sums[mode][f"{name}_sum"] += torch.tensor(local_mean, device=device)
+            self._metric_sums[mode][f"{name}_count"] += torch.ones((), device=device)
         self._pending_metrics.clear()
 
         if images is not None and self.log_multimodal:
@@ -1744,26 +1746,30 @@ def _compute_loss(self, model, inputs):
         # RLOO returns an unscaled loss (the HF Trainer divides by gradient accumulation), so add the aux term unscaled
         if self.aux_loss_enabled:
             loss = loss + self.router_aux_loss_coef * aux_loss
-            self._metrics[mode]["aux_loss"].append(self.accelerator.gather_for_metrics(aux_loss).mean().item())
+            self._metric_sums[mode]["aux_loss_sum"] += aux_loss.detach()
+            self._metric_sums[mode]["aux_loss_count"] += torch.ones_like(aux_loss)
 
         # Entropy
-        entropy_stats = self.accelerator.reduce(
-            torch.stack([(entropies * completion_mask).sum(), completion_mask.sum().float()]), reduction="sum"
-        )
-        self._metrics[mode]["entropy"].append((entropy_stats[0] / entropy_stats[1].clamp(min=1.0)).item())
+        self._metric_sums[mode]["entropy_sum"] += (entropies * completion_mask).sum().detach()
+        self._metric_sums[mode]["entropy_count"] += completion_mask.sum().float()
 
         # Compute the clipped probability ratios
+        sums = self._metric_sums[mode]
+        mins, maxs = self._metric_mins[mode], self._metric_maxs[mode]
         is_low_clipped = (coef_1 < 1 - self.epsilon_low) & (advantages < 0)
         is_high_clipped = (coef_1 > 1 + self.epsilon_high) & (advantages > 0)
         is_region_clipped = is_low_clipped | is_high_clipped
-        gathered_low_clip = self.accelerator.gather(is_low_clipped.float())
-        self._metrics[mode]["clip_ratio/low_mean"].append(gathered_low_clip.nanmean().item())
-        self._metrics[mode]["clip_ratio/low_min"].append(nanmin(gathered_low_clip).item())
-        gathered_high_clip = self.accelerator.gather(is_high_clipped.float())
-        self._metrics[mode]["clip_ratio/high_mean"].append(gathered_high_clip.nanmean().item())
-        self._metrics[mode]["clip_ratio/high_max"].append(nanmax(gathered_high_clip).item())
-        gathered_clip_ratio = self.accelerator.gather(is_region_clipped.float())
-        self._metrics[mode]["clip_ratio/region_mean"].append(gathered_clip_ratio.nanmean().item())
+        clip_count = torch.tensor(float(is_low_clipped.numel()), device=is_low_clipped.device)
+        sums["clip_ratio/low_mean_sum"] += is_low_clipped.float().sum()
+        sums["clip_ratio/low_mean_count"] += clip_count
+        low_min = is_low_clipped.float().min()
+        mins["clip_ratio/low_min"] = torch.minimum(mins.setdefault("clip_ratio/low_min", low_min), low_min)
+        sums["clip_ratio/high_mean_sum"] += is_high_clipped.float().sum()
+        sums["clip_ratio/high_mean_count"] += clip_count
+        high_max = is_high_clipped.float().max()
+        maxs["clip_ratio/high_max"] = torch.maximum(maxs.setdefault("clip_ratio/high_max", high_max), high_max)
+        sums["clip_ratio/region_mean_sum"] += is_region_clipped.float().sum()
+        sums["clip_ratio/region_mean_count"] += clip_count
         return loss
 
     # During eval, Trainer calls prediction_step. If no labels are present in the inputs, it only runs forward and
@@ -1788,6 +1794,31 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
             valid = [v for v in val if not math.isnan(v)]
             metrics[key] = sum(valid) / len(valid) if valid else None
 
+        # Metrics accumulated as per-rank running sums are aggregated across ranks here, in a single collective per
+        # logging window: each `<name>` metric is computed as `<name>_sum / <name>_count`, i.e. weighted by
+        # whatever the count counts (tokens, sequences, batches). Keys are sorted so that every rank stacks them in
+        # the same order.
+        sums = self._metric_sums[mode]
+        if sums:
+            keys = sorted(sums)
+            values = torch.stack([sums[key].double() for key in keys])
+            totals = dict(zip(keys, self.accelerator.reduce(values, reduction="sum").tolist(), strict=True))
+            for key in keys:
+                if key.endswith("_sum"):
+                    name = key.removesuffix("_sum")
+                    count = totals[name + "_count"]
+                    metrics[name] = totals[key] / count if count > 0 else 0.0
+        # Running extrema take the min/max across ranks, one collective each (min(x) is computed as -max(-x)). A
+        # rank with no data contributed +/-inf sentinels; a window with no data at all logs None, like the
+        # NaN-filtered metrics above.
+        for extrema, sign in ((self._metric_mins[mode], -1.0), (self._metric_maxs[mode], 1.0)):
+            if extrema:
+                keys = sorted(extrema)
+                values = torch.stack([sign * extrema[key].double() for key in keys])
+                reduced = sign * self.accelerator.reduce(values, reduction="max")
+                for key, value in zip(keys, reduced.tolist(), strict=True):
+                    metrics[key] = value if math.isfinite(value) else None
+
         # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs`
         # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format.
         if mode == "eval":
@@ -1796,6 +1827,9 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
         logs.update(metrics)
         super().log(logs, start_time)
         self._metrics[mode].clear()
+        self._metric_sums[mode].clear()
+        self._metric_mins[mode].clear()
+        self._metric_maxs[mode].clear()
 
         if self.accelerator.is_main_process and self.log_completions:
             if is_rich_available():

From 5a9b377aa61eb695ce1401716765a7033c98d63d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?=
 <quentin.gallouedec@huggingface.co>
Date: Mon, 24 Aug 2026 18:00:04 +0000
Subject: [PATCH 4/6] Aggregate Reward and DPO metrics at log time

---
 trl/trainer/dpo_trainer.py    | 132 ++++++++++++++++++----------------
 trl/trainer/reward_trainer.py |  82 +++++++++++++++------
 2 files changed, 132 insertions(+), 82 deletions(-)

diff --git a/trl/trainer/dpo_trainer.py b/trl/trainer/dpo_trainer.py
index a43f3764163..3021f73c69b 100644
--- a/trl/trainer/dpo_trainer.py
+++ b/trl/trainer/dpo_trainer.py
@@ -942,7 +942,9 @@ def __init__(
                 disable_dropout_in_model(self.ref_model)
 
         # Initialize the metrics
-        self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)}
+        # Per-mode running sums of on-device metric accumulators. `compute_loss` only adds local tensors here (no
+        # collective, no host sync); `log()` reduces them across ranks in a single collective and resets them.
+        self._metric_sums = {"train": defaultdict(int), "eval": defaultdict(int)}
         self._total_train_tokens = 0
 
         # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the
@@ -1328,31 +1330,33 @@ def _compute_loss_liger(self, model, inputs, return_outputs):
             rejected_rewards,
         ) = metrics
 
+        sums = self._metric_sums[mode]
         if mode == "train":
-            num_tokens_in_batch = self.accelerator.gather_for_metrics(inputs["attention_mask"].sum()).sum().item()
-            self._total_train_tokens += num_tokens_in_batch
-        self._metrics[mode]["num_tokens"] = [self._total_train_tokens]
+            sums["num_tokens_in_batch"] += inputs["attention_mask"].sum()
 
-        avg_chosen_logits = self.accelerator.gather_for_metrics(chosen_logits_mean).mean().item()
-        avg_rejected_logits = self.accelerator.gather_for_metrics(rejected_logits_mean).mean().item()
-        self._metrics[mode]["logits/chosen"].append(avg_chosen_logits)
-        self._metrics[mode]["logits/rejected"].append(avg_rejected_logits)
+        sums["logits/chosen_sum"] += chosen_logits_mean.detach()
+        sums["logits/chosen_count"] += torch.ones_like(chosen_logits_mean)
+        sums["logits/rejected_sum"] += rejected_logits_mean.detach()
+        sums["logits/rejected_count"] += torch.ones_like(rejected_logits_mean)
 
-        agg_chosen_rewards = self.accelerator.gather(chosen_rewards)
-        agg_rejected_rewards = self.accelerator.gather(rejected_rewards)
-        self._metrics[mode]["rewards/chosen"].append(agg_chosen_rewards.mean().item())
-        self._metrics[mode]["rewards/rejected"].append(agg_rejected_rewards.mean().item())
+        num_pairs = torch.tensor(float(chosen_rewards.numel()), device=chosen_rewards.device)
+        sums["rewards/chosen_sum"] += chosen_rewards.detach().sum()
+        sums["rewards/chosen_count"] += num_pairs
+        sums["rewards/rejected_sum"] += rejected_rewards.detach().sum()
+        sums["rewards/rejected_count"] += num_pairs
 
         reward_accuracies = (chosen_rewards > rejected_rewards).float()
-        agg_reward_accuracies = self.accelerator.gather(reward_accuracies)
-        self._metrics[mode]["rewards/accuracies"].append(agg_reward_accuracies.mean().item())
+        sums["rewards/accuracies_sum"] += reward_accuracies.sum()
+        sums["rewards/accuracies_count"] += num_pairs
 
         margins = chosen_rewards - rejected_rewards
-        agg_margins = self.accelerator.gather(margins)
-        self._metrics[mode]["rewards/margins"].append(agg_margins.mean().item())
+        sums["rewards/margins_sum"] += margins.detach().sum()
+        sums["rewards/margins_count"] += num_pairs
 
-        self._metrics[mode]["logps/chosen"].append(self.accelerator.gather(chosen_logps).mean().item())
-        self._metrics[mode]["logps/rejected"].append(self.accelerator.gather(rejected_logps).mean().item())
+        sums["logps/chosen_sum"] += chosen_logps.detach().sum()
+        sums["logps/chosen_count"] += num_pairs
+        sums["logps/rejected_sum"] += rejected_logps.detach().sum()
+        sums["logps/rejected_count"] += num_pairs
 
         return loss
 
@@ -1610,77 +1614,62 @@ def _compute_loss(self, model, inputs, return_outputs):
         if self.aux_loss_enabled:
             aux_loss = outputs.aux_loss
             loss = loss + self.router_aux_loss_coef * aux_loss
-            self._metrics[mode]["aux_loss"].append(self.accelerator.gather_for_metrics(aux_loss).mean().item())
+            self._metric_sums[mode]["aux_loss_sum"] += aux_loss.detach()
+            self._metric_sums[mode]["aux_loss_count"] += torch.ones_like(aux_loss)
 
         # Log the metrics
+        sums = self._metric_sums[mode]
+
         # Entropy
         per_token_entropy = entropy_from_logits(shift_logits.detach())
         mask = shift_completion_mask
-        entropy_sum = (per_token_entropy * mask).sum()
-        total_tokens = mask.sum()
-
-        # Gather counts across ranks and weight-average
-        entropy_sum = self.accelerator.gather_for_metrics(entropy_sum).sum()
-        total_tokens = self.accelerator.gather_for_metrics(total_tokens).sum()
-        entropy = (entropy_sum / total_tokens).item() if total_tokens > 0 else 0.0
-        self._metrics[mode]["entropy"].append(entropy)
+        sums["entropy_sum"] += (per_token_entropy * mask).sum()
+        sums["entropy_count"] += mask.sum()
 
         # Number of tokens
         if mode == "train":
-            num_tokens_in_batch = self.accelerator.gather_for_metrics(inputs["attention_mask"].sum()).sum().item()
-            self._total_train_tokens += num_tokens_in_batch
-        self._metrics[mode]["num_tokens"] = [self._total_train_tokens]
+            sums["num_tokens_in_batch"] += inputs["attention_mask"].sum()
 
         # Average logits for chosen and rejected completions
         chosen_logits, rejected_logits = shift_logits.detach().chunk(2, dim=0)
         chosen_mask, rejected_mask = shift_completion_mask.chunk(2, dim=0)
-        total_chosen_logits = chosen_logits[chosen_mask.bool()].mean(-1).sum()
-        total_chosen_tokens = chosen_mask.sum()
-        total_rejected_logits = rejected_logits[rejected_mask.bool()].mean(-1).sum()
-        total_rejected_tokens = rejected_mask.sum()
-        total_chosen_logits = self.accelerator.gather_for_metrics(total_chosen_logits).sum().item()
-        total_chosen_tokens = self.accelerator.gather_for_metrics(total_chosen_tokens).sum().item()
-        total_rejected_logits = self.accelerator.gather_for_metrics(total_rejected_logits).sum().item()
-        total_rejected_tokens = self.accelerator.gather_for_metrics(total_rejected_tokens).sum().item()
-        avg_chosen_logits = total_chosen_logits / total_chosen_tokens if total_chosen_tokens > 0 else 0.0
-        avg_rejected_logits = total_rejected_logits / total_rejected_tokens if total_rejected_tokens > 0 else 0.0
-        self._metrics[mode]["logits/chosen"].append(avg_chosen_logits)
-        self._metrics[mode]["logits/rejected"].append(avg_rejected_logits)
+        sums["logits/chosen_sum"] += chosen_logits[chosen_mask.bool()].mean(-1).sum()
+        sums["logits/chosen_count"] += chosen_mask.sum()
+        sums["logits/rejected_sum"] += rejected_logits[rejected_mask.bool()].mean(-1).sum()
+        sums["logits/rejected_count"] += rejected_mask.sum()
 
         # Token accuracy for the chosen completions
         predictions = chosen_logits.argmax(dim=-1)
         chosen_mask = shift_completion_mask[: len(shift_completion_mask) // 2].bool()
         chosen_labels = shift_labels[: len(shift_labels) // 2]
         correct_predictions = (predictions == chosen_labels) & chosen_mask
-        total_tokens = chosen_mask.sum()
-        correct_tokens = correct_predictions.sum()
-        correct_tokens = self.accelerator.gather_for_metrics(correct_tokens)
-        total_tokens = self.accelerator.gather_for_metrics(total_tokens)
-        total_sum = total_tokens.sum()
-        accuracy = (correct_tokens.sum() / total_sum).item() if total_sum > 0 else 0.0
-        self._metrics[mode]["mean_token_accuracy"].append(accuracy)
+        sums["mean_token_accuracy_sum"] += correct_predictions.sum()
+        sums["mean_token_accuracy_count"] += chosen_mask.sum()
 
         # Rewards for chosen and rejected completions
         chosen_rewards = self.beta * chosen_logratios.detach()
         rejected_rewards = self.beta * rejected_logratios.detach()
-        agg_chosen_rewards = self.accelerator.gather(chosen_rewards)
-        agg_rejected_rewards = self.accelerator.gather(rejected_rewards)
-        self._metrics[mode]["rewards/chosen"].append(agg_chosen_rewards.mean().item())
-        self._metrics[mode]["rewards/rejected"].append(agg_rejected_rewards.mean().item())
+        num_pairs = torch.tensor(float(chosen_rewards.numel()), device=chosen_rewards.device)
+        sums["rewards/chosen_sum"] += chosen_rewards.sum()
+        sums["rewards/chosen_count"] += num_pairs
+        sums["rewards/rejected_sum"] += rejected_rewards.sum()
+        sums["rewards/rejected_count"] += num_pairs
 
         # Reward accuracy
         reward_accuracies = (chosen_rewards > rejected_rewards).float()
-        agg_reward_accuracies = self.accelerator.gather(reward_accuracies)
-        self._metrics[mode]["rewards/accuracies"].append(agg_reward_accuracies.mean().item())
+        sums["rewards/accuracies_sum"] += reward_accuracies.sum()
+        sums["rewards/accuracies_count"] += num_pairs
 
         # Reward margins
         margins = chosen_rewards - rejected_rewards
-        agg_margins = self.accelerator.gather(margins)
-        self._metrics[mode]["rewards/margins"].append(agg_margins.mean().item())
+        sums["rewards/margins_sum"] += margins.sum()
+        sums["rewards/margins_count"] += num_pairs
 
         # Average log probabilities for chosen and rejected completions
-        self._metrics[mode]["logps/chosen"].append(self.accelerator.gather(chosen_logps).mean().item())
-        self._metrics[mode]["logps/rejected"].append(self.accelerator.gather(rejected_logps).mean().item())
+        sums["logps/chosen_sum"] += chosen_logps.detach().sum()
+        sums["logps/chosen_count"] += num_pairs
+        sums["logps/rejected_sum"] += rejected_logps.detach().sum()
+        sums["logps/rejected_count"] += num_pairs
 
         return (loss, outputs) if return_outputs else loss
 
@@ -1783,14 +1772,35 @@ def training_step(self, *args, **kwargs):
 
     def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
         mode = "train" if self.model.training else "eval"
-        metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()}  # average the metrics
+
+        # Metrics are accumulated in `compute_loss` as per-rank running sums. Aggregate them across ranks here, in
+        # a single collective per logging window, then compute each `<name>` metric as `<name>_sum / <name>_count`,
+        # i.e. weighted by whatever the count counts (tokens, pairs, batches). Keys are sorted so that every rank
+        # stacks them in the same order.
+        metrics = {}
+        sums = self._metric_sums[mode]
+        if sums:
+            keys = sorted(sums)
+            values = torch.stack([sums[key].double() for key in keys])
+            totals = dict(zip(keys, self.accelerator.reduce(values, reduction="sum").tolist(), strict=True))
+            for key in keys:
+                if key.endswith("_sum"):
+                    name = key.removesuffix("_sum")
+                    count = totals[name + "_count"]
+                    metrics[name] = totals[key] / count if count > 0 else 0.0
+            # `num_tokens` advances only when a train-mode log folds in the pending sums, so an eval log between two
+            # train logs can lag by up to one logging window.
+            if mode == "train" and "num_tokens_in_batch" in totals:
+                self._total_train_tokens += int(totals["num_tokens_in_batch"])
+            metrics["num_tokens"] = self._total_train_tokens
+
         # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs`
         # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format.
         if mode == "eval":
             metrics = {f"eval_{key}": val for key, val in metrics.items()}
         logs.update(metrics)
         super().log(logs, start_time)
-        self._metrics[mode].clear()
+        self._metric_sums[mode].clear()
 
     # During eval, Trainer calls prediction_step. If no labels are present in the inputs, it only runs forward and
     # returns logits. We override prediction_step to force compute_loss, because this trainer doesn't involve labels.
diff --git a/trl/trainer/reward_trainer.py b/trl/trainer/reward_trainer.py
index c1ac2977f3d..46a6b983ba1 100644
--- a/trl/trainer/reward_trainer.py
+++ b/trl/trainer/reward_trainer.py
@@ -15,6 +15,7 @@
 import contextlib
 import json
 import logging
+import math
 import os
 import re
 import warnings
@@ -601,8 +602,12 @@ def __init__(
         else:
             self.maybe_activation_offload_context = contextlib.nullcontext()
 
-        # Initialize the metrics
-        self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)}
+        # Per-mode running sums/extrema of on-device metric accumulators. `compute_loss` only adds local tensors
+        # here (no collective, no host sync); `log()` aggregates them across ranks in one collective per kind and
+        # resets them.
+        self._metric_sums = {"train": defaultdict(int), "eval": defaultdict(int)}
+        self._metric_mins = {"train": {}, "eval": {}}
+        self._metric_maxs = {"train": {}, "eval": {}}
         self._total_train_tokens = 0
 
         # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the
@@ -699,8 +704,10 @@ def tokenize_fn(example, processing_class):
                 if isinstance(dataset, Dataset):  # `IterableDataset.map` does not support `desc`
                     map_kwargs["desc"] = f"Filtering {dataset_name} >{args.max_length} tokens"
                 dataset = dataset.filter(
-                    lambda example: len(example["chosen_ids"]) <= args.max_length
-                    and len(example["rejected_ids"]) <= args.max_length,
+                    lambda example: (
+                        len(example["chosen_ids"]) <= args.max_length
+                        and len(example["rejected_ids"]) <= args.max_length
+                    ),
                     **map_kwargs,
                 )
 
@@ -760,24 +767,26 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
             loss += self.args.center_rewards_coefficient * torch.mean((rewards_chosen + rewards_rejected) ** 2)
 
         if mode == "train":
-            num_tokens_in_batch = self.accelerator.gather_for_metrics(inputs["attention_mask"].sum()).sum().item()
-            self._total_train_tokens += num_tokens_in_batch
-        self._metrics[mode]["num_tokens"] = [self._total_train_tokens]
+            self._metric_sums[mode]["num_tokens_in_batch"] += inputs["attention_mask"].sum()
 
         # Compute min, mean, max, accuracy and margin
         with torch.no_grad():
-            all_rewards = self.accelerator.gather(outputs.logits)
-            self._metrics[mode]["min_reward"].append(all_rewards.min().item())
-            self._metrics[mode]["mean_reward"].append(all_rewards.mean().item())
-            self._metrics[mode]["max_reward"].append(all_rewards.max().item())
-
-            mean_accuracy = (rewards_chosen > rewards_rejected).float().mean()
-            mean_accuracy = self.accelerator.gather_for_metrics(mean_accuracy).mean().item()
-            self._metrics[mode]["accuracy"].append(mean_accuracy)
-
-            mean_margin = (rewards_chosen - rewards_rejected).mean()
-            mean_margin = self.accelerator.gather_for_metrics(mean_margin).mean()
-            self._metrics[mode]["margin"].append(mean_margin.item())
+            sums = self._metric_sums[mode]
+            mins, maxs = self._metric_mins[mode], self._metric_maxs[mode]
+            rewards = outputs.logits.detach()
+            num_pairs = torch.tensor(float(rewards_chosen.numel()), device=rewards.device)
+            sums["mean_reward_sum"] += rewards.sum()
+            sums["mean_reward_count"] += torch.tensor(float(rewards.numel()), device=rewards.device)
+            min_reward = rewards.min()
+            mins["min_reward"] = torch.minimum(mins.setdefault("min_reward", min_reward), min_reward)
+            max_reward = rewards.max()
+            maxs["max_reward"] = torch.maximum(maxs.setdefault("max_reward", max_reward), max_reward)
+
+            sums["accuracy_sum"] += (rewards_chosen > rewards_rejected).float().sum()
+            sums["accuracy_count"] += num_pairs
+
+            sums["margin_sum"] += (rewards_chosen - rewards_rejected).sum()
+            sums["margin_count"] += num_pairs
 
         return (loss, outputs) if return_outputs else loss
 
@@ -788,7 +797,36 @@ def training_step(self, *args, **kwargs):
 
     def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
         mode = "train" if self.model.training else "eval"
-        metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()}  # average the metrics
+
+        # Metrics are accumulated in `compute_loss` as per-rank running sums. Aggregate them across ranks here, in
+        # a single collective per logging window, then compute each `<name>` metric as `<name>_sum / <name>_count`,
+        # i.e. weighted by whatever the count counts (tokens, batches). Keys are sorted so that every rank stacks
+        # them in the same order.
+        metrics = {}
+        sums = self._metric_sums[mode]
+        if sums:
+            keys = sorted(sums)
+            values = torch.stack([sums[key].double() for key in keys])
+            totals = dict(zip(keys, self.accelerator.reduce(values, reduction="sum").tolist(), strict=True))
+            for key in keys:
+                if key.endswith("_sum"):
+                    name = key.removesuffix("_sum")
+                    count = totals[name + "_count"]
+                    metrics[name] = totals[key] / count if count > 0 else 0.0
+            # `num_tokens` advances only when a train-mode log folds in the pending sums, so an eval log between two
+            # train logs can lag by up to one logging window.
+            if mode == "train" and "num_tokens_in_batch" in totals:
+                self._total_train_tokens += int(totals["num_tokens_in_batch"])
+            metrics["num_tokens"] = self._total_train_tokens
+        # Running extrema take the min/max across ranks, one collective each (min(x) is computed as -max(-x)). A
+        # rank with no data contributed +/-inf sentinels; a window with no data at all logs None.
+        for extrema, sign in ((self._metric_mins[mode], -1.0), (self._metric_maxs[mode], 1.0)):
+            if extrema:
+                keys = sorted(extrema)
+                values = torch.stack([sign * extrema[key].double() for key in keys])
+                reduced = sign * self.accelerator.reduce(values, reduction="max")
+                for key, value in zip(keys, reduced.tolist(), strict=True):
+                    metrics[key] = value if math.isfinite(value) else None
 
         # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs`
         # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format.
@@ -797,7 +835,9 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
 
         logs.update(metrics)
         super().log(logs, start_time)
-        self._metrics[mode].clear()
+        self._metric_sums[mode].clear()
+        self._metric_mins[mode].clear()
+        self._metric_maxs[mode].clear()
 
     # Ensure the model card is saved along with the checkpoint
     def _save_checkpoint(self, model, trial):

From 14ba632f0597e5131df3bceec0fe33006db6a6d1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?=
 <quentin.gallouedec@huggingface.co>
Date: Mon, 24 Aug 2026 18:06:03 +0000
Subject: [PATCH 5/6] Aggregate KTO metrics at log time

---
 trl/trainer/kto_trainer.py | 156 +++++++++++++++++--------------------
 1 file changed, 70 insertions(+), 86 deletions(-)

diff --git a/trl/trainer/kto_trainer.py b/trl/trainer/kto_trainer.py
index 632cea2423d..9adbc3881c5 100644
--- a/trl/trainer/kto_trainer.py
+++ b/trl/trainer/kto_trainer.py
@@ -945,7 +945,9 @@ def __init__(
                 disable_dropout_in_model(self.ref_model)
 
         # Initialize the metrics
-        self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)}
+        # Per-mode running sums of on-device metric accumulators. `compute_loss` only adds local tensors here (no
+        # collective, no host sync); `log()` reduces them across ranks in a single collective and resets them.
+        self._metric_sums = {"train": defaultdict(int), "eval": defaultdict(int)}
         self._total_train_tokens = 0
 
         # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the
@@ -1449,43 +1451,28 @@ def _compute_loss_liger(self, model, inputs, return_outputs):
                 kl=kl,
             )
 
-        self._metrics[mode]["kl"].append(kl.item())
+        sums = self._metric_sums[mode]
+        # `kl` is already global here (the gather above is required by the loss), so only the window average is left
+        # for `log()` to compute.
+        sums["kl_sum"] += kl.detach().sum()
+        sums["kl_count"] += torch.ones((), device=kl.device)
 
         # Number of tokens
         if mode == "train":
-            num_tokens_in_batch = self.accelerator.gather_for_metrics(batch["attention_mask"].sum()).sum().item()
-            self._total_train_tokens += num_tokens_in_batch
-        self._metrics[mode]["num_tokens"] = [self._total_train_tokens]
-
-        all_num_chosen = self.accelerator.gather_for_metrics(num_chosen).sum().item()
-        all_num_rejected = self.accelerator.gather_for_metrics(num_rejected).sum().item()
-
-        if all_num_chosen > 0:
-            self._metrics[mode]["rewards/chosen"].append(
-                self.accelerator.gather_for_metrics(chosen_rewards_sum.nansum()).nansum().item() / all_num_chosen
-            )
-            self._metrics[mode]["logps/chosen"].append(
-                self.accelerator.gather_for_metrics(chosen_logps_sum.nansum()).nansum().item() / all_num_chosen
-            )
-            self._metrics[mode]["logits/chosen"].append(
-                self.accelerator.gather_for_metrics(chosen_logits_sum.nansum()).nansum().item() / all_num_chosen
-            )
-
-        if all_num_rejected > 0:
-            self._metrics[mode]["rewards/rejected"].append(
-                self.accelerator.gather_for_metrics(rejected_rewards_sum.nansum()).nansum().item() / all_num_rejected
-            )
-            self._metrics[mode]["logps/rejected"].append(
-                self.accelerator.gather_for_metrics(rejected_logps_sum.nansum()).nansum().item() / all_num_rejected
-            )
-            self._metrics[mode]["logits/rejected"].append(
-                self.accelerator.gather_for_metrics(rejected_logits_sum.nansum()).nansum().item() / all_num_rejected
-            )
-
-        if all_num_chosen > 0 and all_num_rejected > 0:
-            self._metrics[mode]["rewards/margins"].append(
-                self._metrics[mode]["rewards/chosen"][-1] - self._metrics[mode]["rewards/rejected"][-1]
-            )
+            sums["num_tokens_in_batch"] += batch["attention_mask"].sum()
+
+        sums["rewards/chosen_sum"] += chosen_rewards_sum.nansum().detach()
+        sums["rewards/chosen_count"] += num_chosen
+        sums["logps/chosen_sum"] += chosen_logps_sum.nansum().detach()
+        sums["logps/chosen_count"] += num_chosen
+        sums["logits/chosen_sum"] += chosen_logits_sum.nansum().detach()
+        sums["logits/chosen_count"] += num_chosen
+        sums["rewards/rejected_sum"] += rejected_rewards_sum.nansum().detach()
+        sums["rewards/rejected_count"] += num_rejected
+        sums["logps/rejected_sum"] += rejected_logps_sum.nansum().detach()
+        sums["logps/rejected_count"] += num_rejected
+        sums["logits/rejected_sum"] += rejected_logits_sum.nansum().detach()
+        sums["logits/rejected_count"] += num_rejected
 
         return loss
 
@@ -1605,25 +1592,21 @@ def _compute_loss(self, model, inputs, return_outputs):
             0,
         )
 
-        self._metrics[mode]["kl"].append(kl.item())
+        sums = self._metric_sums[mode]
+        # `kl` is already global here (the gather above is required by the loss), so only the window average is left
+        # for `log()` to compute.
+        sums["kl_sum"] += kl.detach().sum()
+        sums["kl_count"] += torch.ones((), device=kl.device)
 
         # Entropy
         per_token_entropy = entropy_from_logits(shift_logits.detach())
         mask = batch["completion_mask"][:, 1:]
-        entropy_sum = (per_token_entropy * mask).sum()
-        total_tokens = mask.sum()
-
-        # Gather counts across ranks and weight-average
-        entropy_sum = self.accelerator.gather_for_metrics(entropy_sum).sum()
-        total_tokens = self.accelerator.gather_for_metrics(total_tokens).sum()
-        entropy = (entropy_sum / total_tokens).item() if total_tokens > 0 else 0.0
-        self._metrics[mode]["entropy"].append(entropy)
+        sums["entropy_sum"] += (per_token_entropy * mask).sum()
+        sums["entropy_count"] += mask.sum()
 
         # Number of tokens
         if mode == "train":
-            num_tokens_in_batch = self.accelerator.gather_for_metrics(batch["attention_mask"].sum()).sum().item()
-            self._total_train_tokens += num_tokens_in_batch
-        self._metrics[mode]["num_tokens"] = [self._total_train_tokens]
+            sums["num_tokens_in_batch"] += batch["attention_mask"].sum()
 
         # Average logits for chosen and rejected completions
         shift_completion_mask = batch["completion_mask"][:, 1:]
@@ -1631,48 +1614,26 @@ def _compute_loss(self, model, inputs, return_outputs):
         rejected_logits = shift_logits.detach().index_select(0, rejected_idx)
         chosen_mask = shift_completion_mask.index_select(0, chosen_idx)
         rejected_mask = shift_completion_mask.index_select(0, rejected_idx)
-        total_chosen_logits = chosen_logits[chosen_mask.bool()].mean(-1).sum()
-        total_chosen_tokens = chosen_mask.sum()
-        total_rejected_logits = rejected_logits[rejected_mask.bool()].mean(-1).sum()
-        total_rejected_tokens = rejected_mask.sum()
-        total_chosen_logits = self.accelerator.gather_for_metrics(total_chosen_logits).sum().item()
-        total_chosen_tokens = self.accelerator.gather_for_metrics(total_chosen_tokens).sum().item()
-        total_rejected_logits = self.accelerator.gather_for_metrics(total_rejected_logits).sum().item()
-        total_rejected_tokens = self.accelerator.gather_for_metrics(total_rejected_tokens).sum().item()
-        if total_chosen_tokens > 0:
-            self._metrics[mode]["logits/chosen"].append(total_chosen_logits / total_chosen_tokens)
-        if total_rejected_tokens > 0:
-            self._metrics[mode]["logits/rejected"].append(total_rejected_logits / total_rejected_tokens)
-
-        all_num_chosen = self.accelerator.gather_for_metrics(num_chosen).sum().item()
-        all_num_rejected = self.accelerator.gather_for_metrics(num_rejected).sum().item()
-
-        if all_num_chosen > 0:
-            self._metrics[mode]["rewards/chosen"].append(
-                self.accelerator.gather_for_metrics(chosen_rewards.nansum()).nansum().item() / all_num_chosen
-            )
-            self._metrics[mode]["logps/chosen"].append(
-                self.accelerator.gather_for_metrics(chosen_logps.nansum()).nansum().item() / all_num_chosen
-            )
-
-        if all_num_rejected > 0:
-            self._metrics[mode]["rewards/rejected"].append(
-                self.accelerator.gather_for_metrics(rejected_rewards.nansum()).nansum().item() / all_num_rejected
-            )
-            self._metrics[mode]["logps/rejected"].append(
-                self.accelerator.gather_for_metrics(rejected_logps.nansum()).nansum().item() / all_num_rejected
-            )
-
-        if all_num_chosen > 0 and all_num_rejected > 0:
-            self._metrics[mode]["rewards/margins"].append(
-                self._metrics[mode]["rewards/chosen"][-1] - self._metrics[mode]["rewards/rejected"][-1]
-            )
+        sums["logits/chosen_sum"] += chosen_logits[chosen_mask.bool()].mean(-1).sum()
+        sums["logits/chosen_count"] += chosen_mask.sum()
+        sums["logits/rejected_sum"] += rejected_logits[rejected_mask.bool()].mean(-1).sum()
+        sums["logits/rejected_count"] += rejected_mask.sum()
+
+        sums["rewards/chosen_sum"] += chosen_rewards.nansum()
+        sums["rewards/chosen_count"] += num_chosen
+        sums["logps/chosen_sum"] += chosen_logps.detach().nansum()
+        sums["logps/chosen_count"] += num_chosen
+        sums["rewards/rejected_sum"] += rejected_rewards.nansum()
+        sums["rewards/rejected_count"] += num_rejected
+        sums["logps/rejected_sum"] += rejected_logps.detach().nansum()
+        sums["logps/rejected_count"] += num_rejected
 
         loss = losses.nanmean()
         if self.aux_loss_enabled:
             aux_loss = outputs.aux_loss
             loss = loss + self.router_aux_loss_coef * aux_loss
-            self._metrics[mode]["aux_loss"].append(self.accelerator.gather_for_metrics(aux_loss).mean().item())
+            sums["aux_loss_sum"] += aux_loss.detach()
+            sums["aux_loss_count"] += torch.ones_like(aux_loss)
 
         return (loss, outputs) if return_outputs else loss
 
@@ -1775,14 +1736,37 @@ def training_step(self, *args, **kwargs):
 
     def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
         mode = "train" if self.model.training else "eval"
-        metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()}  # average the metrics
+
+        # Metrics are accumulated in `compute_loss` as per-rank running sums. Aggregate them across ranks here, in
+        # a single collective per logging window, then compute each `<name>` metric as `<name>_sum / <name>_count`,
+        # i.e. weighted by whatever the count counts (tokens, examples, batches). Keys are sorted so that every rank
+        # stacks them in the same order.
+        metrics = {}
+        sums = self._metric_sums[mode]
+        if sums:
+            keys = sorted(sums)
+            values = torch.stack([sums[key].double() for key in keys])
+            totals = dict(zip(keys, self.accelerator.reduce(values, reduction="sum").tolist(), strict=True))
+            for key in keys:
+                if key.endswith("_sum"):
+                    name = key.removesuffix("_sum")
+                    count = totals[name + "_count"]
+                    metrics[name] = totals[key] / count if count > 0 else 0.0
+            if totals.get("rewards/chosen_count", 0) > 0 and totals.get("rewards/rejected_count", 0) > 0:
+                metrics["rewards/margins"] = metrics["rewards/chosen"] - metrics["rewards/rejected"]
+            # `num_tokens` advances only when a train-mode log folds in the pending sums, so an eval log between two
+            # train logs can lag by up to one logging window.
+            if mode == "train" and "num_tokens_in_batch" in totals:
+                self._total_train_tokens += int(totals["num_tokens_in_batch"])
+            metrics["num_tokens"] = self._total_train_tokens
+
         # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs`
         # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format.
         if mode == "eval":
             metrics = {f"eval_{key}": val for key, val in metrics.items()}
         logs.update(metrics)
         super().log(logs, start_time)
-        self._metrics[mode].clear()
+        self._metric_sums[mode].clear()
 
     # During eval, Trainer calls prediction_step. If no labels are present in the inputs, it only runs forward and
     # returns logits. We override prediction_step to force compute_loss, because this trainer doesn't involve labels.

From 985367694df04b1074d891d635db77de22169d76 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?=
 <quentin.gallouedec@huggingface.co>
Date: Thu, 27 Aug 2026 02:23:19 +0000
Subject: [PATCH 6/6] Store metrics as (total, count) pairs, matching the
 existing reduce idiom

---
 tests/test_grpo_trainer.py          |   8 +-
 tests/test_sft_trainer.py           |  16 ++++
 trl/trainer/distillation_trainer.py |  34 ++++-----
 trl/trainer/dpo_trainer.py          | 108 ++++++++++++---------------
 trl/trainer/grpo_trainer.py         |  84 ++++++++++-----------
 trl/trainer/kto_trainer.py          | 109 +++++++++++++---------------
 trl/trainer/reward_trainer.py       |  60 +++++++--------
 trl/trainer/rloo_trainer.py         |  69 ++++++++----------
 trl/trainer/sft_trainer.py          |  75 ++++++++++---------
 9 files changed, 266 insertions(+), 297 deletions(-)

diff --git a/tests/test_grpo_trainer.py b/tests/test_grpo_trainer.py
index 20f30e0d91f..a85742cb542 100644
--- a/tests/test_grpo_trainer.py
+++ b/tests/test_grpo_trainer.py
@@ -1637,12 +1637,8 @@ def generate_with_one_unscorable_token(prompts):
 
         def record_metrics(inputs):
             outputs = original_score(inputs)
-            sums = trainer._metric_sums["train"]
-            mean = (
-                sums["sampling/sampling_logp_difference/mean_sum"]
-                / sums["sampling/sampling_logp_difference/mean_count"]
-            )
-            recorded_metrics.append(("sampling/sampling_logp_difference/mean", mean.item()))
+            total, count = trainer._metric_stats["train"]["sampling/sampling_logp_difference/mean"]
+            recorded_metrics.append(("sampling/sampling_logp_difference/mean", (total / count).item()))
             max_delta = trainer._metric_maxs["train"]["sampling/sampling_logp_difference/max"]
             recorded_metrics.append(("sampling/sampling_logp_difference/max", max_delta.item()))
             return outputs
diff --git a/tests/test_sft_trainer.py b/tests/test_sft_trainer.py
index a0354d27dfa..4ad0acbf35f 100644
--- a/tests/test_sft_trainer.py
+++ b/tests/test_sft_trainer.py
@@ -1607,6 +1607,22 @@ def convert_to_json(example):
             new_param = trainer.model.get_parameter(n)
             assert not torch.equal(param, new_param), f"Parameter {n} has not changed."
 
+    def test_log_averages_over_the_window_weighted_by_count(self):
+        dataset = load_dataset("trl-internal-testing/zen", "standard_language_modeling", split="train")
+        training_args = SFTConfig(output_dir=self.tmp_dir, report_to="none")
+        trainer = SFTTrainer(
+            model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5", args=training_args, train_dataset=dataset
+        )
+
+        # Two steps in one logging window: 1 token with entropy 3.0, then 9 tokens with entropy 1.0. The logged
+        # value must weight them by their token counts (12 / 10), not average the per-step ratios ((3 + 1) / 2).
+        trainer.model.train()
+        trainer._metric_stats["train"]["entropy"] = torch.tensor([3.0 + 9.0, 1.0 + 9.0], device=torch_device)
+        logs = {}
+        trainer.log(logs)
+
+        assert logs["entropy"] == pytest.approx(1.2)
+
     def test_train_with_eval(self):
         dataset = load_dataset("trl-internal-testing/zen", "standard_language_modeling")
 
diff --git a/trl/trainer/distillation_trainer.py b/trl/trainer/distillation_trainer.py
index 680695ea4d6..391bbd3e2eb 100644
--- a/trl/trainer/distillation_trainer.py
+++ b/trl/trainer/distillation_trainer.py
@@ -826,9 +826,10 @@ def __init__(
 
         # Metrics & Logging
         self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)}
-        # Per-mode running sums of on-device metric accumulators. `compute_loss` only adds local tensors here (no
-        # collective, no host sync); `log()` reduces them across ranks in a single collective and resets them.
-        self._metric_sums = {"train": defaultdict(int), "eval": defaultdict(int)}
+        # Each entry is a running `(total, count)` pair of on-device tensors, summed in `compute_loss` with no
+        # collective and no host sync. `log()` reduces the pairs across ranks in a single collective and divides,
+        # so every metric is weighted by whatever its count counts (tokens, batches).
+        self._metric_stats = {"train": defaultdict(int), "eval": defaultdict(int)}
         self._total_train_tokens = 0
         self._current_train_step_time = 0.0
         self.log_completions = args.log_completions
@@ -1860,8 +1861,9 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
         # forward. The Liger path produces no entropy, so it logs none. Mirrors `SFTTrainer.compute_loss`.
         if entropy_sum is not None:
             mode = "train" if self.model.training else "eval"
-            self._metric_sums[mode]["entropy_sum"] += entropy_sum.detach()
-            self._metric_sums[mode]["entropy_count"] += num_valid_tokens.detach()
+            self._metric_stats[mode]["entropy"] += torch.stack(
+                [entropy_sum.detach(), num_valid_tokens.detach().float()]
+            )
 
         return (loss, None) if return_outputs else loss
 
@@ -2020,20 +2022,14 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
             valid = [v for v in val if not math.isnan(v)]
             metrics[key] = sum(valid) / len(valid) if valid else None
 
-        # Metrics are accumulated in `compute_loss` as per-rank running sums. Aggregate them across ranks here, in
-        # a single collective per logging window, then compute each `<name>` metric as `<name>_sum / <name>_count`,
-        # i.e. weighted by whatever the count counts (tokens, batches). Keys are sorted so that every rank stacks
-        # them in the same order.
-        sums = self._metric_sums[mode]
-        if sums:
-            keys = sorted(sums)
-            values = torch.stack([sums[key].double() for key in keys])
+        # Sum every `(total, count)` pair across ranks in a single collective, then divide. Keys are sorted so that
+        # all ranks stack them in the same order.
+        stats = self._metric_stats[mode]
+        if stats:
+            keys = sorted(stats)
+            values = torch.stack([stats[key].double() for key in keys])
             totals = dict(zip(keys, self.accelerator.reduce(values, reduction="sum").tolist(), strict=True))
-            for key in keys:
-                if key.endswith("_sum"):
-                    name = key.removesuffix("_sum")
-                    count = totals[name + "_count"]
-                    metrics[name] = totals[key] / count if count > 0 else 0.0
+            metrics.update({key: total / count if count > 0 else 0.0 for key, (total, count) in totals.items()})
 
         # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs`
         # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format.
@@ -2043,7 +2039,7 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
         logs.update(metrics)
         super().log(logs, start_time)
         self._metrics[mode].clear()
-        self._metric_sums[mode].clear()
+        self._metric_stats[mode].clear()
 
         if self.accelerator.is_main_process and self.log_completions:
             if is_rich_available():
diff --git a/trl/trainer/dpo_trainer.py b/trl/trainer/dpo_trainer.py
index 3021f73c69b..73e104bff6f 100644
--- a/trl/trainer/dpo_trainer.py
+++ b/trl/trainer/dpo_trainer.py
@@ -942,9 +942,10 @@ def __init__(
                 disable_dropout_in_model(self.ref_model)
 
         # Initialize the metrics
-        # Per-mode running sums of on-device metric accumulators. `compute_loss` only adds local tensors here (no
-        # collective, no host sync); `log()` reduces them across ranks in a single collective and resets them.
-        self._metric_sums = {"train": defaultdict(int), "eval": defaultdict(int)}
+        # Each entry is a running `(total, count)` pair of on-device tensors, summed in `compute_loss` with no
+        # collective and no host sync. `log()` reduces the pairs across ranks in a single collective and divides,
+        # so every metric is weighted by whatever its count counts (tokens, pairs, batches).
+        self._metric_stats = {"train": defaultdict(int), "eval": defaultdict(int)}
         self._total_train_tokens = 0
 
         # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the
@@ -1330,33 +1331,26 @@ def _compute_loss_liger(self, model, inputs, return_outputs):
             rejected_rewards,
         ) = metrics
 
-        sums = self._metric_sums[mode]
+        stats = self._metric_stats[mode]
         if mode == "train":
-            sums["num_tokens_in_batch"] += inputs["attention_mask"].sum()
+            num_tokens_in_batch = inputs["attention_mask"].sum()
+            stats["num_tokens"] += torch.stack([num_tokens_in_batch, torch.ones_like(num_tokens_in_batch)])
 
-        sums["logits/chosen_sum"] += chosen_logits_mean.detach()
-        sums["logits/chosen_count"] += torch.ones_like(chosen_logits_mean)
-        sums["logits/rejected_sum"] += rejected_logits_mean.detach()
-        sums["logits/rejected_count"] += torch.ones_like(rejected_logits_mean)
+        stats["logits/chosen"] += torch.stack([chosen_logits_mean.detach(), torch.ones_like(chosen_logits_mean)])
+        stats["logits/rejected"] += torch.stack([rejected_logits_mean.detach(), torch.ones_like(rejected_logits_mean)])
 
         num_pairs = torch.tensor(float(chosen_rewards.numel()), device=chosen_rewards.device)
-        sums["rewards/chosen_sum"] += chosen_rewards.detach().sum()
-        sums["rewards/chosen_count"] += num_pairs
-        sums["rewards/rejected_sum"] += rejected_rewards.detach().sum()
-        sums["rewards/rejected_count"] += num_pairs
+        stats["rewards/chosen"] += torch.stack([chosen_rewards.detach().sum(), num_pairs])
+        stats["rewards/rejected"] += torch.stack([rejected_rewards.detach().sum(), num_pairs])
 
         reward_accuracies = (chosen_rewards > rejected_rewards).float()
-        sums["rewards/accuracies_sum"] += reward_accuracies.sum()
-        sums["rewards/accuracies_count"] += num_pairs
+        stats["rewards/accuracies"] += torch.stack([reward_accuracies.sum(), num_pairs])
 
         margins = chosen_rewards - rejected_rewards
-        sums["rewards/margins_sum"] += margins.detach().sum()
-        sums["rewards/margins_count"] += num_pairs
+        stats["rewards/margins"] += torch.stack([margins.detach().sum(), num_pairs])
 
-        sums["logps/chosen_sum"] += chosen_logps.detach().sum()
-        sums["logps/chosen_count"] += num_pairs
-        sums["logps/rejected_sum"] += rejected_logps.detach().sum()
-        sums["logps/rejected_count"] += num_pairs
+        stats["logps/chosen"] += torch.stack([chosen_logps.detach().sum(), num_pairs])
+        stats["logps/rejected"] += torch.stack([rejected_logps.detach().sum(), num_pairs])
 
         return loss
 
@@ -1614,62 +1608,56 @@ def _compute_loss(self, model, inputs, return_outputs):
         if self.aux_loss_enabled:
             aux_loss = outputs.aux_loss
             loss = loss + self.router_aux_loss_coef * aux_loss
-            self._metric_sums[mode]["aux_loss_sum"] += aux_loss.detach()
-            self._metric_sums[mode]["aux_loss_count"] += torch.ones_like(aux_loss)
+            self._metric_stats[mode]["aux_loss"] += torch.stack([aux_loss.detach(), torch.ones_like(aux_loss)])
 
         # Log the metrics
-        sums = self._metric_sums[mode]
+        stats = self._metric_stats[mode]
 
         # Entropy
         per_token_entropy = entropy_from_logits(shift_logits.detach())
         mask = shift_completion_mask
-        sums["entropy_sum"] += (per_token_entropy * mask).sum()
-        sums["entropy_count"] += mask.sum()
+        stats["entropy"] += torch.stack([(per_token_entropy * mask).sum(), mask.sum().float()])
 
         # Number of tokens
         if mode == "train":
-            sums["num_tokens_in_batch"] += inputs["attention_mask"].sum()
+            num_tokens_in_batch = inputs["attention_mask"].sum()
+            stats["num_tokens"] += torch.stack([num_tokens_in_batch, torch.ones_like(num_tokens_in_batch)])
 
         # Average logits for chosen and rejected completions
         chosen_logits, rejected_logits = shift_logits.detach().chunk(2, dim=0)
         chosen_mask, rejected_mask = shift_completion_mask.chunk(2, dim=0)
-        sums["logits/chosen_sum"] += chosen_logits[chosen_mask.bool()].mean(-1).sum()
-        sums["logits/chosen_count"] += chosen_mask.sum()
-        sums["logits/rejected_sum"] += rejected_logits[rejected_mask.bool()].mean(-1).sum()
-        sums["logits/rejected_count"] += rejected_mask.sum()
+        stats["logits/chosen"] += torch.stack(
+            [chosen_logits[chosen_mask.bool()].mean(-1).sum(), chosen_mask.sum().float()]
+        )
+        stats["logits/rejected"] += torch.stack(
+            [rejected_logits[rejected_mask.bool()].mean(-1).sum(), rejected_mask.sum().float()]
+        )
 
         # Token accuracy for the chosen completions
         predictions = chosen_logits.argmax(dim=-1)
         chosen_mask = shift_completion_mask[: len(shift_completion_mask) // 2].bool()
         chosen_labels = shift_labels[: len(shift_labels) // 2]
         correct_predictions = (predictions == chosen_labels) & chosen_mask
-        sums["mean_token_accuracy_sum"] += correct_predictions.sum()
-        sums["mean_token_accuracy_count"] += chosen_mask.sum()
+        stats["mean_token_accuracy"] += torch.stack([correct_predictions.sum(), chosen_mask.sum()])
 
         # Rewards for chosen and rejected completions
         chosen_rewards = self.beta * chosen_logratios.detach()
         rejected_rewards = self.beta * rejected_logratios.detach()
         num_pairs = torch.tensor(float(chosen_rewards.numel()), device=chosen_rewards.device)
-        sums["rewards/chosen_sum"] += chosen_rewards.sum()
-        sums["rewards/chosen_count"] += num_pairs
-        sums["rewards/rejected_sum"] += rejected_rewards.sum()
-        sums["rewards/rejected_count"] += num_pairs
+        stats["rewards/chosen"] += torch.stack([chosen_rewards.sum(), num_pairs])
+        stats["rewards/rejected"] += torch.stack([rejected_rewards.sum(), num_pairs])
 
         # Reward accuracy
         reward_accuracies = (chosen_rewards > rejected_rewards).float()
-        sums["rewards/accuracies_sum"] += reward_accuracies.sum()
-        sums["rewards/accuracies_count"] += num_pairs
+        stats["rewards/accuracies"] += torch.stack([reward_accuracies.sum(), num_pairs])
 
         # Reward margins
         margins = chosen_rewards - rejected_rewards
-        sums["rewards/margins_sum"] += margins.sum()
-        sums["rewards/margins_count"] += num_pairs
+        stats["rewards/margins"] += torch.stack([margins.sum(), num_pairs])
 
         # Average log probabilities for chosen and rejected completions
-        sums["logps/chosen_sum"] += chosen_logps.detach().sum()
-        sums["logps/chosen_count"] += num_pairs
-        sums["logps/rejected_sum"] += rejected_logps.detach().sum()
-        sums["logps/rejected_count"] += num_pairs
+        stats["logps/chosen"] += torch.stack([chosen_logps.detach().sum(), num_pairs])
+        stats["logps/rejected"] += torch.stack([rejected_logps.detach().sum(), num_pairs])
 
         return (loss, outputs) if return_outputs else loss
 
@@ -1773,25 +1761,19 @@ def training_step(self, *args, **kwargs):
     def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
         mode = "train" if self.model.training else "eval"
 
-        # Metrics are accumulated in `compute_loss` as per-rank running sums. Aggregate them across ranks here, in
-        # a single collective per logging window, then compute each `<name>` metric as `<name>_sum / <name>_count`,
-        # i.e. weighted by whatever the count counts (tokens, pairs, batches). Keys are sorted so that every rank
-        # stacks them in the same order.
+        # Sum every `(total, count)` pair across ranks in a single collective, then divide. Keys are sorted so that
+        # all ranks stack them in the same order.
         metrics = {}
-        sums = self._metric_sums[mode]
-        if sums:
-            keys = sorted(sums)
-            values = torch.stack([sums[key].double() for key in keys])
+        stats = self._metric_stats[mode]
+        if stats:
+            keys = sorted(stats)
+            values = torch.stack([stats[key].double() for key in keys])
             totals = dict(zip(keys, self.accelerator.reduce(values, reduction="sum").tolist(), strict=True))
-            for key in keys:
-                if key.endswith("_sum"):
-                    name = key.removesuffix("_sum")
-                    count = totals[name + "_count"]
-                    metrics[name] = totals[key] / count if count > 0 else 0.0
-            # `num_tokens` advances only when a train-mode log folds in the pending sums, so an eval log between two
-            # train logs can lag by up to one logging window.
-            if mode == "train" and "num_tokens_in_batch" in totals:
-                self._total_train_tokens += int(totals["num_tokens_in_batch"])
+            metrics = {key: total / count if count > 0 else 0.0 for key, (total, count) in totals.items()}
+            # `num_tokens` is a running total, so it takes the pair's total instead of the ratio. It only advances
+            # on a train-mode log, so an eval log in between can lag by up to one logging window.
+            if mode == "train" and "num_tokens" in totals:
+                self._total_train_tokens += int(totals["num_tokens"][0])
             metrics["num_tokens"] = self._total_train_tokens
 
         # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs`
@@ -1800,7 +1782,7 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
             metrics = {f"eval_{key}": val for key, val in metrics.items()}
         logs.update(metrics)
         super().log(logs, start_time)
-        self._metric_sums[mode].clear()
+        self._metric_stats[mode].clear()
 
     # During eval, Trainer calls prediction_step. If no labels are present in the inputs, it only runs forward and
     # returns logits. We override prediction_step to force compute_loss, because this trainer doesn't involve labels.
diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index 3b9ee438086..a8240e04bb6 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -1056,10 +1056,11 @@ def cast_outputs_to_original_dtype(module, args, output):
 
         # Initialize the metrics
         self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)}
-        # Per-mode running sums/extrema of on-device metric accumulators. They only receive local tensors (no
-        # collective, no host sync); `log()` aggregates them across ranks in one collective per kind and resets
-        # them. `_metrics` above keeps plain floats for values that are already identical on every rank.
-        self._metric_sums = {"train": defaultdict(int), "eval": defaultdict(int)}
+        # `_metric_stats` entries are running `(total, count)` pairs and `_metric_mins`/`_metric_maxs` are running
+        # extrema, all on-device: they only receive local tensors (no collective, no host sync), and `log()`
+        # aggregates them across ranks in one collective per kind. `_metrics` above keeps plain floats for values
+        # that are already identical on every rank.
+        self._metric_stats = {"train": defaultdict(int), "eval": defaultdict(int)}
         self._metric_mins = {"train": {}, "eval": {}}
         self._metric_maxs = {"train": {}, "eval": {}}
         self._total_train_tokens = 0
@@ -2875,36 +2876,36 @@ def _generate_and_score_completions(
         for name in self._pending_metrics:
             values = self._pending_metrics[name]
             local_mean = sum(values) / len(values)
-            self._metric_sums[mode][f"{name}_sum"] += torch.tensor(local_mean, device=device)
-            self._metric_sums[mode][f"{name}_count"] += torch.ones((), device=device)
+            local_mean = torch.tensor(local_mean, device=device)
+            self._metric_stats[mode][name] += torch.stack([local_mean, torch.ones_like(local_mean)])
         self._pending_metrics.clear()
 
         if images is not None and self.log_multimodal:
             self._logs["images"].extend(gather_object(images))
 
         if self.use_vllm and self.vllm_importance_sampling_correction:
-            sums = self._metric_sums[mode]
+            stats = self._metric_stats[mode]
             mins, maxs = self._metric_mins[mode], self._metric_maxs[mode]
             delta = torch.abs(old_per_token_logps - sampling_per_token_logps)
             mask = completion_mask.bool() if tool_mask is None else (completion_mask * tool_mask).bool()
             # Tokens vLLM could not score carry NaN, so exclude them rather than let them turn the reported
             # divergence into NaN. Counting them as zero instead would understate the divergence.
             delta_valid = mask & ~torch.isnan(delta)
-            sums["sampling/sampling_logp_difference/mean_sum"] += torch.where(delta_valid, delta, 0.0).sum()
-            sums["sampling/sampling_logp_difference/mean_count"] += delta_valid.sum()
+            stats["sampling/sampling_logp_difference/mean"] += torch.stack(
+                [torch.where(delta_valid, delta, 0.0).sum(), delta_valid.sum().float()]
+            )
             max_delta = torch.where(delta_valid, delta, -torch.inf).max()
             key = "sampling/sampling_logp_difference/max"
-            maxs[key] = torch.maximum(maxs.setdefault(key, max_delta), max_delta)
+            maxs[key] = torch.maximum(maxs.get(key, max_delta), max_delta)
             if sequence_level_is:
                 flat_is_ratio = vllm_importance_sampling_ratio.flatten()
             else:
                 flat_is_ratio = vllm_importance_sampling_ratio[mask]
 
             is_ratio_valid = ~torch.isnan(flat_is_ratio)
-            sums["sampling/importance_sampling_ratio/mean_sum"] += torch.where(
-                is_ratio_valid, flat_is_ratio, 0.0
-            ).sum()
-            sums["sampling/importance_sampling_ratio/mean_count"] += is_ratio_valid.sum()
+            stats["sampling/importance_sampling_ratio/mean"] += torch.stack(
+                [torch.where(is_ratio_valid, flat_is_ratio, 0.0).sum(), is_ratio_valid.sum().float()]
+            )
             if flat_is_ratio.numel() > 0:
                 min_is_ratio = torch.where(is_ratio_valid, flat_is_ratio, torch.inf).min()
                 max_is_ratio = torch.where(is_ratio_valid, flat_is_ratio, -torch.inf).max()
@@ -2912,9 +2913,9 @@ def _generate_and_score_completions(
                 min_is_ratio = torch.tensor(torch.inf, device=device)
                 max_is_ratio = torch.tensor(-torch.inf, device=device)
             key = "sampling/importance_sampling_ratio/min"
-            mins[key] = torch.minimum(mins.setdefault(key, min_is_ratio), min_is_ratio)
+            mins[key] = torch.minimum(mins.get(key, min_is_ratio), min_is_ratio)
             key = "sampling/importance_sampling_ratio/max"
-            maxs[key] = torch.maximum(maxs.setdefault(key, max_is_ratio), max_is_ratio)
+            maxs[key] = torch.maximum(maxs.get(key, max_is_ratio), max_is_ratio)
 
         output = {
             "prompt_ids": prompt_ids,
@@ -3006,10 +3007,10 @@ def compute_liger_loss(self, unwrapped_model, inputs):
 
         mode = "train" if self.model.training else "eval"
         if self.beta != 0.0:
-            self._metric_sums[mode]["kl_sum"] += mean_kl.detach()
-            self._metric_sums[mode]["kl_count"] += torch.ones_like(mean_kl)
-        self._metric_sums[mode]["clip_ratio_sum"] += clip_ratio.detach()
-        self._metric_sums[mode]["clip_ratio_count"] += torch.ones_like(clip_ratio)
+            mean_kl = mean_kl.detach()
+            self._metric_stats[mode]["kl"] += torch.stack([mean_kl, torch.ones_like(mean_kl)])
+        clip_ratio = clip_ratio.detach()
+        self._metric_stats[mode]["clip_ratio"] += torch.stack([clip_ratio, torch.ones_like(clip_ratio)])
         # DAPO/CISPO/VESPO normalize by num_items_in_batch / num_processes (applied internally by
         # the Liger loss), then need a `current_gradient_accumulation_steps / steps_per_generation`
         # rescale to land on the per-window token-mean — matching the non-Liger path
@@ -3302,8 +3303,9 @@ def _compute_loss(self, model, inputs):
             loss = loss - apply_coef * entropy_loss
 
             policy_loss_valid = ~torch.isnan(policy_loss)
-            self._metric_sums[mode]["policy_loss_sum"] += torch.where(policy_loss_valid, policy_loss, 0.0).detach()
-            self._metric_sums[mode]["policy_loss_count"] += policy_loss_valid.float()
+            self._metric_stats[mode]["policy_loss"] += torch.stack(
+                [torch.where(policy_loss_valid, policy_loss, 0.0).detach(), policy_loss_valid.float()]
+            )
 
             # Adaptive update. Gated on train mode so evaluation cannot mutate the entropy controller state.
             if self.use_adaptive_entropy and mode == "train":
@@ -3341,8 +3343,10 @@ def _compute_loss(self, model, inputs):
         if self.aux_loss_enabled:
             normalizer = self.current_gradient_accumulation_steps if mode == "train" else 1.0
             loss = loss + self.router_aux_loss_coef * aux_loss / normalizer
-            self._metric_sums[mode]["aux_loss_sum"] += aux_loss.detach()
-            self._metric_sums[mode]["aux_loss_count"] += torch.ones_like(aux_loss)
+            detached_aux_loss = aux_loss.detach()
+            self._metric_stats[mode]["aux_loss"] += torch.stack(
+                [detached_aux_loss, torch.ones_like(detached_aux_loss)]
+            )
 
         # Log the metrics
         def masked_seq_mean(x):
@@ -3356,8 +3360,7 @@ def accumulate_masked_mean(name, x):
                 local_sum, local_count = x.sum(), torch.tensor(float(x.shape[0]), device=x.device)
             else:
                 local_sum, local_count = (x * mask).sum(), mask.sum().float()
-            self._metric_sums[mode][f"{name}_sum"] += local_sum
-            self._metric_sums[mode][f"{name}_count"] += local_count
+            self._metric_stats[mode][name] += torch.stack([local_sum, local_count])
 
         if self.beta != 0.0:
             accumulate_masked_mean("kl", per_token_kl)
@@ -3375,10 +3378,10 @@ def accumulate_masked_mean(name, x):
             mins, maxs = self._metric_mins[mode], self._metric_maxs[mode]
             low_clip_seq = masked_seq_mean(is_low_clipped.float())
             low_min = torch.where(low_clip_seq.isnan(), torch.inf, low_clip_seq).min()
-            mins["clip_ratio/low_min"] = torch.minimum(mins.setdefault("clip_ratio/low_min", low_min), low_min)
+            mins["clip_ratio/low_min"] = torch.minimum(mins.get("clip_ratio/low_min", low_min), low_min)
             high_clip_seq = masked_seq_mean(is_high_clipped.float())
             high_max = torch.where(high_clip_seq.isnan(), -torch.inf, high_clip_seq).max()
-            maxs["clip_ratio/high_max"] = torch.maximum(maxs.setdefault("clip_ratio/high_max", high_max), high_max)
+            maxs["clip_ratio/high_max"] = torch.maximum(maxs.get("clip_ratio/high_max", high_max), high_max)
         elif self.loss_type == "cispo":
             is_cispo_clipped = (coef_1 > self.epsilon_high) & (advantages > 0)
             accumulate_masked_mean("cispo_clip_ratio", is_cispo_clipped.float())
@@ -3409,23 +3412,16 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
             valid = [v for v in val if not math.isnan(v)]
             metrics[key] = sum(valid) / len(valid) if valid else None
 
-        # Metrics accumulated as per-rank running sums are aggregated across ranks here, in a single collective per
-        # logging window: each `<name>` metric is computed as `<name>_sum / <name>_count`, i.e. weighted by
-        # whatever the count counts (tokens, sequences, batches). Keys are sorted so that every rank stacks them in
-        # the same order.
-        sums = self._metric_sums[mode]
-        if sums:
-            keys = sorted(sums)
-            values = torch.stack([sums[key].double() for key in keys])
+        # Sum every `(total, count)` pair across ranks in a single collective, then divide. Keys are sorted so that
+        # all ranks stack them in the same order.
+        stats = self._metric_stats[mode]
+        if stats:
+            keys = sorted(stats)
+            values = torch.stack([stats[key].double() for key in keys])
             totals = dict(zip(keys, self.accelerator.reduce(values, reduction="sum").tolist(), strict=True))
-            for key in keys:
-                if key.endswith("_sum"):
-                    name = key.removesuffix("_sum")
-                    count = totals[name + "_count"]
-                    metrics[name] = totals[key] / count if count > 0 else 0.0
+            metrics.update({key: total / count if count > 0 else 0.0 for key, (total, count) in totals.items()})
         # Running extrema take the min/max across ranks, one collective each (min(x) is computed as -max(-x)). A
-        # rank with no data contributed +/-inf sentinels; a window with no data at all logs None, like the
-        # NaN-filtered metrics above.
+        # rank with no data contributed +/-inf sentinels; a window with no data at all logs None.
         for extrema, sign in ((self._metric_mins[mode], -1.0), (self._metric_maxs[mode], 1.0)):
             if extrema:
                 keys = sorted(extrema)
@@ -3442,7 +3438,7 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
         logs.update(metrics)
         super().log(logs, start_time)
         self._metrics[mode].clear()
-        self._metric_sums[mode].clear()
+        self._metric_stats[mode].clear()
         self._metric_mins[mode].clear()
         self._metric_maxs[mode].clear()
 
diff --git a/trl/trainer/kto_trainer.py b/trl/trainer/kto_trainer.py
index 9adbc3881c5..519a0acf6ff 100644
--- a/trl/trainer/kto_trainer.py
+++ b/trl/trainer/kto_trainer.py
@@ -945,9 +945,10 @@ def __init__(
                 disable_dropout_in_model(self.ref_model)
 
         # Initialize the metrics
-        # Per-mode running sums of on-device metric accumulators. `compute_loss` only adds local tensors here (no
-        # collective, no host sync); `log()` reduces them across ranks in a single collective and resets them.
-        self._metric_sums = {"train": defaultdict(int), "eval": defaultdict(int)}
+        # Each entry is a running `(total, count)` pair of on-device tensors, summed in `compute_loss` with no
+        # collective and no host sync. `log()` reduces the pairs across ranks in a single collective and divides,
+        # so every metric is weighted by whatever its count counts (tokens, pairs, batches).
+        self._metric_stats = {"train": defaultdict(int), "eval": defaultdict(int)}
         self._total_train_tokens = 0
 
         # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the
@@ -1451,28 +1452,23 @@ def _compute_loss_liger(self, model, inputs, return_outputs):
                 kl=kl,
             )
 
-        sums = self._metric_sums[mode]
+        stats = self._metric_stats[mode]
         # `kl` is already global here (the gather above is required by the loss), so only the window average is left
         # for `log()` to compute.
-        sums["kl_sum"] += kl.detach().sum()
-        sums["kl_count"] += torch.ones((), device=kl.device)
+        kl_total = kl.detach().sum()
+        stats["kl"] += torch.stack([kl_total, torch.ones_like(kl_total)])
 
         # Number of tokens
         if mode == "train":
-            sums["num_tokens_in_batch"] += batch["attention_mask"].sum()
-
-        sums["rewards/chosen_sum"] += chosen_rewards_sum.nansum().detach()
-        sums["rewards/chosen_count"] += num_chosen
-        sums["logps/chosen_sum"] += chosen_logps_sum.nansum().detach()
-        sums["logps/chosen_count"] += num_chosen
-        sums["logits/chosen_sum"] += chosen_logits_sum.nansum().detach()
-        sums["logits/chosen_count"] += num_chosen
-        sums["rewards/rejected_sum"] += rejected_rewards_sum.nansum().detach()
-        sums["rewards/rejected_count"] += num_rejected
-        sums["logps/rejected_sum"] += rejected_logps_sum.nansum().detach()
-        sums["logps/rejected_count"] += num_rejected
-        sums["logits/rejected_sum"] += rejected_logits_sum.nansum().detach()
-        sums["logits/rejected_count"] += num_rejected
+            num_tokens_in_batch = batch["attention_mask"].sum()
+            stats["num_tokens"] += torch.stack([num_tokens_in_batch, torch.ones_like(num_tokens_in_batch)])
+
+        stats["rewards/chosen"] += torch.stack([chosen_rewards_sum.nansum().detach(), num_chosen.float()])
+        stats["logps/chosen"] += torch.stack([chosen_logps_sum.nansum().detach(), num_chosen.float()])
+        stats["logits/chosen"] += torch.stack([chosen_logits_sum.nansum().detach(), num_chosen.float()])
+        stats["rewards/rejected"] += torch.stack([rejected_rewards_sum.nansum().detach(), num_rejected.float()])
+        stats["logps/rejected"] += torch.stack([rejected_logps_sum.nansum().detach(), num_rejected.float()])
+        stats["logits/rejected"] += torch.stack([rejected_logits_sum.nansum().detach(), num_rejected.float()])
 
         return loss
 
@@ -1592,21 +1588,21 @@ def _compute_loss(self, model, inputs, return_outputs):
             0,
         )
 
-        sums = self._metric_sums[mode]
+        stats = self._metric_stats[mode]
         # `kl` is already global here (the gather above is required by the loss), so only the window average is left
         # for `log()` to compute.
-        sums["kl_sum"] += kl.detach().sum()
-        sums["kl_count"] += torch.ones((), device=kl.device)
+        kl_total = kl.detach().sum()
+        stats["kl"] += torch.stack([kl_total, torch.ones_like(kl_total)])
 
         # Entropy
         per_token_entropy = entropy_from_logits(shift_logits.detach())
         mask = batch["completion_mask"][:, 1:]
-        sums["entropy_sum"] += (per_token_entropy * mask).sum()
-        sums["entropy_count"] += mask.sum()
+        stats["entropy"] += torch.stack([(per_token_entropy * mask).sum(), mask.sum().float()])
 
         # Number of tokens
         if mode == "train":
-            sums["num_tokens_in_batch"] += batch["attention_mask"].sum()
+            num_tokens_in_batch = batch["attention_mask"].sum()
+            stats["num_tokens"] += torch.stack([num_tokens_in_batch, torch.ones_like(num_tokens_in_batch)])
 
         # Average logits for chosen and rejected completions
         shift_completion_mask = batch["completion_mask"][:, 1:]
@@ -1614,26 +1610,23 @@ def _compute_loss(self, model, inputs, return_outputs):
         rejected_logits = shift_logits.detach().index_select(0, rejected_idx)
         chosen_mask = shift_completion_mask.index_select(0, chosen_idx)
         rejected_mask = shift_completion_mask.index_select(0, rejected_idx)
-        sums["logits/chosen_sum"] += chosen_logits[chosen_mask.bool()].mean(-1).sum()
-        sums["logits/chosen_count"] += chosen_mask.sum()
-        sums["logits/rejected_sum"] += rejected_logits[rejected_mask.bool()].mean(-1).sum()
-        sums["logits/rejected_count"] += rejected_mask.sum()
-
-        sums["rewards/chosen_sum"] += chosen_rewards.nansum()
-        sums["rewards/chosen_count"] += num_chosen
-        sums["logps/chosen_sum"] += chosen_logps.detach().nansum()
-        sums["logps/chosen_count"] += num_chosen
-        sums["rewards/rejected_sum"] += rejected_rewards.nansum()
-        sums["rewards/rejected_count"] += num_rejected
-        sums["logps/rejected_sum"] += rejected_logps.detach().nansum()
-        sums["logps/rejected_count"] += num_rejected
+        stats["logits/chosen"] += torch.stack(
+            [chosen_logits[chosen_mask.bool()].mean(-1).sum(), chosen_mask.sum().float()]
+        )
+        stats["logits/rejected"] += torch.stack(
+            [rejected_logits[rejected_mask.bool()].mean(-1).sum(), rejected_mask.sum().float()]
+        )
+
+        stats["rewards/chosen"] += torch.stack([chosen_rewards.nansum(), num_chosen.float()])
+        stats["logps/chosen"] += torch.stack([chosen_logps.detach().nansum(), num_chosen.float()])
+        stats["rewards/rejected"] += torch.stack([rejected_rewards.nansum(), num_rejected.float()])
+        stats["logps/rejected"] += torch.stack([rejected_logps.detach().nansum(), num_rejected.float()])
 
         loss = losses.nanmean()
         if self.aux_loss_enabled:
             aux_loss = outputs.aux_loss
             loss = loss + self.router_aux_loss_coef * aux_loss
-            sums["aux_loss_sum"] += aux_loss.detach()
-            sums["aux_loss_count"] += torch.ones_like(aux_loss)
+            stats["aux_loss"] += torch.stack([aux_loss.detach(), torch.ones_like(aux_loss)])
 
         return (loss, outputs) if return_outputs else loss
 
@@ -1737,27 +1730,23 @@ def training_step(self, *args, **kwargs):
     def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
         mode = "train" if self.model.training else "eval"
 
-        # Metrics are accumulated in `compute_loss` as per-rank running sums. Aggregate them across ranks here, in
-        # a single collective per logging window, then compute each `<name>` metric as `<name>_sum / <name>_count`,
-        # i.e. weighted by whatever the count counts (tokens, examples, batches). Keys are sorted so that every rank
-        # stacks them in the same order.
+        # Sum every `(total, count)` pair across ranks in a single collective, then divide. Keys are sorted so that
+        # all ranks stack them in the same order.
         metrics = {}
-        sums = self._metric_sums[mode]
-        if sums:
-            keys = sorted(sums)
-            values = torch.stack([sums[key].double() for key in keys])
+        stats = self._metric_stats[mode]
+        if stats:
+            keys = sorted(stats)
+            values = torch.stack([stats[key].double() for key in keys])
             totals = dict(zip(keys, self.accelerator.reduce(values, reduction="sum").tolist(), strict=True))
-            for key in keys:
-                if key.endswith("_sum"):
-                    name = key.removesuffix("_sum")
-                    count = totals[name + "_count"]
-                    metrics[name] = totals[key] / count if count > 0 else 0.0
-            if totals.get("rewards/chosen_count", 0) > 0 and totals.get("rewards/rejected_count", 0) > 0:
+            metrics = {key: total / count if count > 0 else 0.0 for key, (total, count) in totals.items()}
+            # KTO sees chosen and rejected completions in separate, independently sized groups, so the margin is a
+            # difference of the two window means rather than a metric with a count of its own.
+            if totals["rewards/chosen"][1] > 0 and totals["rewards/rejected"][1] > 0:
                 metrics["rewards/margins"] = metrics["rewards/chosen"] - metrics["rewards/rejected"]
-            # `num_tokens` advances only when a train-mode log folds in the pending sums, so an eval log between two
-            # train logs can lag by up to one logging window.
-            if mode == "train" and "num_tokens_in_batch" in totals:
-                self._total_train_tokens += int(totals["num_tokens_in_batch"])
+            # `num_tokens` is a running total, so it takes the pair's total instead of the ratio. It only advances
+            # on a train-mode log, so an eval log in between can lag by up to one logging window.
+            if mode == "train" and "num_tokens" in totals:
+                self._total_train_tokens += int(totals["num_tokens"][0])
             metrics["num_tokens"] = self._total_train_tokens
 
         # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs`
@@ -1766,7 +1755,7 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
             metrics = {f"eval_{key}": val for key, val in metrics.items()}
         logs.update(metrics)
         super().log(logs, start_time)
-        self._metric_sums[mode].clear()
+        self._metric_stats[mode].clear()
 
     # During eval, Trainer calls prediction_step. If no labels are present in the inputs, it only runs forward and
     # returns logits. We override prediction_step to force compute_loss, because this trainer doesn't involve labels.
diff --git a/trl/trainer/reward_trainer.py b/trl/trainer/reward_trainer.py
index 46a6b983ba1..31b3d12dc4c 100644
--- a/trl/trainer/reward_trainer.py
+++ b/trl/trainer/reward_trainer.py
@@ -602,10 +602,10 @@ def __init__(
         else:
             self.maybe_activation_offload_context = contextlib.nullcontext()
 
-        # Per-mode running sums/extrema of on-device metric accumulators. `compute_loss` only adds local tensors
-        # here (no collective, no host sync); `log()` aggregates them across ranks in one collective per kind and
-        # resets them.
-        self._metric_sums = {"train": defaultdict(int), "eval": defaultdict(int)}
+        # `_metric_stats` entries are running `(total, count)` pairs and `_metric_mins`/`_metric_maxs` are running
+        # extrema, all on-device: `compute_loss` only adds local tensors here (no collective, no host sync), and
+        # `log()` aggregates them across ranks in one collective per kind.
+        self._metric_stats = {"train": defaultdict(int), "eval": defaultdict(int)}
         self._metric_mins = {"train": {}, "eval": {}}
         self._metric_maxs = {"train": {}, "eval": {}}
         self._total_train_tokens = 0
@@ -767,26 +767,28 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
             loss += self.args.center_rewards_coefficient * torch.mean((rewards_chosen + rewards_rejected) ** 2)
 
         if mode == "train":
-            self._metric_sums[mode]["num_tokens_in_batch"] += inputs["attention_mask"].sum()
+            num_tokens_in_batch = inputs["attention_mask"].sum()
+            self._metric_stats[mode]["num_tokens"] += torch.stack(
+                [num_tokens_in_batch, torch.ones_like(num_tokens_in_batch)]
+            )
 
         # Compute min, mean, max, accuracy and margin
         with torch.no_grad():
-            sums = self._metric_sums[mode]
+            stats = self._metric_stats[mode]
             mins, maxs = self._metric_mins[mode], self._metric_maxs[mode]
             rewards = outputs.logits.detach()
             num_pairs = torch.tensor(float(rewards_chosen.numel()), device=rewards.device)
-            sums["mean_reward_sum"] += rewards.sum()
-            sums["mean_reward_count"] += torch.tensor(float(rewards.numel()), device=rewards.device)
+            stats["mean_reward"] += torch.stack(
+                [rewards.sum(), torch.tensor(float(rewards.numel()), device=rewards.device)]
+            )
             min_reward = rewards.min()
-            mins["min_reward"] = torch.minimum(mins.setdefault("min_reward", min_reward), min_reward)
+            mins["min_reward"] = torch.minimum(mins.get("min_reward", min_reward), min_reward)
             max_reward = rewards.max()
-            maxs["max_reward"] = torch.maximum(maxs.setdefault("max_reward", max_reward), max_reward)
+            maxs["max_reward"] = torch.maximum(maxs.get("max_reward", max_reward), max_reward)
 
-            sums["accuracy_sum"] += (rewards_chosen > rewards_rejected).float().sum()
-            sums["accuracy_count"] += num_pairs
+            stats["accuracy"] += torch.stack([(rewards_chosen > rewards_rejected).float().sum(), num_pairs])
 
-            sums["margin_sum"] += (rewards_chosen - rewards_rejected).sum()
-            sums["margin_count"] += num_pairs
+            stats["margin"] += torch.stack([(rewards_chosen - rewards_rejected).sum(), num_pairs])
 
         return (loss, outputs) if return_outputs else loss
 
@@ -798,25 +800,19 @@ def training_step(self, *args, **kwargs):
     def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
         mode = "train" if self.model.training else "eval"
 
-        # Metrics are accumulated in `compute_loss` as per-rank running sums. Aggregate them across ranks here, in
-        # a single collective per logging window, then compute each `<name>` metric as `<name>_sum / <name>_count`,
-        # i.e. weighted by whatever the count counts (tokens, batches). Keys are sorted so that every rank stacks
-        # them in the same order.
+        # Sum every `(total, count)` pair across ranks in a single collective, then divide. Keys are sorted so that
+        # all ranks stack them in the same order.
         metrics = {}
-        sums = self._metric_sums[mode]
-        if sums:
-            keys = sorted(sums)
-            values = torch.stack([sums[key].double() for key in keys])
+        stats = self._metric_stats[mode]
+        if stats:
+            keys = sorted(stats)
+            values = torch.stack([stats[key].double() for key in keys])
             totals = dict(zip(keys, self.accelerator.reduce(values, reduction="sum").tolist(), strict=True))
-            for key in keys:
-                if key.endswith("_sum"):
-                    name = key.removesuffix("_sum")
-                    count = totals[name + "_count"]
-                    metrics[name] = totals[key] / count if count > 0 else 0.0
-            # `num_tokens` advances only when a train-mode log folds in the pending sums, so an eval log between two
-            # train logs can lag by up to one logging window.
-            if mode == "train" and "num_tokens_in_batch" in totals:
-                self._total_train_tokens += int(totals["num_tokens_in_batch"])
+            metrics = {key: total / count if count > 0 else 0.0 for key, (total, count) in totals.items()}
+            # `num_tokens` is a running total, so it takes the pair's total instead of the ratio. It only advances
+            # on a train-mode log, so an eval log in between can lag by up to one logging window.
+            if mode == "train" and "num_tokens" in totals:
+                self._total_train_tokens += int(totals["num_tokens"][0])
             metrics["num_tokens"] = self._total_train_tokens
         # Running extrema take the min/max across ranks, one collective each (min(x) is computed as -max(-x)). A
         # rank with no data contributed +/-inf sentinels; a window with no data at all logs None.
@@ -835,7 +831,7 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
 
         logs.update(metrics)
         super().log(logs, start_time)
-        self._metric_sums[mode].clear()
+        self._metric_stats[mode].clear()
         self._metric_mins[mode].clear()
         self._metric_maxs[mode].clear()
 
diff --git a/trl/trainer/rloo_trainer.py b/trl/trainer/rloo_trainer.py
index 6f280b0cf8f..a532d697495 100644
--- a/trl/trainer/rloo_trainer.py
+++ b/trl/trainer/rloo_trainer.py
@@ -624,10 +624,11 @@ def __init__(
 
         # Initialize the metrics
         self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)}
-        # Per-mode running sums/extrema of on-device metric accumulators. They only receive local tensors (no
-        # collective, no host sync); `log()` aggregates them across ranks in one collective per kind and resets
-        # them. `_metrics` above keeps plain floats for values that are already identical on every rank.
-        self._metric_sums = {"train": defaultdict(int), "eval": defaultdict(int)}
+        # `_metric_stats` entries are running `(total, count)` pairs and `_metric_mins`/`_metric_maxs` are running
+        # extrema, all on-device: they only receive local tensors (no collective, no host sync), and `log()`
+        # aggregates them across ranks in one collective per kind. `_metrics` above keeps plain floats for values
+        # that are already identical on every rank.
+        self._metric_stats = {"train": defaultdict(int), "eval": defaultdict(int)}
         self._metric_mins = {"train": {}, "eval": {}}
         self._metric_maxs = {"train": {}, "eval": {}}
         self._total_train_tokens = 0
@@ -1621,8 +1622,9 @@ def _generate_and_score_completions(
 
         # Calculate and log the mean KL divergence between current and reference model
         if self.beta != 0.0:
-            self._metric_sums[mode]["kl_sum"] += (per_token_kl * completion_mask).sum().detach()
-            self._metric_sums[mode]["kl_count"] += completion_mask.sum().float()
+            self._metric_stats[mode]["kl"] += torch.stack(
+                [(per_token_kl * completion_mask).sum().detach(), completion_mask.sum().float()]
+            )
 
         # Calculate mean reward per function, but only for samples where the function was applied (non-NaN values)
         for i, reward_func_name in enumerate(self.reward_func_names):
@@ -1656,8 +1658,8 @@ def _generate_and_score_completions(
         for name in self._pending_metrics:
             values = self._pending_metrics[name]
             local_mean = sum(values) / len(values)
-            self._metric_sums[mode][f"{name}_sum"] += torch.tensor(local_mean, device=device)
-            self._metric_sums[mode][f"{name}_count"] += torch.ones((), device=device)
+            local_mean = torch.tensor(local_mean, device=device)
+            self._metric_stats[mode][name] += torch.stack([local_mean, torch.ones_like(local_mean)])
         self._pending_metrics.clear()
 
         if images is not None and self.log_multimodal:
@@ -1746,30 +1748,30 @@ def _compute_loss(self, model, inputs):
         # RLOO returns an unscaled loss (the HF Trainer divides by gradient accumulation), so add the aux term unscaled
         if self.aux_loss_enabled:
             loss = loss + self.router_aux_loss_coef * aux_loss
-            self._metric_sums[mode]["aux_loss_sum"] += aux_loss.detach()
-            self._metric_sums[mode]["aux_loss_count"] += torch.ones_like(aux_loss)
+            detached_aux_loss = aux_loss.detach()
+            self._metric_stats[mode]["aux_loss"] += torch.stack(
+                [detached_aux_loss, torch.ones_like(detached_aux_loss)]
+            )
 
         # Entropy
-        self._metric_sums[mode]["entropy_sum"] += (entropies * completion_mask).sum().detach()
-        self._metric_sums[mode]["entropy_count"] += completion_mask.sum().float()
+        self._metric_stats[mode]["entropy"] += torch.stack(
+            [(entropies * completion_mask).sum().detach(), completion_mask.sum().float()]
+        )
 
         # Compute the clipped probability ratios
-        sums = self._metric_sums[mode]
+        stats = self._metric_stats[mode]
         mins, maxs = self._metric_mins[mode], self._metric_maxs[mode]
         is_low_clipped = (coef_1 < 1 - self.epsilon_low) & (advantages < 0)
         is_high_clipped = (coef_1 > 1 + self.epsilon_high) & (advantages > 0)
         is_region_clipped = is_low_clipped | is_high_clipped
         clip_count = torch.tensor(float(is_low_clipped.numel()), device=is_low_clipped.device)
-        sums["clip_ratio/low_mean_sum"] += is_low_clipped.float().sum()
-        sums["clip_ratio/low_mean_count"] += clip_count
+        stats["clip_ratio/low_mean"] += torch.stack([is_low_clipped.float().sum(), clip_count])
         low_min = is_low_clipped.float().min()
-        mins["clip_ratio/low_min"] = torch.minimum(mins.setdefault("clip_ratio/low_min", low_min), low_min)
-        sums["clip_ratio/high_mean_sum"] += is_high_clipped.float().sum()
-        sums["clip_ratio/high_mean_count"] += clip_count
+        mins["clip_ratio/low_min"] = torch.minimum(mins.get("clip_ratio/low_min", low_min), low_min)
+        stats["clip_ratio/high_mean"] += torch.stack([is_high_clipped.float().sum(), clip_count])
         high_max = is_high_clipped.float().max()
-        maxs["clip_ratio/high_max"] = torch.maximum(maxs.setdefault("clip_ratio/high_max", high_max), high_max)
-        sums["clip_ratio/region_mean_sum"] += is_region_clipped.float().sum()
-        sums["clip_ratio/region_mean_count"] += clip_count
+        maxs["clip_ratio/high_max"] = torch.maximum(maxs.get("clip_ratio/high_max", high_max), high_max)
+        stats["clip_ratio/region_mean"] += torch.stack([is_region_clipped.float().sum(), clip_count])
         return loss
 
     # During eval, Trainer calls prediction_step. If no labels are present in the inputs, it only runs forward and
@@ -1794,23 +1796,16 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
             valid = [v for v in val if not math.isnan(v)]
             metrics[key] = sum(valid) / len(valid) if valid else None
 
-        # Metrics accumulated as per-rank running sums are aggregated across ranks here, in a single collective per
-        # logging window: each `<name>` metric is computed as `<name>_sum / <name>_count`, i.e. weighted by
-        # whatever the count counts (tokens, sequences, batches). Keys are sorted so that every rank stacks them in
-        # the same order.
-        sums = self._metric_sums[mode]
-        if sums:
-            keys = sorted(sums)
-            values = torch.stack([sums[key].double() for key in keys])
+        # Sum every `(total, count)` pair across ranks in a single collective, then divide. Keys are sorted so that
+        # all ranks stack them in the same order.
+        stats = self._metric_stats[mode]
+        if stats:
+            keys = sorted(stats)
+            values = torch.stack([stats[key].double() for key in keys])
             totals = dict(zip(keys, self.accelerator.reduce(values, reduction="sum").tolist(), strict=True))
-            for key in keys:
-                if key.endswith("_sum"):
-                    name = key.removesuffix("_sum")
-                    count = totals[name + "_count"]
-                    metrics[name] = totals[key] / count if count > 0 else 0.0
+            metrics.update({key: total / count if count > 0 else 0.0 for key, (total, count) in totals.items()})
         # Running extrema take the min/max across ranks, one collective each (min(x) is computed as -max(-x)). A
-        # rank with no data contributed +/-inf sentinels; a window with no data at all logs None, like the
-        # NaN-filtered metrics above.
+        # rank with no data contributed +/-inf sentinels; a window with no data at all logs None.
         for extrema, sign in ((self._metric_mins[mode], -1.0), (self._metric_maxs[mode], 1.0)):
             if extrema:
                 keys = sorted(extrema)
@@ -1827,7 +1822,7 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
         logs.update(metrics)
         super().log(logs, start_time)
         self._metrics[mode].clear()
-        self._metric_sums[mode].clear()
+        self._metric_stats[mode].clear()
         self._metric_mins[mode].clear()
         self._metric_maxs[mode].clear()
 
diff --git a/trl/trainer/sft_trainer.py b/trl/trainer/sft_trainer.py
index c48ff843c20..0770bbf076a 100644
--- a/trl/trainer/sft_trainer.py
+++ b/trl/trainer/sft_trainer.py
@@ -1414,10 +1414,10 @@ def __init__(
             text_config.output_router_logits = self.aux_loss_enabled
             text_config.router_aux_loss_coef = self.args.router_aux_loss_coef
 
-        # Initialize the metrics
-        # Per-mode running sums of on-device metric accumulators. `compute_loss` only adds local tensors here (no
-        # collective, no host sync); `log()` reduces them across ranks in a single collective and resets them.
-        self._metric_sums = {"train": defaultdict(int), "eval": defaultdict(int)}
+        # Initialize the metrics. Each entry is a running `(total, count)` pair of on-device tensors, summed in
+        # `compute_loss` with no collective and no host sync. `log()` reduces the pairs across ranks in a single
+        # collective and divides, so every metric is weighted by whatever its count counts (tokens, batches).
+        self._metric_stats = {"train": defaultdict(int), "eval": defaultdict(int)}
         self._total_train_tokens = 0
 
         # Add tags to the model
@@ -1806,8 +1806,8 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
             # PEFT (PromptTuning, P-Tuning) prepends `-100`-padded virtual tokens before delegating into the patched
             # forward, so the valid-token count over the padded labels can differ from the un-padded `labels[..., 1:]`
             # count by up to one per sequence; using the patched output keeps numerator and denominator aligned.
-            self._metric_sums[mode]["entropy_sum"] += outputs.entropy_sum.detach()
-            self._metric_sums[mode]["entropy_count"] += outputs.num_valid_tokens.detach()
+            num_valid = outputs.num_valid_tokens.detach().float()
+            self._metric_stats[mode]["entropy"] += torch.stack([outputs.entropy_sum.detach(), num_valid])
         elif not self.args.use_liger_kernel:  # liger doesn't return logits
             with torch.no_grad():
                 if "shift_labels" in inputs:
@@ -1833,10 +1833,10 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
                 total_tokens = mask.sum()
                 correct_predictions = (predictions == shift_labels) & mask
                 correct_tokens = correct_predictions.sum()
-            self._metric_sums[mode]["entropy_sum"] += entropy_sum
-            self._metric_sums[mode]["entropy_count"] += total_tokens
-            self._metric_sums[mode]["mean_token_accuracy_sum"] += correct_tokens
-            self._metric_sums[mode]["mean_token_accuracy_count"] += total_tokens
+            self._metric_stats[mode]["entropy"] += torch.stack([entropy_sum, total_tokens.float()])
+            self._metric_stats[mode]["mean_token_accuracy"] += torch.stack(
+                [correct_tokens.float(), total_tokens.float()]
+            )
 
         if mode == "train":
             # When using padding-free, the attention_mask is not present in the inputs, instead we have cu_seq_lens_q,
@@ -1844,18 +1844,27 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
             if "attention_mask" in inputs:
                 num_tokens_in_batch = inputs["attention_mask"].sum()
             elif "position_ids" in inputs:
-                num_tokens_in_batch = torch.tensor(inputs["position_ids"].size(1), device=inputs["position_ids"].device)
+                num_tokens_in_batch = torch.tensor(
+                    inputs["position_ids"].size(1), device=inputs["position_ids"].device
+                )
             else:
                 raise ValueError("Expected 'attention_mask' or 'position_ids' in inputs.")
-            self._metric_sums[mode]["num_tokens_in_batch"] += num_tokens_in_batch
+            # `num_tokens` is a running total rather than a window average, so `log()` uses the pair's total and
+            # ignores its count.
+            self._metric_stats[mode]["num_tokens"] += torch.stack(
+                [num_tokens_in_batch, torch.ones_like(num_tokens_in_batch)]
+            )
 
         if self.args.loss_type == "chunked_nll":
-            self._metric_sums[mode]["mean_token_accuracy_sum"] += outputs.num_correct_tokens.detach()
-            self._metric_sums[mode]["mean_token_accuracy_count"] += outputs.num_valid_tokens.detach()
+            self._metric_stats[mode]["mean_token_accuracy"] += torch.stack(
+                [outputs.num_correct_tokens.detach(), num_valid]
+            )
         elif self.args.use_liger_kernel:
             if hasattr(outputs, "token_accuracy") and outputs.token_accuracy is not None:
-                self._metric_sums[mode]["mean_token_accuracy_sum"] += outputs.token_accuracy.detach()
-                self._metric_sums[mode]["mean_token_accuracy_count"] += torch.ones_like(outputs.token_accuracy)
+                token_accuracy = outputs.token_accuracy.detach()
+                self._metric_stats[mode]["mean_token_accuracy"] += torch.stack(
+                    [token_accuracy, torch.ones_like(token_accuracy)]
+                )
             else:
                 warnings.warn(
                     "liger-kernel did not return token_accuracy when requested. The mean_token_accuracy metric will "
@@ -1864,8 +1873,8 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
                 )
         # Log auxiliary loss if enabled (applies to both Liger and non-Liger)
         if self.aux_loss_enabled:
-            self._metric_sums[mode]["aux_loss_sum"] += outputs.aux_loss.detach()
-            self._metric_sums[mode]["aux_loss_count"] += torch.ones_like(outputs.aux_loss)
+            aux_loss = outputs.aux_loss.detach()
+            self._metric_stats[mode]["aux_loss"] += torch.stack([aux_loss, torch.ones_like(aux_loss)])
 
         return (loss, outputs) if return_outputs else loss
 
@@ -1882,25 +1891,19 @@ def training_step(self, *args, **kwargs):
     def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
         mode = "train" if self.model.training else "eval"
 
-        # Metrics are accumulated in `compute_loss` as per-rank running sums. Aggregate them across ranks here, in
-        # a single collective per logging window, then compute each `<name>` metric as `<name>_sum / <name>_count`,
-        # i.e. weighted by whatever the count counts (tokens, batches). Keys are sorted so that every rank stacks
-        # them in the same order.
+        # Sum every `(total, count)` pair across ranks in a single collective, then divide. Keys are sorted so that
+        # all ranks stack them in the same order.
         metrics = {}
-        sums = self._metric_sums[mode]
-        if sums:
-            keys = sorted(sums)
-            values = torch.stack([sums[key].double() for key in keys])
+        stats = self._metric_stats[mode]
+        if stats:
+            keys = sorted(stats)
+            values = torch.stack([stats[key].double() for key in keys])
             totals = dict(zip(keys, self.accelerator.reduce(values, reduction="sum").tolist(), strict=True))
-            for key in keys:
-                if key.endswith("_sum"):
-                    name = key.removesuffix("_sum")
-                    count = totals[name + "_count"]
-                    metrics[name] = totals[key] / count if count > 0 else 0.0
-            # `num_tokens` advances only when a train-mode log folds in the pending sums, so an eval log between two
-            # train logs can lag by up to one logging window.
-            if mode == "train" and "num_tokens_in_batch" in totals:
-                self._total_train_tokens += int(totals["num_tokens_in_batch"])
+            metrics = {key: total / count if count > 0 else 0.0 for key, (total, count) in totals.items()}
+            # `num_tokens` is a running total, so it takes the pair's total instead of the ratio. It only advances
+            # on a train-mode log, so an eval log in between can lag by up to one logging window.
+            if mode == "train" and "num_tokens" in totals:
+                self._total_train_tokens += int(totals["num_tokens"][0])
             metrics["num_tokens"] = self._total_train_tokens
 
         # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs`
@@ -1910,7 +1913,7 @@ def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
 
         logs.update(metrics)
         super().log(logs, start_time)
-        self._metric_sums[mode].clear()
+        self._metric_stats[mode].clear()
 
     # Ensure the model card is saved along with the checkpoint
     def _save_checkpoint(self, model, trial):