enable average tokens across devices

src/transformers/trainer.py

@@ -3636,7 +3636,15 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
                 model_name = unwrapped_model._get_name()
             # User-defined compute_loss function
             if self.compute_loss_func is not 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())
+                    device_count_for_loss = self.args.world_size
+                else:
+                    device_count_for_loss = 1
                 loss = self.compute_loss_func(outputs, labels, num_items_in_batch=num_items_in_batch)
+                loss *= device_count_for_loss
             elif model_name in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.values():
                 loss = self.label_smoother(outputs, labels, shift_labels=True)
             else:

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: