The loss and gradient explosion caused by the trainer
Author: rangehow
Pull request: no
Labels: name: bug


- `transformers` version: 4.49.0
- Platform: Linux-4.18.0-147.mt20200626.413.el8_1.x86_64-x86_64-with-glibc2.17
- Python version: 3.12.9
- Huggingface_hub version: 0.26.3
- Safetensors version: 0.4.5
- Accelerate version: 1.4.0
- Accelerate config:    not found
- DeepSpeed version: 0.15.4
- PyTorch version (GPU?): 2.5.1+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: yes
- Using GPU in script?: yes
- GPU type: NVIDIA Graphics Device

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [x] My own task or dataset (give details below)

### Reproduction

To observe the loss for each device in DDP, I overrode the trainer's compute_loss as follows:
```python
from transformers.loss.loss_utils import ForCausalLMLoss

class NTPTrainer(Trainer):

    
    def compute_loss(
        self,
        model: torch.nn.Module,
        inputs: Dict[str, Union[torch.Tensor, Any]],
        return_outputs: bool = False,
        num_items_in_batch=None,
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, Dict]]:
        
        labels = inputs.pop("labels", None)
        outputs = model(**inputs)

        logits = outputs.logits
        loss_kwargs = {
            "logits": logits,
            "labels": labels,
            "vocab_size": self.model.config.vocab_size 
        }


        loss = ForCausalLMLoss(**loss_kwargs)
        print(loss)

        return (loss, outputs) if return_outputs else loss


    def _get_train_sampler(self):

        return SequentialSampler(self.train_dataset)

```

I compared the loss using my overridden trainer(left) and the default trainer (right). The loss and gradients have both increased by dozens of times, but the way I implemented them is consistent with the loss calculation method provided in the official modeling file.

<img width="822" alt="Image" src="https://github.com/user-attachments/assets/d0426809-0b24-4c11-89c7-0b6c9117a203" />

I am using 8 GPUs, with gradient accumulation set to 16. This print the loss for each device that is output by the trainer.

<img width="449" alt="Image" src="https://github.com/user-attachments/assets/f666a056-3a02-493d-aa6d-9807a38bb7c5" />

The loss for each device is in the single digits, so I'm wondering if gradient accumulation or DDP is not correctly averaging?

### Expected behavior

Aligned with the loss of the standard trainer.

Comments:
cc @zach-huggingface @SunMarc 

The crux of the issue is that passing num_items_in_batch to ForCausalLMLoss triggers token batch loss (sum reduction); otherwise, it's batch mean loss.

Thanks for the detailed report @rangehow ! To not have this issue in the future, maybe it would be better to give more details on what users needs to be careful at when overwriting `compute_loss` ? If you are down to, would you like to submit a small PR ? We also have `compute_loss_func` that you can pass if you don't want to overwrite `compute_loss`.


The issue arises here: if `num_items_in_batch` is passed, it triggers sum reduction; otherwise, it triggers mean reduction.

[https://github.com/huggingface/transformers/blob/786d9c5ed920a099573ea7b6dbf265f1aeb32fc0/src/transformers/loss/loss_utils.py#L27-L38](https://github.com/huggingface/transformers/blob/786d9c5ed920a099573ea7b6dbf265f1aeb32fc0/src/transformers/loss/loss_utils.py#L27-L38)

This shouldn't lead to such a significant gradient difference as expected. The problem lies with `num_items_in_batch`. If I remember correctly, this parameter was introduced in a relatively recent version, around 3.8.x? During gradient accumulation, this parameter represents the total number of tokens in the accumulated batch, thus dividing each loss by a number much larger than that used in mean reduction.

However, this sum behavior seems correct. So the question is, if the user wishes to customize the loss and use mean reduction, how should they do it? (Theoretically, should they add a normalization to the large batch formed by gradient accumulation and DDP?）I'm still a newbie and don't know the answer to this question.

It's indeed not recommended to override `compute_loss` anymore, due to the following reason:

[https://github.com/huggingface/transformers/blob/786d9c5ed920a099573ea7b6dbf265f1aeb32fc0/src/transformers/trainer.py#L3824-L3829](https://github.com/huggingface/transformers/blob/786d9c5ed920a099573ea7b6dbf265f1aeb32fc0/src/transformers/trainer.py#L3824-L3829)

There is a more complex behavior added here, which seems to attempt to expand the token batch loss's batch from a narrow gradient accumulation batch to an equivalent batch formed by gradient accumulation and DDP combined. Therefore, passing compute_loss_func is a more granular and seemingly safer behavior.


> This shouldn't lead to such a significant gradient difference as expected. The problem lies with num_items_in_batch. If I remember correctly, this parameter was introduced in a relatively recent version, around 3.8.x? During gradient accumulation, this parameter represents the total number of tokens in the accumulated batch, thus dividing each loss by a number much larger than that used in mean reduction.

For mean reduction, actually, we divide again by the gradient accumulation size so it shouldn't be that different. However, in your case, you were doing mean reduction but it was not dividing by the grad_acc size since you overwrote `compute_loss`. This is probably why there is such a big difference. 

> There is a more complex behavior added here, which seems to attempt to expand the token batch loss's batch from a narrow gradient accumulation batch to an equivalent batch formed by gradient accumulation and DDP combined. Therefore, passing compute_loss_func is a more granular and seemingly safer behavior.

`average_tokens_across_devices` default is `False` so it should not be a big issue to overwrite the `compute_loss`. The only big issue is with `num_items_in_batch`

> However, this sum behavior seems correct. So the question is, if the user wishes to customize the loss and use mean reduction, how should they do it? (Theoretically, should they add a normalization to the large batch formed by gradient accumulation and DDP?）I'm still a newbie and don't know the answer to this question.

If they wish to use mean reduction, the only way right now is to overwrite `accepts_loss_kwargs` attribute in the model. Maybe it can be helpful to introduce an argument in trainer to control this behavior, the same as `average_tokens_across_devices` ?