Repository: huggingface/transformers
Pull request: fix bug in distributed loss test

# What does this PR do?

Fixes distributed loss test failed as mentioned in https://github.com/huggingface/transformers/pull/35743#issuecomment-2883577747


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?

@ydshieh @muellerzr @ @SunMarc 

## Tests on 2 GPUs:
base_loss: [8.4486, 9.0954, 7.8702, 7.9928, 7.2161, 8.3778, 8.0487, 6.1117, 7.283, 6.4819]
broken_loss: [8.4486, 9.3632, 7.6003, 7.4984, 6.6267, 7.4892, 7.9702, 5.6609, 6.8272, 6.3275]
fixed_loss: [8.4486, 9.0954, 7.8702, 7.9928, 7.2161, 8.3778, 8.0487, 6.1117, 7.283, 6.4819]

## Tests on 8 GPUs:
base_loss: [9.1648, 8.8716, 7.6154, 7.4581, 6.6029, 6.2143, 5.2902, 5.4008, 5.253, 4.9834]
broken_loss: [9.387, 8.7627, 7.7227, 9.42, 7.2081, 5.9345, 5.0765, 6.0079, 4.9359, 5.4696]
fixed_loss: [9.1648, 8.8716, 7.6154, 7.4581, 6.6029, 6.2143, 5.2902, 5.4008, 5.253, 4.9834]

## Others
I accidentally closed https://github.com/huggingface/transformers/pull/36987. So I add the doc fix in this PR.


Changed files:
filename: src/transformers/trainer.py
patch: @@ -3784,7 +3784,7 @@ def training_step(
             with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                 scaled_loss.backward()
         else:
-            # Finally we need to normalize the loss for reporting
+            # Finally we need to normalize the loss for reporting if GA loss bug is not fixed during compute loss
             if not self.model_accepts_loss_kwargs and self.compute_loss_func is None:
                 loss = loss / self.args.gradient_accumulation_steps
 

filename: tests/trainer/test_trainer_distributed_loss.py
patch: @@ -26,7 +26,7 @@ class TestTrainerDistributedLoss(TestCasePlus):
     @require_torch_multi_accelerator
     def test_trainer(self):
         device_count = backend_device_count(torch_device)
-        min_bs = 1
+        min_bs = 2
         output_dir = self.get_auto_remove_tmp_dir()
         for gpu_num, enable, bs, name in (
             (1, True, min_bs * device_count, "base"),
@@ -50,9 +50,10 @@ def test_trainer(self):
         broken_diff = [abs(base_loss[i] - broken_loss[i]) for i in range(len(base_loss))]
         fixed_diff = [abs(base_loss[i] - fixed_loss[i]) for i in range(len(base_loss))]
         sum_base = sum(base_loss)
-        sum_broken = sum(broken_diff)
+        sum_broken = sum(broken_loss)
         relative_broken = abs(sum_base - sum_broken) / max(sum_base, sum_broken)
 
+        # the gap may be smaller for other models, but it still ok.
         self.assertGreater(max(broken_diff), 0.5)
         self.assertLess(max(fixed_diff), 0.005)
         self.assertLess(relative_broken, 0.1)
@@ -63,7 +64,7 @@ def run_distributed_training(training_args):
     model_name = "nickypro/tinyllama-15M"
     dataset_name = "wikitext"
     dataset_config = "wikitext-2-raw-v1"
-    dataset = datasets.load_dataset(dataset_name, dataset_config, split="train[:17]")
+    dataset = datasets.load_dataset(dataset_name, dataset_config, split="train[:100]")
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     tokenizer.pad_token = tokenizer.eos_token