

TITLE
[Trainer]: Resume training with `save_strategy="epoch"` does not load RNG state

COMMENTS
2

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
### System Info

```shell
- `transformers` version: 4.19.0.dev0
- Platform: Linux-5.15.36-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. `git clone https://github.com/atreyasha/transformers`
2. Create a virtual environment and install the `[dev-torch]` extras
3. `pytest tests/trainer/test_trainer.py::TrainerIntegrationTest::test_resume_training_with_randomness_from_epoch`

**Edit**: I removed the forked repository as the diff has been incorporated in the PR mentioned below.

Here is the relevant test snippet where I added `save_strategy="epoch"` and adjusted the checkpoint number to reflect the steps in one epoch:

```python
    @require_torch_non_multi_gpu
    def test_resume_training_with_randomness_from_epoch(self):
        # This test will fail flakily for more than 1 GPUs since the result will be slightly more different
        # TODO: investigate why it fails for 2 GPUs?

        if torch.cuda.is_available():
            torch.backends.cudnn.deterministic = True
        train_dataset = RegressionDataset(length=128)
        eval_dataset = RegressionDataset()

        config = RegressionModelConfig(a=0, b=2)
        model = RegressionRandomPreTrainedModel(config)

        tmp_dir = self.get_auto_remove_tmp_dir()
        args = RegressionTrainingArguments(tmp_dir, save_strategy="epoch", learning_rate=0.1)
        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-16"))
        (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: 2.0825276374816895 != 2.081479072570801 within 1e-08 delta (0.0010485649108886719 difference)
```

### Cause

The RNG state is only loaded when resuming a checkpoint that completed non-zero steps in the current epoch. If the checkpoint was saved at the end of the epoch, `steps_trained_in_current_epoch` would be `0` for the new epoch and the saved RNG state would not be loaded.

https://github.com/huggingface/transformers/blob/da47c264f9a881f5db5f6fbb59a30c95e428571f/src/transformers/trainer.py#L1423-L1435

### Possible fix

Check if the checkpoint to resume is a whole-number multiple of steps per epoch. If this is true, then load the RNG state once before entering the `epoch_iterator` loop above.

### Expected behavior

The test case above should pass, meaning that the regression variables should be the same or similar (within the delta).