From 79d03df21019af641884ab38ffc83abac0cd2928 Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Mon, 23 Dec 2024 21:15:19 +0530 Subject: [PATCH 1/2] Handle `None` before gathering num items in batch across devices --- src/transformers/trainer.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index c878d2b345cc31..803b3c4ef429d8 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -2465,6 +2465,7 @@ def _inner_training_loop( if remainder == 0: remainder = args.gradient_accumulation_steps update_step = -1 + # Note: total_updates can be larger than dataloader length total_updates = steps_in_epoch // args.gradient_accumulation_steps + 1 for _ in range(total_updates): update_step += 1 @@ -5152,6 +5153,9 @@ def get_batch_samples(self, epoch_iterator, num_batches): except (TypeError, AttributeError): pass + if num_items_in_batch is None: + num_items_in_batch = torch.tensor(0).to(self.args.device) + if self.args.average_tokens_across_devices: num_items_in_batch = self.accelerator.gather(num_items_in_batch).sum().item() From ebffbc7a963ab1c4d471392c252ef07a54ab559b Mon Sep 17 00:00:00 2001 From: Chirag Jain Date: Mon, 23 Dec 2024 21:33:58 +0530 Subject: [PATCH 2/2] Address review comments --- 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 803b3c4ef429d8..93b67a239fcc8e 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -5153,10 +5153,9 @@ def get_batch_samples(self, epoch_iterator, num_batches): except (TypeError, AttributeError): pass - if num_items_in_batch is None: - num_items_in_batch = torch.tensor(0).to(self.args.device) - if self.args.average_tokens_across_devices: + if num_items_in_batch is None: + num_items_in_batch = torch.tensor(0).to(self.args.device) num_items_in_batch = self.accelerator.gather(num_items_in_batch).sum().item() if torch.is_tensor(num_items_in_batch):