dpo_trainer gather metrics across ranks before logging #2474
Conversation
| @@ -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]) | |||
There was a problem hiding this comment.
maybe self.accelerator.gather(metrics).mean().item() would be simpler?
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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:
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.
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
according to #2468
What does this PR do?
Fixes #2468
Before submitting
Pull Request section?
to it if that's the case.
documentation guidelines.
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.