Repository: huggingface/transformers (Python)
Date: 2020-11-18 16:42:19+00:00
Task type: Bug Fix

Issue:
Here maybe a bug, when we load staged checkpoint
https://github.com/huggingface/transformers/blob/40d98ebf50c4662bcd6dce6395bbed0b2142ea52/src/transformers/trainer.py#L458

I met this bug when I used the setting below:

global_steps = 2748
len(train_dataloader) = 27484
gradient_accumulation_steps = 4

In the original code, "steps_trained_in_current_epoch" will be 2748 !  BUT this variable should be  2748*4 = 10,992

the code I suggested is below:

```
epochs_trained = (self.global_step * self.args.gradient_accumulation_steps) // len(train_dataloader)
steps_trained_in_current_epoch = (self.global_step * self.args.gradient_accumulation_steps) % len(train_dataloader)
```


Issue discussion:
This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

I'm also puzzled by this. The calculations here seems incorrect.
To me these calculations are not incorrect if we take `step` as optimization steps, however `steps_trained_in_current_epoch` is wrongly used to skip training batches without considering gradient accumulation. 

+1 for the proposed calculation for `steps_trained_in_current_epoch` as the number of batches to be skipped.
@sgugger might be interested in this.
There is indeed a problem, but only with `steps_trained_in_current_epoch`. The `global_step` variable represents the number of optimization steps, not the number of batches seen. The variable `num_update_steps_per_epoch` take this into account so `epochs_trained` is correct. `steps_trained_in_current_epoch` represents the number of update steps to skip but is used as the number of batches to skip, so either need to multiply it by the `gradient_accumulation_steps` (and rename it for clarity) or skip  `gradient_accumulation_steps` batches before subtracting 1 to it later in the loop.

This also shows that we direly miss a test to check resuming training works with gradient accumulation. I can look into this when I have a bit of time, but will be fairly busy with the preparation for v4.

Solution patch:
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -676,11 +676,12 @@ def train(self, model_path: Optional[str] = None, trial: Union["optuna.Trial", D
             self.state = TrainerState.load_from_json(os.path.join(model_path, "trainer_state.json"))
             epochs_trained = self.state.global_step // num_update_steps_per_epoch
             steps_trained_in_current_epoch = self.state.global_step % (num_update_steps_per_epoch)
+            steps_trained_in_current_epoch *= self.args.gradient_accumulation_steps
 
             logger.info("  Continuing training from checkpoint, will skip to saved global_step")
             logger.info("  Continuing training from epoch %d", epochs_trained)
             logger.info("  Continuing training from global step %d", self.state.global_step)
-            logger.info("  Will skip the first %d steps in the first epoch", steps_trained_in_current_epoch)
+            logger.info("  Will skip the first %d batches in the first epoch", steps_trained_in_current_epoch)
 
         # Update the references
         self.callback_handler.model = self.model


Test patch:
diff --git a/tests/test_trainer.py b/tests/test_trainer.py
--- a/tests/test_trainer.py
+++ b/tests/test_trainer.py
@@ -465,6 +465,14 @@ def test_save_checkpoints(self):
             trainer.train()
             self.check_saved_checkpoints(tmpdir, 5, int(self.n_epochs * 64 / self.batch_size), False)
 
+    def test_gradient_accumulation(self):
+        # Training with half the batch size but accumulation steps as 2 should give the same results.
+        trainer = get_regression_trainer(
+            gradient_accumulation_steps=2, per_device_train_batch_size=4, learning_rate=0.1
+        )
+        trainer.train()
+        self.check_trained_model(trainer.model)
+
     def test_can_resume_training(self):
         if torch.cuda.device_count() > 2:
             # This test will fail for more than 2 GPUs since the batch size will get bigger and with the number of
@@ -514,6 +522,38 @@ def test_can_resume_training(self):
             self.assertEqual(b, b1)
             self.assertEqual(state, state1)
 
+    def test_resume_training_with_gradient_accumulation(self):
+        if torch.cuda.device_count() > 2:
+            # This test will fail for more than 2 GPUs since the batch size will get bigger and with the number of
+            # save_steps, the checkpoint will resume training at epoch 2 or more (so the data seen by the model
+            # won't be the same since the training dataloader is shuffled).
+            return
+        with tempfile.TemporaryDirectory() as tmpdir:
+            trainer = get_regression_trainer(
+                output_dir=tmpdir,
+                train_len=128,
+                gradient_accumulation_steps=2,
+                per_device_train_batch_size=4,
+                save_steps=5,
+                learning_rate=0.1,
+            )
+            trainer.train()
+            (a, b) = trainer.model.a.item(), trainer.model.b.item()
+            state = dataclasses.asdict(trainer.state)
+
+            checkpoint = os.path.join(tmpdir, "checkpoint-5")
+
+            # Reinitialize trainer and load model
+            model = RegressionPreTrainedModel.from_pretrained(checkpoint)
+            trainer = Trainer(model, trainer.args, train_dataset=trainer.train_dataset)
+
+            trainer.train(model_path=checkpoint)
+            (a1, b1) = trainer.model.a.item(), trainer.model.b.item()
+            state1 = dataclasses.asdict(trainer.state)
+            self.assertEqual(a, a1)
+            self.assertEqual(b, b1)
+            self.assertEqual(state, state1)
+
     def test_load_best_model_at_end(self):
         total = int(self.n_epochs * 64 / self.batch_size)
         with tempfile.TemporaryDirectory() as tmpdir: