Gemma3 Gradient Accumulation loss
Author: winglian
Pull request: no
Labels: name: bug

### System Info

transformers==4.50.3

### Who can help?

@osanseviero @ArthurZucker 


### Information

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

### Tasks

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

### Reproduction


repro script with SFTTrainer here: https://gist.github.com/winglian/569924fe154824c8ce148f6e185cd4cd

When using `gradient_accumulation_steps=1`, the magnitude of the loss for gemma3 family of models is approximately correct ~1 or less depending on the model size. When changing `gradient_accumulation_steps=4`, the loss is scaled up by a factor of the `gradient_accumulation_steps` and in this case is ~8-10.

Digging into the [Gemma3TextModel.forward method ](https://github.com/huggingface/transformers/blob/800510c67bfc5cedd0bb7635648a07f39719be43/src/transformers/models/gemma3/modular_gemma3.py#L599), we see that it doesn't accept the typical `LossKwargs` to handle gradient accumulation. 

Compare the signature to [llama](https://github.com/huggingface/transformers/blob/800510c67bfc5cedd0bb7635648a07f39719be43/src/transformers/models/llama/modeling_llama.py#L809)  where it's defined as
```python
class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ...
```

### Expected behavior

loss should be ~1.0 regardless of gradient accumulation steps.

Comments:
I made a patch in Axolotl to address this with charts before/after https://github.com/axolotl-ai-cloud/axolotl/pull/2473

Lines of interest: https://github.com/axolotl-ai-cloud/axolotl/pull/2473/files#diff-9feac1e1ed5234c4effcd90c4d7ddf62264add2863f588d6c71be6e4a8cf15e9R190-R218

However, I'm not completely sure my patch is done correctly. Would love for feedback on it before I make the PR to here too!

cc @sunmarc @muellerzr as well

This should be fixed in this [PR](https://github.com/huggingface/transformers/pull/37208).  As you saw @NanoCode012, in the code, they are filetring the logits/labels so I decided to simply not use `num_items_in_batch` to calculate the loss. Otherwise, the loss won't be correctly calculated for one of the cases.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

> This should be fixed in this [PR](https://github.com/huggingface/transformers/pull/37208). As you saw [@NanoCode012](https://github.com/NanoCode012), in the code, they are filetring the logits/labels so I decided to simply not use `num_items_in_batch` to calculate the loss. Otherwise, the loss won't be correctly calculated for one of the cases.

Hi, I am still a bit confused about the cause. Could you explain why the loss becomes **proportional to** `grad_accum_steps`?

In my understanding, the bug causes the loss to be computed as `mean(loss_i / num_items_i)` instead of the correct `sum(loss_i) / sum(num_items_i)`. While this creates a discrepancy when sequence lengths differ significantly within a batch, I don't see how it leads to proportional scaling. Am I missing something?

I am asking because I found that reproducing this "bug" actually aligns the performance with baselines, whereas fixing it seems to degrade performance of my network. : (