Step shifting using total_batched_samples for gradient_accumulation_steps counting
Author: kibitzing
Pull request: no
Labels: name: trainer

name: bug

### System Info

- `transformers` version: 4.39.0
- Platform: Linux-5.4.239-1.el7.elrepo.x86_64-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.24.6
- Safetensors version: 0.4.1
- Accelerate version: 0.28.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.2.0+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 GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

_No response_

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

### Issue Analysis
The `on_step_begin` callback is invoked when the `step` is divisible by `args.gradient_accumulation_steps` (i.e.,` step % args.gradient_accumulation_steps == 0`).
However, the `on_step_end` callback behaves differently. Its condition is as follows:
`total_batched_samples % args.gradient_accumulation_steps == 0 or is_last_step_and_steps_less_than_grad_acc`

Here, the `on_step_end` callback is triggered when `total_batched_samples` is divisible by `args.gradient_accumulation_steps`. It’s important to note that `step` is reset at the beginning of each epoch, whereas `total_batched_samples` is initialized to 0 at the start of training and persists across all epochs until training ends. 

### Expected Behavior:

When `gradient_accumulation_steps = N`, there should be exactly N sub-steps between the `on_step_begin` and `on_step_end` callbacks. This ensures that gradients are accumulated correctly before an optimization step occurs. The only exception to this rule is the last step in an epoch or the training run, where fewer sub-steps might exist.

### Problematic Behavior Example
The issue arises when `total_batched_samples` is not divisible by args.gradient_accumulation_steps. For example, if `steps_per_epoch = 3` and `gradient_accumulation_steps = 2`, we observe the following behavior:

Epoch 1:
* step 0: `on_step_begin` called, `on_step_end` not called (expected behavior)
* step 1: `on_step_begin` not called, `on_step_end` called (expected behavior)
* step 2: `on_step_begin` called, `on_step_end` not called (expected behavior)

Epoch 2:
* step 3: `on_step_begin` called (0 % 2 == 0), `on_step_end` called (4 % 2 == 0) (incorrect because on_step_end is called after only one sub step)
* step 4: `on_step_begin` not called (1 % 2 != 0), `on_step_end` not called (5 % 2 != 0) (incorrect)
* step 5: `on_step_begin` called (2 % 2 == 0), `on_step_end` called (6 % 2 == 0) (incorrect because on_step_end is called after only one sub step)

Epoch 3:
* step 6: `on_step_begin` called, `on_step_end` not called (expected behavior)
* step 7: `on_step_begin` not called, `on_step_end` called (expected behavior)
* step 8: `on_step_begin` called, `on_step_end` not called (expected behavior)

Note: `total_batched_samples` is incremented by 1 at the start of each step loop.

In this case, when the number of steps per epoch is not divisible by gradient_accumulation_steps, the callbacks only function correctly at intervals, leading to incorrect behavior during other epochs.

Comments:
Upon further investigation, I found that this issue is not solely related to the callbacks. When performing gradient accumulation for updates, if the total number of batches in an epoch is not divisible by `gradient_accumulation_steps`, a shifting phenomenon occurs with gradient accumulation.

Specifically, after updating with the last non-divisible batch, the gradient accumulation should be counted from 0. However, since total_batched_samples is not refreshed with each epoch, it can lead to this shifting issue.

For example, if the number of data `batches_per_epoch = 7` and `gradient_accumulation_steps = 4`, the updates would proceed as follows:
* 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)




cc @SunMarc and @muellerzr

Thanks for this clear report @kibitzing ! I left a [comment](https://github.com/huggingface/transformers/pull/33673#pullrequestreview-2329873116) on your PR

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.