enable average tokens across devices

src/transformers/trainer.py

@@ -3603,6 +3603,9 @@ def training_step(
                 scaled_loss.backward()
         else:
             loss *= self.args.gradient_accumulation_steps
+            if (self.args.average_tokens_across_devices and num_items_in_batch is not None and
+                    self.args.world_size > 1):
+                loss *= self.args.world_size
             self.accelerator.backward(loss, **kwargs)
 
         return loss.detach() / self.args.gradient_accumulation_steps
@@ -3617,6 +3620,10 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
             labels = inputs.pop("labels")
         else:
             labels = None
+        if (self.args.average_tokens_across_devices and num_items_in_batch is not None and
+                self.args.world_size > 1):
+            num_items_in_batch_tensor = torch.tensor(num_items_in_batch, device=self.args.device)
+            num_items_in_batch = int(self.accelerator.gather(num_items_in_batch_tensor).sum().cpu())
         if self.model_accepts_loss_kwargs:
             loss_kwargs = {}
             if num_items_in_batch is not None:

src/transformers/training_args.py

@@ -1530,6 +1530,13 @@ class TrainingArguments:
         },
     )
 
+    average_tokens_across_devices: Optional[bool] = field(
+        default=False,
+        metadata={
+            "help": "Whether or not to average tokens across devices."
+        }
+    )
+
     def __post_init__(self):
         # Parse in args that could be `dict` sent in from the CLI as a string
         for field in _VALID_DICT_FIELDS: