Repository: huggingface/trl
Pull request: 🧑‍🤝‍🧑 Proper metrics gathering across ranks before logging
according to https://github.com/huggingface/trl/issues/2468

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly. They may suggest changes to make the code even better.
-->

<!-- Remove if not applicable -->

Fixes #2468

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/trl/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a GitHub issue? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/trl/tree/main/docs).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

Problem statement:
DPOTrainer log metrics are not gathered and meaned across ranks
### Feature request

synchronize and average metrics across ranks.

### Motivation

current metrics reported are only numbers on rank 0.

```python
        metrics[f"{prefix}rewards/chosen"] = chosen_rewards.mean().cpu()
        metrics[f"{prefix}rewards/rejected"] = rejected_rewards.mean().cpu()
        metrics[f"{prefix}rewards/accuracies"] = reward_accuracies.mean().cpu()
        metrics[f"{prefix}rewards/margins"] = (chosen_rewards - rejected_rewards).mean().cpu()
        metrics[f"{prefix}logps/chosen"] = model_output["chosen_logps"].detach().mean().cpu()
        metrics[f"{prefix}logps/rejected"] = model_output["rejected_logps"].detach().mean().cpu()
        metrics[f"{prefix}logits/chosen"] = model_output["mean_chosen_logits"].detach().cpu()
        metrics[f"{prefix}logits/rejected"] = model_output["mean_rejected_logits"].detach().cpu()
```

all of these aren't synced.

### Your contribution

current log function looks like:

```python
    def log(self, logs: Dict[str, float]) -> None:
        """
        Log `logs` on the various objects watching training, including stored metrics.

        Args:
            logs (`Dict[str, float]`):
                The values to log.
        """
        # logs either has 'loss' or 'eval_loss'
        train_eval = "train" if "loss" in logs else "eval"
        # Add averaged stored metrics to logs
        for key, metrics in self._stored_metrics[train_eval].items():
            logs[key] = torch.tensor(metrics).mean().item()
        del self._stored_metrics[train_eval]
        return super().log(logs)
  ```
  
  it would have this feature if it looks like:
  
  ```python
    def log(self, logs: Dict[str, float]) -> None:
        """
        Log `logs` on the various objects watching training, including stored metrics.

        Args:
            logs (`Dict[str, float]`):
                The values to log.
        """
        # logs either has 'loss' or 'eval_loss'
        train_eval = "train" if "loss" in logs else "eval"
        # Add averaged stored metrics to logs
        for key, metrics in self._stored_metrics[train_eval].items():
            if isinstance(metrics[0], torch.Tensor):
                gathered = self._nested_gather([m.cuda() for m in metrics])
                metrics = [g.mean() for g in gathered]
            meaned = torch.tensor(metrics).mean()
            logs[key] = meaned.item()
        del self._stored_metrics[train_eval]
        return super().log(logs)
```

I'm happy to submit a pr.

Hints from issue discussion:
That's a good point! Feel free to open a PR to fix this. I don't think adding a unittest for this is relevant. If possible, add plots (eg, with wandb) before/after to ensure that we aren't introducing a regression
Ofcourse!
![image](https://github.com/user-attachments/assets/2da93fdf-a29d-41a1-974a-2b640e3a6ee6)
here's a graph for the same training with and without the modification. You can see the pink line is a lot more smoother. Especially the accuracy graph. My `per_device_batch_size` is 2 so the accuracy per device can only be 1, 0.5 or 0.

Commit message: dpo_trainer gather metrics across ranks before logging

according to https://github.com/huggingface/trl/issues/2468

Patch to review:
diff --git a/trl/trainer/dpo_trainer.py b/trl/trainer/dpo_trainer.py
index 7ed0ac387f..f26ae9b4b2 100644
--- a/trl/trainer/dpo_trainer.py
+++ b/trl/trainer/dpo_trainer.py
@@ -1424,7 +1424,11 @@ def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> Non
         train_eval = "train" if "loss" in logs else "eval"
         # Add averaged stored metrics to logs
         for key, metrics in self._stored_metrics[train_eval].items():
-            logs[key] = torch.tensor(metrics).mean().item()
+            if isinstance(metrics[0], torch.Tensor):
+                gathered = self._nested_gather([m.cuda() for m in metrics])
+                metrics = [g.mean() for g in gathered]
+            meaned = torch.tensor(metrics).mean()
+            logs[key] = meaned.item()
         del self._stored_metrics[train_eval]
 
         if version.parse(transformers.__version__) >= version.parse("4.47.0.dev0"):


Code review comments:
diff_hunk: @@ -1424,7 +1424,11 @@ def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> Non
         train_eval = "train" if "loss" in logs else "eval"
         # Add averaged stored metrics to logs
         for key, metrics in self._stored_metrics[train_eval].items():
-            logs[key] = torch.tensor(metrics).mean().item()
+            if isinstance(metrics[0], torch.Tensor):
+                gathered = self._nested_gather([m.cuda() for m in metrics])
path: trl/trainer/dpo_trainer.py
text: @user1:
do you need `.cuda()`?

@user1:
maybe `self.accelerator.gather(metrics).mean().item()` would be simpler?

@author:
> do you need `.cuda()`?

metrics are moved cpu before, some backends (e.g. nccl) does not support gathering tensors on cpu. but I admit `.cuda` here loses some genrality. maybe `.to(self.accelerator.device)` is better?

@author:
> maybe `self.accelerator.gather(metrics).mean().item()` would be simpler?

I agree `self.accelerator.gather` is better. but `metrics` in the loop is a `list[torch.Tensor]` or `list[float]`, so `gather` actually returns a `list[torch.Tensor]`. So I think I should change it into:

```python
            if isinstance(metrics[0], torch.Tensor):
                gathered = self.accelerator.gather([m.to(self.accelerator.device) for m in metrics])
                metrics = [g.mean() for g in gathered]
            meaned = torch.tensor(metrics).mean()
            logs[key] = meaned.item()
```

I know creating a new tensor on metrics seems a little weird, but that is how it was originally written. I don't why but I don't why to break anything so I left it there.