fix(trainer): Correct loss scaling for incomplete gradient accumulation steps
Author: hutaiHang
Pull request: yes

# What does this PR do?

This PR addresses an issue where the loss scaling during gradient accumulation is incorrect for the final optimizer step of an epoch if the total number of batches is not perfectly divisible by `gradient_accumulation_steps`.

Currently, the loss for each micro-batch is always divided by the configured `args.gradient_accumulation_steps`. This leads to the accumulated loss for the final, incomplete cycle being scaled down too much, resulting in an improperly small gradient update for that step.

This fix resolves the issue by dynamically tracking the number of micro-batches processed in each accumulation cycle and using this actual count for loss scaling.

**The changes are as follows:**
1.  In the `_inner_training_loop`, a new instance variable `self.cur_gradient_accumulation_steps` is introduced. It is updated at the start of each optimizer step with the actual number of batches being processed (i.e., `len(batch_samples)`).
2.  In the `training_step` method, the loss scaling logic now uses this dynamic `self.cur_gradient_accumulation_steps` value instead of the fixed `self.args.gradient_accumulation_steps`.

This ensures that the loss is correctly averaged over the number of batches that actually contributed to the gradient accumulation, regardless of whether the cycle was complete or not. This change has no new dependencies.

Fixes #38837

## 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. (This PR is for issue #38837)
- [ ] 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?


## Who can review?

- trainer: @SunMarc  @zach-huggingface @qgallouedec


Comments:
Hi @SunMarc,

Thanks for your guidance in issue #38837. As you suggested, I've created this Pull Request to address the gradient accumulation scaling issue.

All CI checks have now passed. Could you please take a look when you have a moment?

Thank you!

Hi @qgallouedec, thanks for the feedback! I've just pushed the requested changes. Could you please take another look when you have a chance?

Looks good, let's see if the CI is happy

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

@qgallouedec All checks have passed, thanks for the review. 😊

Note for myself: it will silently break the grad accumulation in GRPOtrainer because this trainer oversample from the dataloader. I'll have to find a way to solve that 

Hi @hutaiHang

There is a question. What if the last batch only has 1 data ? This only one data will contribute to all of the gradient, which means that the model will update to a direction that is not **general** (Cause we wanna batch size to be *big* to make each update to be **general**). That will cause training not stable.

I thinks we should keep the loss divided by `self.args.gradient_accumulation_steps` before backward process. At the end return the loss (right here loss is just for log) that is sacled to the correct one.

```python
# trainer.py
        if self.use_apex:
            with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                scaled_loss.backward()
        else:
            # 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

            # Turning off loss scaling w.r.t. gradient accumulation when DeepSpeed is enabled
            # https://github.com/huggingface/transformers/pull/35808
            if self.accelerator.distributed_type == DistributedType.DEEPSPEED:
                kwargs["scale_wrt_gas"] = False

            self.accelerator.backward(loss, **kwargs)

            rescale_loss = loss.detach() * self.args.gradient_accumulation_steps / self.current_gradient_accumulation_steps

            return rescale_loss
```


Hi @kaln27 

Thank you for bringing up this excellent point about training dynamics! It's a very important consideration.

You are absolutely right to be concerned about the potential instability if an optimizer step is updated based on a very small final batch (e.g., with a single sample). The gradient from such a small batch can indeed be noisy.

However, I believe the proposed solution in the PR is the correct way to fix the underlying mathematical bug in the `Trainer`, and the stability concern should be addressed at the data loading level. Let me break down my reasoning:

**1. Gradient Magnitude Correctness (The goal of this PR)**

The purpose of gradient accumulation is to simulate a larger batch size. The final gradient update should have a magnitude that is the *average* of the gradients from the processed micro-batches.

- If we process **4** batches, the loss is summed and then divided by **4** before `backward()`.
- If we process **1** batch, the loss should be divided by **1** before `backward()`.

This PR ensures this mathematical correctness. If we were to always divide by `self.args.gradient_accumulation_steps` as you suggested, the gradient for the final, incomplete step would be artificially suppressed (e.g., 1/4 of its correct magnitude). This would mean the model barely learns from those last few samples, which is also undesirable.

**2. Training Stability (The concern you raised)**

The question of whether one *wants* to perform an update based on a noisy, small batch is a matter of training strategy.

The standard and recommended way to handle this in `transformers` is to use the `dataloader_drop_last=True` argument in `TrainingArguments`. This tells the `DataLoader` to simply discard the last, incomplete batch, ensuring that every single optimizer step sees a full batch of data.

**Conclusion:**

This PR is focused on making the `Trainer`'s default behavior mathematically correct. Users who are concerned about the stability implications of small final batches already have a direct and explicit tool (`dataloader_drop_last`) to manage this. The `Trainer` itself should not silently suppress gradients as a "feature," as that hides the true contribution of the data.

I hope this clarifies my approach. Let me know what you think!