[Trainer]: Resume training not consistent with large effective batch size
Author: atreyasha
Pull request: no
Labels: name: bug

### System Info

```shell
- `transformers` version: 4.19.0.dev0
- Platform: Linux-5.15.37-1-lts-x86_64-with-glibc2.33
- Python version: 3.8.12
- Huggingface_hub version: 0.5.1
- PyTorch version (GPU?): 1.11.0+cu102 (False)
- 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?: No
- Using distributed or parallel set-up in script?: No
```

### Who can help?

@sgugger

### Information

- [X] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

I provide a MWE for this issue by forking `transformers` and writing a failing test case. This can be reproduced via the steps below:

1. Run `git clone https://github.com/atreyasha/transformers`
2. Create a virtual environment and install the `[dev-torch]` extras
3. Run `pytest tests/trainer/test_trainer.py::TrainerIntegrationTest::test_resume_training_with_randomness_large_accumulation`

Following is a snippet of the failing test:

```python
def test_resume_training_with_randomness_large_accumulation(self):
    # For more than 1 GPUs, since the randomness is introduced in the model and with DataParallel (which is used
    # in this test for more than 2 GPUs), the calls to the torch RNG will happen in a random order (sometimes
    # GPU 0 will call first and sometimes GPU 1).
    random_torch = not torch.cuda.is_available() or torch.cuda.device_count() <= 1

    if torch.cuda.is_available():
        torch.backends.cudnn.deterministic = True
    train_dataset = RegressionDataset(length=10)
    eval_dataset = RegressionDataset()

    config = RegressionModelConfig(a=0, b=2, random_torch=random_torch)
    model = RegressionRandomPreTrainedModel(config)

    tmp_dir = self.get_auto_remove_tmp_dir()
    args = RegressionTrainingArguments(tmp_dir,
                                       save_steps=1,
                                       per_device_train_batch_size=1,
                                       gradient_accumulation_steps=6,
                                       learning_rate=0.1,
                                       num_train_epochs=2)
    trainer = Trainer(model, args, train_dataset=train_dataset, eval_dataset=eval_dataset)

    trainer.train()
    (a, b) = trainer.model.a.item(), trainer.model.b.item()

    model = RegressionRandomPreTrainedModel(config)
    trainer = Trainer(model, args, train_dataset=train_dataset, eval_dataset=eval_dataset)
    trainer.train(resume_from_checkpoint=os.path.join(tmp_dir, "checkpoint-1"))
    (a1, b1) = trainer.model.a.item(), trainer.model.b.item()

    self.assertAlmostEqual(a, a1, delta=1e-8)
    self.assertAlmostEqual(b, b1, delta=1e-8)
```

This should produce an error because the regression variables are not the same or similar:

```console
> self.assertAlmostEqual(a, a1, delta=1e-8)
E AssertionError: 0.1493121236562729 != 0.1493326723575592 within 1e-08 delta (2.0548701286315918e-05 difference)
```

I observed that this issue occurs when the effective batch size (`per_device_train_batch_size` * `gradient_accumulation_steps`) is strictly greater than half of the training set size. In practice, we normally don't encounter such scenarios since effective batch sizes tend to be much smaller than even half the size of the training set.

But in any case, I came across this bug while playing with some settings and thought it to be worth reporting; especially since resuming training consistently should (AFAICT) not be affected by the effective batch size.

### Cause

I did some debugging and found a cause. So basically if the accumulation steps are greater than half of the training set, then ` num_update_steps_per_epoch` would equal `1` as per below:

https://github.com/huggingface/transformers/blob/1766fa21599fc99442027cab15139f146b2301a3/src/transformers/trainer.py#L1300

Correspondingly, the partially completed epoch where the checkpoint was saved will be treated as complete because of the exact divison below:

https://github.com/huggingface/transformers/blob/1766fa21599fc99442027cab15139f146b2301a3/src/transformers/trainer.py#L1398

As a result when training is resumed from the checkpoint, it will be treated as if it starts in a new epoch where it actually should have started in the partial epoch before.

### Possible fix

The following diff fixes this issue:

```diff
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
index aa54f2af1..6fad35ff0 100755
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -1297,7 +1297,7 @@ class Trainer:
         len_dataloader = None
         if has_length(train_dataloader):
             len_dataloader = len(train_dataloader)
-            num_update_steps_per_epoch = len_dataloader // args.gradient_accumulation_steps
+            num_update_steps_per_epoch = math.ceil(len_dataloader / args.gradient_accumulation_steps)
             num_update_steps_per_epoch = max(num_update_steps_per_epoch, 1)
             num_examples = self.num_examples(train_dataloader)
             if args.max_steps > 0:
@@ -1531,8 +1531,7 @@ class Trainer:

                 if (step + 1) % args.gradient_accumulation_steps == 0 or (
                     # last step in epoch but step is always smaller than gradient_accumulation_steps
-                    steps_in_epoch <= args.gradient_accumulation_steps
-                    and (step + 1) == steps_in_epoch
+                    (step + 1) == steps_in_epoch
                 ):
                     # Gradient clipping
                     if args.max_grad_norm is not None and args.max_grad_norm > 0 and not self.deepspeed:
```

But this then changes the behaviour of gradient accumulation such that any final steps before the end of the epoch are still used for training, even if they are explicitly smaller than the accumulation steps. Maybe there is another better way that I can't see.

### Expected behavior

The test case above should pass, meaning that the regression variables should be the same or similar (within the delta).


Comments:
I don't see exactly what fails at the top of my mind for this degenerate case, but I won't investigate this anytime soon as I have more pressing items on my TODO (like fixing bugs that happen in real life ;-) ).

Thanks @sgugger for the prompt response. I did some debugging and found a possible cause and fix. Updated the description above.

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.