From dd59ef013b9652aedb7169e45e35852a7b5ea974 Mon Sep 17 00:00:00 2001 From: CodeCreator Date: Mon, 15 Jul 2024 07:49:36 -0400 Subject: [PATCH 1/2] Move token count to device before gathering --- src/transformers/trainer.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index 71c3ee43af2c..6715e4886c90 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -2239,12 +2239,11 @@ 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) + 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 From cbbca13ed3f514a7250bbee17f6d0e0aed6672e2 Mon Sep 17 00:00:00 2001 From: CodeCreator Date: Tue, 16 Jul 2024 10:24:51 -0400 Subject: [PATCH 2/2] Run 'make style; make quality' --- src/transformers/trainer.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index 7d9a960c260a..25b7b6993092 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -2245,11 +2245,17 @@ def _inner_training_loop( "a `main_input_name` attribute to the model class you are using." ) else: - 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) + 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 + ) + ) ) - ).cpu().item() + .cpu() + .item() + ) if rng_to_sync: self._load_rng_state(resume_from_checkpoint) rng_to_sync = False