From e391706420934f6c87cebe9997fc85a757aa4353 Mon Sep 17 00:00:00 2001 From: Alexander Wettig Date: Tue, 16 Jul 2024 20:35:10 +0200 Subject: [PATCH] Fix gather when collecting 'num_input_tokens_seen' (#31974) * Move token count to device before gathering * Run 'make style; make quality' --- src/transformers/trainer.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index 485f6cd61e0e7b..25b7b69930921d 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -2245,12 +2245,17 @@ def _inner_training_loop( "a `main_input_name` attribute to the model class you are using." ) else: - input_device = inputs[main_input_name].device - self.state.num_input_tokens_seen += torch.sum( - self.accelerator.gather( - torch.tensor(inputs[main_input_name].numel(), device=input_device, dtype=torch.int64) + self.state.num_input_tokens_seen += ( + torch.sum( + self.accelerator.gather( + torch.tensor( + inputs[main_input_name].numel(), device=self.args.device, dtype=torch.int64 + ) + ) ) - ).item() + .cpu() + .item() + ) if rng_to_sync: self._load_rng_state(resume_from_checkpoint) rng_to_sync = False