Fix Gradient Accumulation issue
Author: ArthurZucker
Pull request: yes

# What does this PR do?
First draft

End goal is to make it easy for anyone to:
- change the loss for his model 
- contribute a new loss for a model (like vision model, ENCODEC etc)
- allow passing arbitrary kwargs, interfacing 

TODO:
- [ ] Fix deformable detr loss computation

Comments:
Continuing discussion from https://github.com/huggingface/transformers/pull/34198 - @muellerzr is correct if we truncate and partition all sequences to block sizes of 128 for Wikitext, the GA fix and without the GA fix should nearly be equivalent. I can also repro this:
![image](https://github.com/user-attachments/assets/8fde0302-6059-4d28-aae0-d96ef0f78a34)

But unfortunately if we don't pad it and use varying sequence lengths, the loss for `bsz = 1, ga = 32` actually seems to go haywire:
![image](https://github.com/user-attachments/assets/7dda264d-97df-4770-803d-05d5db14a056)

If you plot the training loss for `bsz = 1, ga = 32` with and without the fix (blue and yellow), and also `bsz = 32, ga = 1` (red) under a log scale, we get:
![image](https://github.com/user-attachments/assets/ad71b3d4-b0c5-4ffa-bce2-0277beb0b725)

And if we plot them on the left and right axes:
![image](https://github.com/user-attachments/assets/ab4f7588-4fec-4254-aeaf-ab6e1595660c)

All experiments here: https://docs.google.com/spreadsheets/d/1RUiVuFNfnl9eBAa3JhvkKb0hm20m4NqnUO-OWDPpNos/edit?usp=sharing - also the `bsz = 32` varying sequence lengths unfortunately OOMed on a A100 40GB, so only 50% of the training went through.

TLDR: The fix works as expected, making GA and full batch training equivalent. The effect on truncated or merged sequence lengths is less pronounced, but varying sequence lengths makes the issue become very obvious

Can also confirm, and confirm the fix in #34198 solves it

![W B Chart 10_17_2024, 5_45_37 AM](https://github.com/user-attachments/assets/7d97354c-44e8-462c-9f00-b369b844a927)



![W B Chart 10_17_2024, 5_45_37 AM(1)](https://github.com/user-attachments/assets/104e0727-55cc-4c06-8dc4-5fe39e0e74e1)


Quick question: Is Distributed Data Parallel (DDP) training affected by this issue in the same way?

@Hannibal046 If you used gradient accumulation along with DDP, yes.

Consider the following scenario with two samples of different lengths (wo gradient acc):

1. Sample A: 10 tokens
2. Sample B: 5 tokens

- Case 1: Single GPU (batch_size=2)

```python
loss = average_loss(10 tokens + 5 tokens)
gradient_1 = backprop(loss)  # Based on average of 15 tokens
```
- Case 2: Distributed Training (2 GPUs, local_batch_size=1)
```python
# GPU 1
loss_1 = average_loss(10 tokens)
gradient_1 = backprop(loss_1)

# GPU 2
loss_2 = average_loss(5 tokens)
gradient_2 = backprop(loss_2)

final_gradient = 0.5 * gradient_1 + 0.5 * gradient_2
```
so `gradient_1 != final_gradient`.


The issue fixed in this PR is about gradient accumulation, not DDP.
What you explained about is orthogonal with gradient accumulation, and may be intended behavior.

After using this fix, the following problem occurred:
the grad norm and loss (custom loss from trl cpo trainer) increased dramatically

![image](https://github.com/user-attachments/assets/64251d7a-10f3-4193-91eb-529360b9930f)
green：llama34b-deepspeed-zero3-8gpus-fullparams-bs4-grad8-wo_fix
orange: llama34b-deepspeed-zero3-8gpus-fullparams-bs4-grad8-wo_fix


> After using this fix, the following problem occurred: the grad norm and loss (custom loss from trl cpo trainer) increased dramatically
> 
> ![image](https://private-user-images.githubusercontent.com/738834/378013278-64251d7a-10f3-4193-91eb-529360b9930f.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MjkzMDE2NDUsIm5iZiI6MTcyOTMwMTM0NSwicGF0aCI6Ii83Mzg4MzQvMzc4MDEzMjc4LTY0MjUxZDdhLTEwZjMtNDE5My05MWViLTUyOTM2MGI5OTMwZi5wbmc_WC1BbXotQWxnb3JpdGhtPUFXUzQtSE1BQy1TSEEyNTYmWC1BbXotQ3JlZGVudGlhbD1BS0lBVkNPRFlMU0E1M1BRSzRaQSUyRjIwMjQxMDE5JTJGdXMtZWFzdC0xJTJGczMlMkZhd3M0X3JlcXVlc3QmWC1BbXotRGF0ZT0yMDI0MTAxOVQwMTI5MDVaJlgtQW16LUV4cGlyZXM9MzAwJlgtQW16LVNpZ25hdHVyZT04ZWVhYjkwYWQ2NzY2NjZkN2IzYjRiY2EyMzM5MWU5NDk0ZDI1ZjRkYzllYWMzM2E3OGUyMjNmNmIzODM4NjcwJlgtQW16LVNpZ25lZEhlYWRlcnM9aG9zdCJ9.wZ_6eX1qKsC0rQhCCXzGyTaJJTdaZ5H7LVgktlQPlNM) green：llama34b-deepspeed-zero3-8gpus-fullparams-bs4-grad8-wo_fix orange: llama34b-deepspeed-zero3-8gpus-fullparams-bs4-grad8-wo_fix

Can confirm the same issue

I think the fix isnt complete, another fix is needed to make it normal. See https://github.com/hiyouga/LLaMA-Factory/issues/5747#issuecomment-2425488226

> I think the fix isnt complete, another fix is needed to make it normal. See [hiyouga/LLaMA-Factory#5747 (comment)](https://github.com/hiyouga/LLaMA-Factory/issues/5747#issuecomment-2425488226)

Hi @ArthurZucker , would you mind check the further fix?

Just so we can note this on the right issue please put comments here: https://github.com/huggingface/transformers/pull/34198

This PR only modified the loss functions, that one is actually what made modifications in the trainer

Note though: I would likely not expect old LR's to behave similar. I'd expect old LR's to behave like as though GA wasn't used. 

As [hiyouga/LLaMA-Factory#5747 (comment)](https://github.com/hiyouga/LLaMA-Factory/issues/5747#issuecomment-2425488226) suggested, maybe a few changes will make this fix work again?

> > After using this fix, the following problem occurred: the grad norm and loss (custom loss from trl cpo trainer) increased dramatically
> > ![image](https://private-user-images.githubusercontent.com/738834/378013278-64251d7a-10f3-4193-91eb-529360b9930f.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MjkzMDE2NDUsIm5iZiI6MTcyOTMwMTM0NSwicGF0aCI6Ii83Mzg4MzQvMzc4MDEzMjc4LTY0MjUxZDdhLTEwZjMtNDE5My05MWViLTUyOTM2MGI5OTMwZi5wbmc_WC1BbXotQWxnb3JpdGhtPUFXUzQtSE1BQy1TSEEyNTYmWC1BbXotQ3JlZGVudGlhbD1BS0lBVkNPRFlMU0E1M1BRSzRaQSUyRjIwMjQxMDE5JTJGdXMtZWFzdC0xJTJGczMlMkZhd3M0X3JlcXVlc3QmWC1BbXotRGF0ZT0yMDI0MTAxOVQwMTI5MDVaJlgtQW16LUV4cGlyZXM9MzAwJlgtQW16LVNpZ25hdHVyZT04ZWVhYjkwYWQ2NzY2NjZkN2IzYjRiY2EyMzM5MWU5NDk0ZDI1ZjRkYzllYWMzM2E3OGUyMjNmNmIzODM4NjcwJlgtQW16LVNpZ25lZEhlYWRlcnM9aG9zdCJ9.wZ_6eX1qKsC0rQhCCXzGyTaJJTdaZ5H7LVgktlQPlNM) green：llama34b-deepspeed-zero3-8gpus-fullparams-bs4-grad8-wo_fix orange: llama34b-deepspeed-zero3-8gpus-fullparams-bs4-grad8-wo_fix
> 
> Can confirm the same issue
same. encountered when fine-tuning using qwen2.5. 

Fix is being worked on here: https://github.com/huggingface/transformers/pull/34283

Oh my....

Hello all, I could not follow the whole context of the PR here, but I am guessing a fix is underway regarding the loss becoming huge due to [this code](https://github.com/huggingface/transformers/blob/1d063793318b20654ebb850f48f43e0a247ab7bb/src/transformers/trainer.py#L3605). The code scales the loss by gradient accumulation step in the latest trainer (not in previous version). The snippet is included in a single training iteration step, for which I think scaling the loss by grad_acc_step is not appropriate, if the loss is not de-scaled by grad_acc_step inside the accelerator.backward call in the following line.

Does the latest accelerator backward de-scale by grad_acc_step (or is being fixed to do so)? If so, maybe the output loss of training_step() should be de-scaled one more time by grad_acc_step?
For now, is it correct to scale down the learning rate by the grad_acc_step for the original functionality?

Yes accelerate de-scales the loss when calling `backward()` which is why we do so. (Accelerate has always done this)

@muellerzr Thank you for the reply. I have just one more question to clear things up.

You've mentioned that accelerate had always internally de-scaled the loss by gradient accumulation steps.
However, the previous trainer.py implementations do not scale (scale up) the loss by grad_acc_steps before accelerator.backward. 
If accelerator.backward had been internally de-scaling the loss by grad_acc_steps, does that mean the previous trainer had been training a model on 1/grad_acc_steps * sum(loss) instead of sum(loss) when grad_acc_steps > 1?


FYI: As also noted by @cphillippi-stripe, the `@property` `PreTrainedModel.loss_function` that was added by this PR breaks all software packages that have models which (1) inherit from `PreTrainedModel` while (2) setting a `loss_function` beforehand. 

I maintain [such a package](https://github.com/bauwenst/ArchIt/blob/4018ead262e4a6551bb6f1ff21bf6f52af6fce70/src/archit/instantiation/abstracts.py#L234-L236), and indeed, training is now broken and my IDE complains that `loss_function` cannot be set. 

![image](https://github.com/user-attachments/assets/94b1a179-e24f-4862-8177-944c104fcc0d)

The default behaviour of setting the loss to *causal LM* loss when no `loss_type` is defined in the config seems **highly** unwarranted to me. It is assumed now that either (1) models define this brand new `loss_type` field in the config and otherwise (2) the loss type can be inferred from the `ForXYZ` suffix:

https://github.com/huggingface/transformers/blob/63766abe362b930522dd073c9173499ba5fde02a/src/transformers/modeling_utils.py#L5110-L5120

Yet, this fails to take into account the following cases where `transformers` is used as a dependency:
1. Model classes defined for *tasks that don't have a known auto class* in `transformers`. For example: if your package defines a `ModelForDependencyParsing`, the above regex search will not find its loss and default to causal loss, which is wrong.
 
```python
  File ".../ArchIt/src/archit/instantiation/tasks.py", line 328, in computeLoss
    return self.loss_function(arc_scores, arc_labels) + self.loss_function(rel_scores, rel_labels)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: ForCausalLMLoss() missing 1 required positional argument: 'vocab_size'
```

2. Model classes that have a known task, but it's not in the name (e.g. `ModelForSentenceClassifying`). Unless the config is from after October 2024, the config does not define its loss type.

3. Loss is not inherent to a `PreTrainedModel`. There are plenty of papers that start from the model weights of one model and then change its loss function for continued pretraining (take [this one](https://aclanthology.org/2024.acl-long.136/) as a recent example).

TL;DR: Seems like a design flaw to make `loss_function` a `@property` and seems like a mistake to overwrite it to `ForCausalLMLoss` for all models that don't fit the very narrow range of tasks supported by `transformers`.

> FYI: As also noted by @cphillippi-stripe, the `@property` `PreTrainedModel.loss_function` that was added by this PR breaks all software packages that have models which (1) inherit from `PreTrainedModel` while (2) setting a `loss_function` beforehand.
> 
> I maintain [such a package](https://github.com/bauwenst/ArchIt/blob/4018ead262e4a6551bb6f1ff21bf6f52af6fce70/src/archit/instantiation/abstracts.py#L234-L236), and indeed, training is now broken and my IDE complains that `loss_function` cannot be set.
> 
> ![image](https://private-user-images.githubusercontent.com/145220868/395201187-94b1a179-e24f-4862-8177-944c104fcc0d.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MzQwMzk3MzMsIm5iZiI6MTczNDAzOTQzMywicGF0aCI6Ii8xNDUyMjA4NjgvMzk1MjAxMTg3LTk0YjFhMTc5LWUyNGYtNDg2Mi04MTc3LTk0NGMxMDRmY2MwZC5wbmc_WC1BbXotQWxnb3JpdGhtPUFXUzQtSE1BQy1TSEEyNTYmWC1BbXotQ3JlZGVudGlhbD1BS0lBVkNPRFlMU0E1M1BRSzRaQSUyRjIwMjQxMjEyJTJGdXMtZWFzdC0xJTJGczMlMkZhd3M0X3JlcXVlc3QmWC1BbXotRGF0ZT0yMDI0MTIxMlQyMTM3MTNaJlgtQW16LUV4cGlyZXM9MzAwJlgtQW16LVNpZ25hdHVyZT1lMjRmZWU1OWNkMDJiN2Q0ZGI3YmI5YTc5MTFiMzIxYzkzNmYwNTZhYzE3NzYyNTZlMDhkNzUyNjViMzE3ZDAwJlgtQW16LVNpZ25lZEhlYWRlcnM9aG9zdCJ9.tDkfmLNbPhn8zFCYx3vh2KIaykrCzBXoBQi59duoHsA)
> 
> The default behaviour of setting the loss to _causal LM_ loss when no `loss_type` is defined in the config seems **highly** unwarranted to me. It is assumed now that either models define this brand new `loss_type` field in the config and that otherwise the loss type can be inferred from the `ForXYZ` suffix:
> 
> https://github.com/huggingface/transformers/blob/63766abe362b930522dd073c9173499ba5fde02a/src/transformers/modeling_utils.py#L5110-L5120
> 
> Yet, this fails to take into account the following cases where `transformers` is used as a dependency:
> 
> 1. Model classes defined for _tasks that don't have a known auto class_ in `transformers`. For example: if your package defines a `ModelForDependencyParsing`, the above regex search will not find its loss and default to causal loss, which is wrong.
> 
> ```python
>   File ".../ArchIt/src/archit/instantiation/tasks.py", line 328, in computeLoss
>     return self.loss_function(arc_scores, arc_labels) + self.loss_function(rel_scores, rel_labels)
>            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
> TypeError: ForCausalLMLoss() missing 1 required positional argument: 'vocab_size'
> ```
> 
> 2. Model classes that have a known task, but it's not in the name (e.g. `ModelForSentenceClassifying`). Unless the config is from pre-October 2024, the config does not define its loss type.
> 3. Loss is not inherent to a `PreTrainedModel`. There are plenty of papers that start from the model weights of one model and then change its loss function for continued pretraining (take [this one](https://aclanthology.org/2024.acl-long.136/) as a recent example).
> 
> TL;DR: Seems like a design flaw to make `loss_function` a `@property` and seems like a mistake to overwrite it to `ForCausalLMLoss` for all models that don't fit the very narrow range of tasks supported by `transformers`.

Not a fix but a temporary workaround for me was this patch after loading a mdoel:
```
model = AutoModelForSequenceClassification.from_pretrained(...)

old_loss_function = model.loss_function
@lru_cache
def get_loss_function(self):
    def fixed_loss(labels, pooled_logits, config, **kwargs):
        labels = labels.to(pooled_logits.device)
        return old_loss_function(labels, pooled_logits, config, **kwargs)
    return fixed_loss
    
model.__class__.loss_function = property(get_loss_function)
```

Also, @ArthurZucker, I believe this fix confirms removing the `labels = labels.to(pooled_logits.device)` line breaks the multi-gpu setup. I would be very surprised if only Mistral is affected here given similar lines were removed from other models. Initially, I tried patching this at the `PreTrainedModel.loss_function` level, but for whatever reason the loaded models don't seem to pick up the change, and I'm not sure why.

Feedback taken, will open a PR to get to something better. It's hard to take everything into account, and support important to get your feedbacks! 🤗 
Will merge the fix to loss on multi-GPU (I don't think we test all tasks for training and that is a flaw in our test environnement for sure! ) cc @ydshieh as for a TODO 😉