Scale loss before backward
Author: qgallouedec
Pull request: yes

# What does this PR do?


Fixes https://github.com/huggingface/trl/issues/2456

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] 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.
- [ ] 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).
- [ ] Did you write any new necessary tests?


Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_35207). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

> Do the slow tests pass under this?

They pass locally. I get errors but I think they aren't related:

```
============================== short test summary info ==============================
FAILED tests/trainer/test_trainer.py::TrainerIntegrationTest::test_auto_batch_size_finder - AttributeError: 'DownloadConfig' object has no attribute 'use_auth_token'
FAILED tests/trainer/test_trainer.py::TrainerIntegrationTest::test_auto_batch_size_with_deepspeed - ModuleNotFoundError: No module named 'mpi4py'
FAILED tests/trainer/test_trainer.py::TrainerIntegrationTest::test_use_liger_kernel_patching - AssertionError: False is not true
FAILED tests/trainer/test_trainer_seq2seq.py::Seq2seqTrainerTester::test_finetune_bert2bert - TypeError: BertModel.forward() got an unexpected keyword argument 'num_items_in_batch'
======== 4 failed, 250 passed, 82 skipped, 89 warnings in 203.93s (0:03:23) =========

@ArthurZucker @muellerzr @SunMarc I'll let you merge when you think it's good.

Merging! 🤗 

@qgallouedec @SunMarc @ArthurZucker @muellerzr
Hi, I've identified that there still exists a potential issue with the loss reporting in the Trainer class. 
The problem appears to stem from changes introduced after version 4.45.2, specifically in this line:
https://github.com/huggingface/transformers/blob/24c91f095fec4d90fa6901ef17146b4f4c21d0a3/src/transformers/trainer.py#L3703
The current implementation only divides the loss scalar when `num_items_in_batch` is `None`. However, I believe we should divide the loss scalar regardless of the `num_items_in_batch value`, as was the case in `transformers<=4.45.2`.
This change has led to an unexpected doubling of the logged loss when `gradient_accumulated_steps > 1`. 
You can observe the differences by comparing versions:
https://github.com/huggingface/transformers/compare/v4.45.2...v4.46.0.



@qgallouedec Hello, not sure If I understand correctly, does this commit mean `gradient_accumulated_steps` will never take effect, as `num_items_in_batch` always be calculated?

@qgallouedec Just wondering have you observed consistant hehaviors in

https://github.com/huggingface/trl/issues/2456

when switching from transformers v4.45.2 to v4.46
I think this two are quite different regarding loss with gradient accumulation.

I think you're wrong @yzhangcs:


if I apply your modification:

```diff
- if num_items_in_batch is None:
-     loss = loss / self.args.gradient_accumulation_steps
+ loss = loss / self.args.gradient_accumulation_steps
```

and run the following script with

- `gradient_accumulation_steps = 2` and   `per_device_train_batch_size = 3`
- `gradient_accumulation_steps = 1` and   `per_device_train_batch_size = 6`

```python
import torch
from transformers import AutoModelForCausalLM, Trainer, TrainingArguments
from datasets import Dataset

num_batch = 32
gradient_accumulation_steps = 2  # or 1
per_device_train_batch_size = 3  # or 6
seq_len = 5

eff_batch_size = per_device_train_batch_size * gradient_accumulation_steps
dataset_len = num_batch * eff_batch_size

data = torch.arange(0, dataset_len * seq_len)
data = data.reshape(dataset_len, seq_len)
data = data.tolist()

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-1.5B").to("cuda")
dataset = Dataset.from_dict({"input_ids": data, "labels": data})

args = TrainingArguments(
    output_dir=f"out_bs_{batch_size}_grad_{grad_accum_steps}_before",
    per_device_train_batch_size= per_device_train_batch_size,
    gradient_accumulation_steps= gradient_accumulation_steps,
    logging_steps=2,
)

trainer = Trainer(model=model, args=args, train_dataset=dataset)

