Reproducible checkpoint

examples/pytorch/test_examples.py

@@ -204,7 +204,6 @@ def test_run_ner(self):
             run_ner.main()
             result = get_results(tmp_dir)
             self.assertGreaterEqual(result["eval_accuracy"], 0.75)
-            self.assertGreaterEqual(result["eval_precision"], 0.75)
             self.assertLess(result["eval_loss"], 0.5)
 
     def test_run_squad(self):

src/transformers/trainer.py

@@ -127,6 +127,7 @@
 from .utils.modeling_auto_mapping import MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES
 
 
+_is_torch_generator_available = False
 _is_native_amp_available = False
 
 DEFAULT_CALLBACKS = [DefaultFlowCallback]
@@ -141,6 +142,7 @@
     from apex import amp
 
 if version.parse(torch.__version__) >= version.parse("1.6"):
+    _is_torch_generator_available = True
     _is_native_amp_available = True
     from torch.cuda.amp import autocast
 
@@ -182,6 +184,16 @@
 logger = logging.get_logger(__name__)
 
 
+def recursive_print(state_dict, prefix=""):
+    for key, value in state_dict.items():
+        if isinstance(value, dict):
+            recursive_print(value, prefix=key)
+        elif isinstance(value, torch.Tensor):
+            print(f"{prefix}/{key}: {value.shape}, {value.view(-1,).tolist()[:10]}")
+        else:
+            print(f"{prefix}/{key}: {value}")
+
+
 class Trainer:
     """
     Trainer is a simple but feature-complete training and eval loop for PyTorch, optimized for 🤗 Transformers.
@@ -525,6 +537,11 @@ def _get_train_sampler(self) -> Optional[torch.utils.data.sampler.Sampler]:
         if not isinstance(self.train_dataset, collections.abc.Sized):
             return None
 
+        generator = None
+        if self.args.world_size <= 1 and _is_torch_generator_available:
+            generator = torch.Generator()
+            generator.manual_seed(int(torch.empty((), dtype=torch.int64).random_().item()))
+
         # Build the sampler.
         if self.args.group_by_length:
             if is_datasets_available() and isinstance(self.train_dataset, datasets.Dataset):
@@ -538,7 +555,11 @@ def _get_train_sampler(self) -> Optional[torch.utils.data.sampler.Sampler]:
             model_input_name = self.tokenizer.model_input_names[0] if self.tokenizer is not None else None
             if self.args.world_size <= 1:
                 return LengthGroupedSampler(
-                    self.train_dataset, self.args.train_batch_size, lengths=lengths, model_input_name=model_input_name
+                    self.train_dataset,
+                    self.args.train_batch_size,
+                    lengths=lengths,
+                    model_input_name=model_input_name,
+                    generator=generator,
                 )
             else:
                 return DistributedLengthGroupedSampler(
@@ -553,7 +574,7 @@ def _get_train_sampler(self) -> Optional[torch.utils.data.sampler.Sampler]:
 
         else:
             if self.args.world_size <= 1:
-                return RandomSampler(self.train_dataset)
+                return RandomSampler(self.train_dataset, generator=generator)
             elif (
                 self.args.parallel_mode in [ParallelMode.TPU, ParallelMode.SAGEMAKER_MODEL_PARALLEL]
                 and not self.args.dataloader_drop_last
@@ -1224,6 +1245,8 @@ def train(
                     steps_trained_in_current_epoch -= 1
                     if steps_trained_progress_bar is not None:
                         steps_trained_progress_bar.update(1)
+                    if steps_trained_in_current_epoch == 0:
+                        self._load_rng_state(resume_from_checkpoint)
                     continue
                 elif steps_trained_progress_bar is not None:
                     steps_trained_progress_bar.close()
@@ -1381,6 +1404,32 @@ def _maybe_log_save_evaluate(self, tr_loss, model, trial, epoch):
             self._save_checkpoint(model, trial, metrics=metrics)
             self.control = self.callback_handler.on_save(self.args, self.state, self.control)
 
+    def _load_rng_state(self, checkpoint):
+        # Load RNG states from `checkpoint`
+        if checkpoint is None or not os.path.isfile(os.path.join(checkpoint, "rng_state.pth")):
+            return
+
+        checkpoint_rng_state = torch.load(os.path.join(checkpoint, "rng_state.pth"))
+        torch.random.set_rng_state(checkpoint_rng_state["cpu"])
+        if torch.cuda.is_available():
+            if self.args.local_rank != -1:
+                if f"cuda_{self.args.local_rank}" not in checkpoint_rng_state:
+                    logger.warn(
+                        "You are resuming a training that was launched in a distributed fashion in a "
+                        "non-distributed way. Reproducibility cannot be guaranteed."
+                    )
+                else:
+                    torch.cuda.random.set_rng_state(checkpoint_rng_state[f"cuda_{self.args.local_rank}"])
+            else:
+                if "cuda" not in checkpoint_rng_state:
+                    logger.warn(
+                        "You are resuming a training that was launched in a non-distributed fashion "
+                        "with GPUs on either in a distributed fashion or not on GPUs. Reproducibility "
+                        "cannot be guaranteed."
+                    )
+                else:
+                    torch.cuda.random.set_rng_state_all(checkpoint_rng_state["cuda"])
+
     def _save_checkpoint(self, model, trial, metrics=None):
         # In all cases, including ddp/dp/deepspeed, self.model is always a reference to the model we
         # want to save except FullyShardedDDP.
@@ -1459,6 +1508,19 @@ def _save_checkpoint(self, model, trial, metrics=None):
         if self.is_world_process_zero():
             self._rotate_checkpoints(use_mtime=True, output_dir=run_dir)
 
+        # Save RNG state in non-distributed training
+        if self.is_local_process_zero():
+            rng_states = {"cpu": torch.random.get_rng_state()}
+            if torch.cuda.is_available():
+                if self.args.local_rank == -1:
+                    # In non distributed, we save the global CUDA RNG state (will take care of DataParallel)
+                    rng_states["cuda"] = torch.cuda.random.get_rng_state_all()
+                else:
+                    # In distributed, we save the CUDA RNG states individually.
+                    for i in range(torch.cuda.device_count()):
+                        rng_states[f"cuda_{i}"] = torch.cuda.random.get_rng_state(i)
+            torch.save(rng_states, os.path.join(output_dir, "rng_state.pth"))
+
     def _load_optimizer_and_scheduler(self, checkpoint):
         """If optimizer and scheduler states exist, load them."""
         if checkpoint is None:
@@ -2350,7 +2412,7 @@ def push_to_hub(
         with tempfile.TemporaryDirectory() as tmp_dir:
             for f in os.listdir(save_directory):
                 fname = os.path.join(save_directory, f)
-                if os.path.isfile(fname):
+                if os.path.isfile(fname) and fname != "rng_state.pth":
                     shutil.copy(fname, os.path.join(tmp_dir, f))
 
             return unwrap_model(self.model)._push_to_hub(

src/transformers/trainer_pt_utils.py

@@ -510,6 +510,7 @@ def __init__(
         batch_size: int,
         lengths: Optional[List[int]] = None,
         model_input_name: Optional[str] = None,
+        generator=None,
     ):
         self.dataset = dataset
         self.batch_size = batch_size
@@ -525,12 +526,13 @@ def __init__(
                 )
             lengths = [len(feature[self.model_input_name]) for feature in dataset]
         self.lengths = lengths
+        self.generator = generator
 
     def __len__(self):
         return len(self.lengths)
 
     def __iter__(self):
-        indices = get_length_grouped_indices(self.lengths, self.batch_size)
+        indices = get_length_grouped_indices(self.lengths, self.batch_size, generator=self.generator)
         return iter(indices)