Repository: huggingface/transformers (Python) — Bug Fix

Issue:
[Trainer]: Resume training with `save_strategy="epoch"` does not load RNG state
### 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).


Discussion:

Resolving patch:
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -732,13 +732,16 @@ def estimate_tokens(self, input_dict: Dict[str, Union[torch.Tensor, Any]]) -> in
         Returns:
             `int`: The total number of tokens.
         """
+        if not hasattr(self, "warnings_issued"):
+            self.warnings_issued = {}
         if self.main_input_name in input_dict:
             return input_dict[self.main_input_name].numel()
-        else:
+        elif "estimate_tokens" not in self.warnings_issued:
             logger.warning(
                 "Could not estimate the number of tokens of the input, floating-point operations will not be computed"
             )
-            return 0
+            self.warnings_issued["estimate_tokens"] = True
+        return 0
 
     def floating_point_ops(
         self, input_dict: Dict[str, Union[torch.Tensor, Any]], exclude_embeddings: bool = True
@@ -838,6 +841,7 @@ def __init__(self, config: PretrainedConfig, *inputs, **kwargs):
         # Save config and origin of the pretrained weights if given in model
         self.config = config
         self.name_or_path = config.name_or_path
+        self.warnings_issued = {}
 
     def post_init(self):
         """
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -1150,7 +1150,8 @@ def train(
             kwargs:
                 Additional keyword arguments used to hide deprecated arguments
         """
-        resume_from_checkpoint = None if not resume_from_checkpoint else resume_from_checkpoint
+        if resume_from_checkpoint is False:
+            resume_from_checkpoint = None
 
         # memory metrics - must set up as early as possible
         self._memory_tracker.start()
@@ -1394,6 +1395,9 @@ def train(
             )
             self.control = self.callback_handler.on_epoch_begin(args, self.state, self.control)
 
+            if epoch == epochs_trained and resume_from_checkpoint is not None and steps_trained_in_current_epoch == 0:
+                self._load_rng_state(resume_from_checkpoint)
+
             step = -1
             for step, inputs in enumerate(epoch_iterator):
 


Test patch:
diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py
--- a/tests/trainer/test_trainer.py
+++ b/tests/trainer/test_trainer.py
@@ -57,7 +57,6 @@
     require_torch_bf16,
     require_torch_gpu,
     require_torch_multi_gpu,
-    require_torch_non_multi_gpu,
     require_torch_tf32,
     require_torch_up_to_2_gpus,
     require_wandb,
@@ -161,11 +160,12 @@ def __call__(self, eval_pred):
 
 
 class RegressionModelConfig(PretrainedConfig):
-    def __init__(self, a=0, b=0, double_output=False, **kwargs):
+    def __init__(self, a=0, b=0, double_output=False, random_torch=True, **kwargs):
         super().__init__(**kwargs)
         self.a = a
         self.b = b
         self.double_output = double_output
+        self.random_torch = random_torch
         self.hidden_size = 1
 
 
@@ -263,14 +263,18 @@ def __init__(self, config):
             super().__init__(config)
             self.a = nn.Parameter(torch.tensor(config.a).float())
             self.b = nn.Parameter(torch.tensor(config.b).float())
+            self.random_torch = config.random_torch
 
         def forward(self, input_x, labels=None, **kwargs):
             y = input_x * self.a + self.b
-            torch_rand = torch.randn(1).squeeze()
+            if self.random_torch:
+                torch_rand = torch.randn(1).squeeze()
             np_rand = np.random.rand()
             rand_rand = random.random()
 
-            y += 0.05 * torch_rand + 0.05 * torch.tensor(np_rand + rand_rand)
+            if self.random_torch:
+                y += 0.05 * torch_rand
+            y += 0.05 * torch.tensor(np_rand + rand_rand)
 
             if labels is None:
                 return (y,)
@@ -996,33 +1000,60 @@ def test_can_resume_training(self):
             trainer.train(resume_from_checkpoint=True)
         self.assertTrue("No valid checkpoint found in output directory" in str(context.exception))
 
-    @require_torch_non_multi_gpu
     def test_resume_training_with_randomness(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?
+        # 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=128)
         eval_dataset = RegressionDataset()
 
-        config = RegressionModelConfig(a=0, b=2)
-        model = RegressionRandomPreTrainedModel(config)
+        with self.subTest("Test every step"):
+            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=5, learning_rate=0.1)
-        trainer = Trainer(model, args, train_dataset=train_dataset, eval_dataset=eval_dataset)
+            tmp_dir = self.get_auto_remove_tmp_dir()
+            args = RegressionTrainingArguments(tmp_dir, save_steps=5, 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()
+            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-15"))
-        (a1, b1) = 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-15"))
+            (a1, b1) = trainer.model.a.item(), trainer.model.b.item()
+
+            self.assertAlmostEqual(a, a1, delta=1e-8)
+            self.assertAlmostEqual(b, b1, delta=1e-8)
+
+        with self.subTest("Test every epoch"):
+            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_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)
+
+            checkpoints = [d for d in os.listdir(tmp_dir) if d.startswith("checkpoint-")]
+            # There should be one checkpoint per epoch.
+            self.assertEqual(len(checkpoints), 3)
+            checkpoint_dir = sorted(checkpoints, key=lambda x: int(x.replace("checkpoint-", "")))[0]
+
+            trainer.train(resume_from_checkpoint=os.path.join(tmp_dir, checkpoint_dir))
+            (a1, b1) = trainer.model.a.item(), trainer.model.b.item()
 
-        self.assertAlmostEqual(a, a1, delta=1e-8)
-        self.assertAlmostEqual(b, b1, delta=1e-8)
+            self.assertAlmostEqual(a, a1, delta=1e-8)
+            self.assertAlmostEqual(b, b1, delta=1e-8)
 
     # regression for this issue: https://github.com/huggingface/transformers/issues/12970
     def test_training_with_resume_from_checkpoint_false(self):