Fix step shifting when accumulate gradient
Author: kibitzing
Pull request: yes

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes #33671


## 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.
- [ ] 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?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@muellerzr and @SunMarc

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

Hello @SunMarc ,

Thank you for taking the time to review my PR! I appreciate your suggestion, it's a great idea and would indeed simplify the code.

However, after investigating and experimenting with this, I found that `self.gradient_state.end_of_dataloader` is always `False` in my case, preventing the step of `self.accelerator` from being reset. 
This happens because:
*  `self.gradient_state.end_of_dataloader` depends on `self.gradient_state.active_dataloader` 
    * `end_of_dataloader` [depends on](https://github.com/huggingface/accelerate/blob/4305033f8035defad0a87cd38e5c918e78510ba5/src/accelerate/state.py#L1197) `in_dataloader` 
    * `in_dataloader` [depends on](https://github.com/huggingface/accelerate/blob/4305033f8035defad0a87cd38e5c918e78510ba5/src/accelerate/state.py#L1250) `active_dataloader`
* and `active_dataloader` is generally [set to `None`](https://github.com/huggingface/accelerate/blob/4305033f8035defad0a87cd38e5c918e78510ba5/src/accelerate/state.py#L1163), unless explicitly configured otherwise.

While it's been interesting to dig deeper, I believe this issue might be outside the scope of the current PR. Hence, I suggest we stick with the previous approach for now.

Additionally, during my testing, I found a bug in the condition `is_last_step_and_steps_less_than_grad_acc`. It should use `or` instead of `and` for the expected behavior. If `and` is used, the last steps are always dropped if` args.gradient_accumulation_steps` is larger than `step_in_epoch`, which is highly likely to happen. I just removed the comparing `step_in_epoch` and `grad_accum_step` part because updating at the last step covers that case as well. I’ve fixed this in the PR because I believe it is directly related to gradient accumulation.

---

\+ It is not modified now but to use `self.accelerator.sync_gradients`, we should also remove https://github.com/huggingface/transformers/blob/f2c388e3f946862f657acc1e21b272ec946fc66c/src/transformers/trainer.py#L4820  because `_do_sync` function check this as well.

Regarding the failing test ([run_swag_no_trainer.py](https://github.com/huggingface/transformers/blob/main/examples/pytorch/multiple-choice/run_swag_no_trainer.py)), I checked that it is not using the Trainer, which I have fixed. 

Do you have any suggestions or comments on it?

After reading the code a bit more, I don't think we are performing a gradient update at the end of each epoch no (see your example) ? I said that because as you saw the condition using `and` in `is_last_step_and_steps_less_than_grad_acc` + we have set `grad_acc_kwargs["sync_with_dataloader"] = False`. So, I feel like we are doing the right number of update across epoch. The only issue if that we indeed have an issue since there is no N sub-steps between the on_step_begin and on_step_end callbacks. 

```
Epoch 1

Steps 1, 2, 3, 4 (update parameter because (4 % 4 = 0)
Steps 5, 6, 7 (update because it's the last step)
Epoch 2

Step 8 (update parameter because 8 % 4 =0)
Steps 9, 10, 11, 12 (update parameter because 12 % 4 = 0)
Steps 13, 14 (update because it's the last step)
Epoch 3

Step 15, 16 (update parameter because 16 % 4 = 0)
Steps 17, 18, 19, 20 (update parameter because 16 % 4 = 0)
Steps 21 (update because it's the last step)
```

> 
> However, after investigating and experimenting with this, I found that self.gradient_state.end_of_dataloader is always False in my case, preventing the step of self.accelerator from being reset.
> This happens because:
> 
> self.gradient_state.end_of_dataloader depends on self.gradient_state.active_dataloader
> end_of_dataloader [depends on](https://github.com/huggingface/accelerate/blob/4305033f8035defad0a87cd38e5c918e78510ba5/src/accelerate/state.py#L1197) in_dataloader
> in_dataloader [depends on](https://github.com/huggingface/accelerate/blob/4305033f8035defad0a87cd38e5c918e78510ba5/src/accelerate/state.py#L1250) active_dataloader
> and active_dataloader is generally [set to None](https://github.com/huggingface/accelerate/blob/4305033f8035defad0a87cd38e5c918e78510ba5/src/accelerate/state.py#L1163), unless explicitly configured otherwise.

Thanks for exploring ! Could you share a minimal reproducer ? This might indeed require a fix on accelerate side.

> Regarding the failing test ([run_swag_no_trainer.py](https://github.com/huggingface/transformers/blob/main/examples/pytorch/multiple-choice/run_swag_no_trainer.py)), I checked that it is not using the Trainer, which I have fixed.

This is probably a flaky test ! 

Do you have any suggestions or comments on it?

Thank you for reply @SunMarc, 
I may have misinterpreted the code initially. I agree that there's no issue with the update logic. I will go ahead and revert my changes to keep the condition `is_last_step_and_steps_less_than_grad_acc`.

> Thanks for exploring ! Could you share a minimal reproducer ? This might indeed require a fix on accelerate side.

Sure, I will!

> Thank you for reply @SunMarc,
> I may have misinterpreted the code initially. I agree that there's no issue with the update logic. I will go ahead and revert my changes to keep the condition is_last_step_and_steps_less_than_grad_acc.

I mean that the original issue you had (too many updates) might not exist

Yes, you are right, that issue does not exist. 
I thought it always update at last, but it doesn't because it is generally blocked by "steps_less_than_than_grad_acc" condition. 
Sorry for the confusion.

Not an issue ! I was confused also. Still, I will discuss with @muellerzr if it makes sense to switch to back to your idea with the update at each last step, just like how it is coded in accelerate. cc@mueller I'll keep you updated. 

Now, the issue is why self.gradient_state.end_of_dataloader doesn't work as expected. Feel free to open an issue on accelerate library with the reproducer ! Thanks a lot ! 

Okay, I will create a new issue regarding the `self.gradient_state.end_of_dataloader` and wait for updates on this PR. 
Thanks!

Hello @SunMarc, 

I revisited the issue I previously reported regarding `accelerate` and, embarrassingly, it turned out to be a problem with my own codebase. I was using `accelerate` together with a custom `Trainer` but had overridden `get_train_dataloader` without calling `accelerate.prepare(dataloader)`.

To explain the flow in more detail: 
1. `active_dataloader` is set to `None` at first. 
2. When `prepare(dataloader)` is called, it goes into `prepare_data_loader`, where it returns either a `DataLoaderShard` or `DataLoaderShard` that inherits from `DataLoaderStateMixin` as a new dataloader. 
3. These prepared dataloaders add an `active_dataloader` at every `begin()`. 
* Without calling `prepare`, it is clear that the `GradientState` doesn't have an active dataloader

Therefore, **there is no issue with the `accelerate` code itself**. Everything works correctly when `prepare` is called, and I confirmed that at the last batch, `self.gradient_state.end_of_dataloader=True` is set as expected. I also checked this with the pytorch `run_glue.py` example.

Thus, as you suggested, using `if self.accelerator.sync_gradients:` is perfectly fine and results in simpler code.

I’m going to change the current condition to this one and push the update.

Hello, I have updated the if condition as per @SunMarc's suggestion.
Do we have any updates regarding the update at each last step ?

@SunMarc 
Is there any update on this PR? 
If any additional work is needed, please let me know!

(Last bit is resolving the conflicts)

Hello, @SunMarc @muellerzr 

Here is a summary of the new commits:

1. I merged the main branch and resolved conflicts.
2. I replaced `total_batched_samples` with `(step + 1)` and removed `steps_in_epoch <= args.gradient_accumulation_steps` condition to simplify, as we discussed before
> I see that in [accelerate](https://github.com/huggingface/accelerate/blob/4305033f8035defad0a87cd38e5c918e78510ba5/src/accelerate/accelerator.py#L1038), we only turn the gradient sync when we are at the end of the dataloader or step+1 % args.gradient_accumulation_steps == 0. Potentially, we can even remove the condition that we placed in transformers.

3. I reverted the previous changes and set `grad_acc_kwargs["sync_with_dataloader"] = False` as SunMarc mentioned, 
> we will set self.accelerator.gradient_state._set_sync_gradients by ourselves and not rely on the values set by accelerate.accumulate

If there are any issues or additional modifications needed, please let me know!


failing tests come from main, we are in limbo 🫠 

cc @ydshieh 

Hello @muellerzr, I merged the main branch and it looks like it passed the tests! 😄