trainer.train()
```

You get the following results:

#### Without the modification (current main)
<img width="1028" alt="Screenshot 2024-12-26 at 12 22 24" src="https://github.com/user-attachments/assets/3fdd7f7d-ab81-40ac-bb8f-f1557980793a" />


#### With the modif

<img width="1028" alt="Screenshot 2024-12-26 at 12 22 32" src="https://github.com/user-attachments/assets/3e0c6395-7c10-4dba-81f3-7f0752a35375" />



After the suggested modification, altering the gradient accumulation steps alter the results. While before the modification, everything overlap nicely.



@qgallouedec Hi, thank you for sharing these easy-to-run scripts! 
I successfully ran them on my machine with 4 GPUs.
The good news is that I found the loss to be consistent when varying the grad_accum and bsz. 
However, the bad news is that I noticed significant differences in the loss when switching between versions v4.45.2 and v4.47.

FYI

**For v4.45.2, everything seems fine across different settings:**

`gradient_accumulation_steps=1, per_device_train_batch_size=6` 
```
{'loss': 10.5949, 'grad_norm': 188.29049682617188, 'learning_rate': 4.5833333333333334e-05, 'epoch': 0.25}                                                                                                             
{'loss': 8.2312, 'grad_norm': 87.74673461914062, 'learning_rate': 4.166666666666667e-05, 'epoch': 0.5}                                                                                                                 
{'loss': 7.8988, 'grad_norm': 74.35655212402344, 'learning_rate': 3.7500000000000003e-05, 'epoch': 0.75}                                                                                                               
{'loss': 8.0197, 'grad_norm': 52.913124084472656, 'learning_rate': 3.3333333333333335e-05, 'epoch': 1.0}                                                                                                               
{'loss': 6.0704, 'grad_norm': 51.57878875732422, 'learning_rate': 2.916666666666667e-05, 'epoch': 1.25}                                                                                                                
{'loss': 5.8247, 'grad_norm': 56.05608367919922, 'learning_rate': 2.5e-05, 'epoch': 1.5}                                                                                                                               
{'loss': 5.9932, 'grad_norm': 50.10258483886719, 'learning_rate': 2.0833333333333336e-05, 'epoch': 1.75}                                                                                                               
{'loss': 5.3769, 'grad_norm': 65.85218811035156, 'learning_rate': 1.6666666666666667e-05, 'epoch': 2.0}                                                                                                                
{'loss': 4.1261, 'grad_norm': 62.63885498046875, 'learning_rate': 1.25e-05, 'epoch': 2.25}                                                                                                                             
{'loss': 3.5956, 'grad_norm': 67.01679229736328, 'learning_rate': 8.333333333333334e-06, 'epoch': 2.5}                                                                                                                 
{'loss': 3.5619, 'grad_norm': 72.61600494384766, 'learning_rate': 4.166666666666667e-06, 'epoch': 2.75}                                                                                                                
{'loss': 3.6611, 'grad_norm': 67.8575439453125, 'learning_rate': 0.0, 'epoch': 3.0}                                                                                                                                    
{'train_runtime': 56.5843, 'train_samples_per_second': 10.179, 'train_steps_per_second': 0.424, 'train_loss': 6.079538067181905, 'epoch': 3.0}                                                                         
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 24/24 [00:55<00:00,  2.30s/it]
```
`gradient_accumulation_steps=2, per_device_train_batch_size=3` 
```
{'loss': 10.5949, 'grad_norm': 188.28981018066406, 'learning_rate': 4.5833333333333334e-05, 'epoch': 0.25}                                                                                                             
{'loss': 8.2312, 'grad_norm': 87.7468032836914, 'learning_rate': 4.166666666666667e-05, 'epoch': 0.5}                                                                                                                  
{'loss': 7.8988, 'grad_norm': 74.3565444946289, 'learning_rate': 3.7500000000000003e-05, 'epoch': 0.75}                                                                                                                
{'loss': 8.0197, 'grad_norm': 52.91275405883789, 'learning_rate': 3.3333333333333335e-05, 'epoch': 1.0}                                                                                                                
{'loss': 6.0704, 'grad_norm': 51.57986068725586, 'learning_rate': 2.916666666666667e-05, 'epoch': 1.25}                                                                                                                
{'loss': 5.8247, 'grad_norm': 56.05683898925781, 'learning_rate': 2.5e-05, 'epoch': 1.5}                                                                                                                               
{'loss': 5.9932, 'grad_norm': 50.103092193603516, 'learning_rate': 2.0833333333333336e-05, 'epoch': 1.75}                                                                                                              
{'loss': 5.3769, 'grad_norm': 65.85319519042969, 'learning_rate': 1.6666666666666667e-05, 'epoch': 2.0}                                                                                                                
{'loss': 4.1261, 'grad_norm': 62.63907241821289, 'learning_rate': 1.25e-05, 'epoch': 2.25}                                                                                                                             
{'loss': 3.5956, 'grad_norm': 67.01693725585938, 'learning_rate': 8.333333333333334e-06, 'epoch': 2.5}                                                                                                                 
{'loss': 3.562, 'grad_norm': 72.61691284179688, 'learning_rate': 4.166666666666667e-06, 'epoch': 2.75}                                                                                                                 
{'loss': 3.6611, 'grad_norm': 67.85787963867188, 'learning_rate': 0.0, 'epoch': 3.0}                                                                                                                                   
{'train_runtime': 62.5724, 'train_samples_per_second': 9.205, 'train_steps_per_second': 0.384, 'train_loss': 6.079543709754944, 'epoch': 3.0}                                                                          
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 24/24 [01:01<00:00,  2.55s/it]
```

**When I switched to v4.47.0, the first setting remained the same as in v4.45.2:**

`gradient_accumulation_steps=1, per_device_train_batch_size=6` 
```
{'loss': 10.5949, 'grad_norm': 188.29049682617188, 'learning_rate': 4.5833333333333334e-05, 'epoch': 0.25}                                                                                                             
{'loss': 8.2312, 'grad_norm': 87.74673461914062, 'learning_rate': 4.166666666666667e-05, 'epoch': 0.5}                                                                                                                 
{'loss': 7.8988, 'grad_norm': 74.35655212402344, 'learning_rate': 3.7500000000000003e-05, 'epoch': 0.75}                                                                                                               
{'loss': 8.0197, 'grad_norm': 52.913124084472656, 'learning_rate': 3.3333333333333335e-05, 'epoch': 1.0}                                                                                                               
{'loss': 6.0704, 'grad_norm': 51.57878875732422, 'learning_rate': 2.916666666666667e-05, 'epoch': 1.25}                                                                                                                
{'loss': 5.8247, 'grad_norm': 56.05608367919922, 'learning_rate': 2.5e-05, 'epoch': 1.5}                                                                                                                               
{'loss': 5.9932, 'grad_norm': 50.10258483886719, 'learning_rate': 2.0833333333333336e-05, 'epoch': 1.75}                                                                                                               
{'loss': 5.3769, 'grad_norm': 65.85218811035156, 'learning_rate': 1.6666666666666667e-05, 'epoch': 2.0}                                                                                                                
{'loss': 4.1261, 'grad_norm': 62.63885498046875, 'learning_rate': 1.25e-05, 'epoch': 2.25}                                                                                                                             
{'loss': 3.5956, 'grad_norm': 67.01679229736328, 'learning_rate': 8.333333333333334e-06, 'epoch': 2.5}                                                                                                                 
{'loss': 3.5619, 'grad_norm': 72.61600494384766, 'learning_rate': 4.166666666666667e-06, 'epoch': 2.75}                                                                                                                
{'loss': 3.6611, 'grad_norm': 67.8575439453125, 'learning_rate': 0.0, 'epoch': 3.0}                                                                                                                                    
{'train_runtime': 56.4532, 'train_samples_per_second': 10.203, 'train_steps_per_second': 0.425, 'train_loss': 6.079538067181905, 'epoch': 3.0}                                                                         
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 24/24 [00:55<00:00,  2.29s/it]
```

**However, the results changed significantly for larger gradient accumulation steps:**

`gradient_accumulation_steps=2, per_device_train_batch_size=3` 
```
{'loss': 21.1898, 'grad_norm': 376.5796203613281, 'learning_rate': 4.5833333333333334e-05, 'epoch': 0.25}                                                                                                              
{'loss': 16.4623, 'grad_norm': 175.4936065673828, 'learning_rate': 4.166666666666667e-05, 'epoch': 0.5}                                                                                                                
{'loss': 15.7976, 'grad_norm': 148.7130889892578, 'learning_rate': 3.7500000000000003e-05, 'epoch': 0.75}                                                                                                              
{'loss': 16.0394, 'grad_norm': 105.82550811767578, 'learning_rate': 3.3333333333333335e-05, 'epoch': 1.0}                                                                                                              
{'loss': 12.1408, 'grad_norm': 103.15972137451172, 'learning_rate': 2.916666666666667e-05, 'epoch': 1.25}                                                                                                              
{'loss': 11.6494, 'grad_norm': 112.11367797851562, 'learning_rate': 2.5e-05, 'epoch': 1.5}                                                                                                                             
{'loss': 11.9865, 'grad_norm': 100.20618438720703, 'learning_rate': 2.0833333333333336e-05, 'epoch': 1.75}                                                                                                             
{'loss': 10.7537, 'grad_norm': 131.70639038085938, 'learning_rate': 1.6666666666666667e-05, 'epoch': 2.0}                                                                                                              
{'loss': 8.2522, 'grad_norm': 125.27814483642578, 'learning_rate': 1.25e-05, 'epoch': 2.25}                                                                                                                            
{'loss': 7.1912, 'grad_norm': 134.03387451171875, 'learning_rate': 8.333333333333334e-06, 'epoch': 2.5}                                                                                                                
{'loss': 7.1239, 'grad_norm': 145.23382568359375, 'learning_rate': 4.166666666666667e-06, 'epoch': 2.75}                                                                                                               
{'loss': 7.3222, 'grad_norm': 135.71575927734375, 'learning_rate': 0.0, 'epoch': 3.0}                                                                                                                                  
{'train_runtime': 64.6658, 'train_samples_per_second': 8.907, 'train_steps_per_second': 0.371, 'train_loss': 12.159087419509888, 'epoch': 3.0}                                                                         
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 24/24 [01:03<00:00,  2.63s/it]
```

**After upgrading to the latest v4.47.1 and higher (including your new commit), the loss is the same across different grad accum steps**
`gradient_accumulation_steps=1, per_device_train_batch_size=6` 
```
{'loss': 2.119, 'grad_norm': 37.658103942871094, 'learning_rate': 4.5833333333333334e-05, 'epoch': 0.25}                                                                                                               
{'loss': 1.6462, 'grad_norm': 17.549333572387695, 'learning_rate': 4.166666666666667e-05, 'epoch': 0.5}                                                                                                                
{'loss': 1.5798, 'grad_norm': 14.87132740020752, 'learning_rate': 3.7500000000000003e-05, 'epoch': 0.75}                                                                                                               
{'loss': 1.6039, 'grad_norm': 10.582618713378906, 'learning_rate': 3.3333333333333335e-05, 'epoch': 1.0}                                                                                                               
{'loss': 1.2141, 'grad_norm': 10.315776824951172, 'learning_rate': 2.916666666666667e-05, 'epoch': 1.25}                                                                                                               
{'loss': 1.1649, 'grad_norm': 11.211206436157227, 'learning_rate': 2.5e-05, 'epoch': 1.5}                                                                                                                              
{'loss': 1.1986, 'grad_norm': 10.020527839660645, 'learning_rate': 2.0833333333333336e-05, 'epoch': 1.75}                                                                                                              
{'loss': 1.0754, 'grad_norm': 13.170440673828125, 'learning_rate': 1.6666666666666667e-05, 'epoch': 2.0}                                                                                                               
{'loss': 0.8252, 'grad_norm': 12.527780532836914, 'learning_rate': 1.25e-05, 'epoch': 2.25}                                                                                                                            
{'loss': 0.7191, 'grad_norm': 13.403361320495605, 'learning_rate': 8.333333333333334e-06, 'epoch': 2.5}                                                                                                                
{'loss': 0.7124, 'grad_norm': 14.523214340209961, 'learning_rate': 4.166666666666667e-06, 'epoch': 2.75}                                                                                                               
{'loss': 0.7322, 'grad_norm': 13.57152271270752, 'learning_rate': 0.0, 'epoch': 3.0}                                                                                                                                   
{'train_runtime': 58.3558, 'train_samples_per_second': 9.87, 'train_steps_per_second': 0.411, 'train_loss': 1.2159071018298466, 'epoch': 3.0}                                                                          
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 24/24 [00:56<00:00,  2.37s/it]
``` 
`gradient_accumulation_steps=2, per_device_train_batch_size=3` 
```
{'loss': 2.119, 'grad_norm': 37.657958984375, 'learning_rate': 4.5833333333333334e-05, 'epoch': 0.25}                                                                                                                  
{'loss': 1.6462, 'grad_norm': 17.54932403564453, 'learning_rate': 4.166666666666667e-05, 'epoch': 0.5}                                                                                                                 
{'loss': 1.5798, 'grad_norm': 14.871335983276367, 'learning_rate': 3.7500000000000003e-05, 'epoch': 0.75}                                                                                                              
{'loss': 1.6039, 'grad_norm': 10.582595825195312, 'learning_rate': 3.3333333333333335e-05, 'epoch': 1.0}                                                                                                               
{'loss': 1.2141, 'grad_norm': 10.315934181213379, 'learning_rate': 2.916666666666667e-05, 'epoch': 1.25}                                                                                                               
{'loss': 1.1649, 'grad_norm': 11.211311340332031, 'learning_rate': 2.5e-05, 'epoch': 1.5}                                                                                                                              
{'loss': 1.1986, 'grad_norm': 10.020581245422363, 'learning_rate': 2.0833333333333336e-05, 'epoch': 1.75}                                                                                                              
{'loss': 1.0754, 'grad_norm': 13.170597076416016, 'learning_rate': 1.6666666666666667e-05, 'epoch': 2.0}                                                                                                               
{'loss': 0.8252, 'grad_norm': 12.527819633483887, 'learning_rate': 1.25e-05, 'epoch': 2.25}                                                                                                                            
{'loss': 0.7191, 'grad_norm': 13.403372764587402, 'learning_rate': 8.333333333333334e-06, 'epoch': 2.5}                                                                                                                
{'loss': 0.7124, 'grad_norm': 14.523296356201172, 'learning_rate': 4.166666666666667e-06, 'epoch': 2.75}                                                                                                               
{'loss': 0.7322, 'grad_norm': 13.571563720703125, 'learning_rate': 0.0, 'epoch': 3.0}                                                                                                                                  
{'train_runtime': 64.584, 'train_samples_per_second': 8.919, 'train_steps_per_second': 0.372, 'train_loss': 1.2159080157677333, 'epoch': 3.0}                                                                          
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 24/24 [01:03<00:00,  2.63s/it]
```

**But is quite different compared to v4.45.2**

I'm not entirely sure where the questions are arising, so I'm diving into the details of the Trainer. I'll keep you updated as soon as I have more information.

Encountered a similar issue when running the above script.
`transformers=4.47.0`:
```
{'loss': 22.5655, 'grad_norm': 495.53765869140625, 'learning_rate': 4.791666666666667e-05, 'epoch': 0.12}                                                                                                                               
{'loss': 16.8421, 'grad_norm': 284.22906494140625, 'learning_rate': 4.5833333333333334e-05, 'epoch': 0.25}                                                                                                                              
{'loss': 16.2376, 'grad_norm': 177.2020721435547, 'learning_rate': 4.375e-05, 'epoch': 0.38}                                                                                                                                            
{'loss': 15.8786, 'grad_norm': 199.05421447753906, 'learning_rate': 4.166666666666667e-05, 'epoch': 0.5}                                                                                                                                
{'loss': 16.4592, 'grad_norm': 178.1647491455078, 'learning_rate': 3.958333333333333e-05, 'epoch': 0.62}                                                                                                                                
{'loss': 15.2518, 'grad_norm': 179.19195556640625, 'learning_rate': 3.7500000000000003e-05, 'epoch': 0.75}                                                                                                                              
{'loss': 15.9623, 'grad_norm': 126.81049346923828, 'learning_rate': 3.541666666666667e-05, 'epoch': 0.88}                                                                                                                               
{'loss': 15.7668, 'grad_norm': 120.61547088623047, 'learning_rate': 3.3333333333333335e-05, 'epoch': 1.0}
```
`transformers=4.44.2`:
```
{'loss': 11.2828, 'grad_norm': 247.76889038085938, 'learning_rate': 4.791666666666667e-05, 'epoch': 0.12}                                                                                                                               
{'loss': 8.421, 'grad_norm': 142.11448669433594, 'learning_rate': 4.5833333333333334e-05, 'epoch': 0.25}                                                                                                                                
{'loss': 8.1188, 'grad_norm': 88.60110473632812, 'learning_rate': 4.375e-05, 'epoch': 0.38}                                                                                                                                             
{'loss': 7.9393, 'grad_norm': 99.52731323242188, 'learning_rate': 4.166666666666667e-05, 'epoch': 0.5}                                                                                                                                  
{'loss': 8.2296, 'grad_norm': 89.08261108398438, 'learning_rate': 3.958333333333333e-05, 'epoch': 0.62}                                                                                                                                 
{'loss': 7.6259, 'grad_norm': 89.5963134765625, 'learning_rate': 3.7500000000000003e-05, 'epoch': 0.75}                                                                                                                                 
{'loss': 7.9811, 'grad_norm': 63.40541458129883, 'learning_rate': 3.541666666666667e-05, 'epoch': 0.88}                                                                                                                                 
{'loss': 7.8834, 'grad_norm': 60.307289123535156, 'learning_rate': 3.3333333333333335e-05, 'epoch': 1.0}  
```


Yes, probably related to the grad accum issue see #34198. As mentionned above, I think that everything works as expected with the dev version