Reproducible checkpoint

[sgugger]

What does this PR do?

This PR fixes reproducibility when resuming training from a checkpoint for PyTorch >= 1.6 and non-TPU trainings. The current implementation in the Trainer had two flaws:

First when going over the data at the beginning (skipping epochs and batches), the first shuffling was the same in the original training or resuming from the checkpoint, but not the subsequents shufflings. This is because those shufflings were determined by the global torch RNG which had a different state at the end of epoch 1 in a full training vs just after the shuffle of epoch 1 in the resumed training. I did not identify the exact reason but there was one or several calls to the global torch RNG during epoch 1 of the full training and when resuming we only called for the shuffle of epoch 1.

To fix this, a generator is now used to determine the shuffling, this way the generator is set with the same seed in both cases (full training or resuming from checkpoint) and then has the exact same number of calls (one shuffle per epoch) so that generator ends up in the same state in the full training or when resuming from a checkpoint.

The second thing that was different is the CUDAs RNG states, which are used in the dropout layer of the models. To fix this, all RNG states are saved when saving a checkpoint and set after all the data skip phase when resuming training.

Fixes #11504
Fixes #11323

[stas00]

Wonderful!

Thank you for working on this, @sgugger!

[stas00]

Suggested change

	if version.parse(torch.__version__) < version.parse("1.6.0"):
	return RandomSampler(self.train_dataset)
	# Torch generator were introduced in PyTorch 1.6.0.
	generator = torch.Generator()
	generator.manual_seed(int(torch.empty((), dtype=torch.int64).random_().item()))
	return RandomSampler(self.train_dataset, generator=generator)
	kwargs = {}
	if has_torch_generator():
	generator = torch.Generator()
	generator.manual_seed(int(torch.empty((), dtype=torch.int64).random_().item()))
	kwargs.update(dict(generator=generator))
	return RandomSampler(self.train_dataset, **kwargs)

Here is another possible way (including the suggested earlier has_torch_generator which doesn't currently exist)

[stas00]

I'm just wondering whether you checked that different processes in dist env will have the same cpu RNG state?

[stas00]

And another question - what about python's main RNG?

[sgugger]

In distributed training, the torch generator should have the same state on all processes since it's set with the same seed initially (then all processes execute the same code). Unless someone goes out of their way to only execute a random instruction on one process.

The code has been tested one GPU, 2 GPUs with DataParallel and 2 GPUs distributed and the same result are obtained with a full training and resuming from the last checkpoint.

As for python main RNG and numpy main RNG I have no way of extracting the seed (I can set it but not get the current state easily) so any code that wants to be 100% reproducible when resuming from a checkpoint needs to use torch RNGs only.

[stas00]

As for python main RNG and numpy main RNG I have no way of extracting the seed (I can set it but not get the current state easily)

Won't these work?

py_rng_state = random.getstate()
random.setstate(py_rng_state)

np_rng_state = numpy.random.get_state()
numpy.random.set_state(np_rng_state)

so any code that wants to be 100% reproducible when resuming from a checkpoint needs to use torch RNGs only.

Do you mean to say that they should stick to torch's random APIs? and convert to other formats from there?

[sgugger]

Oh TIL! I didn't think it was possible to extract and set those. Will add that to the PR.

[LysandreJik]

Cool, this should come in handy!

[LysandreJik]

Could this method also print a warning in case we're on TPU as we don't expect reproducibility when on TPUs?

[sgugger]

Mmm, maybe just a comment in the README where we document resuming from checkpoint? I don't really want to issue a warning for each run on TPU using a checkpoint.

[sgugger]

@stas00 Could you double check the last changes are all good with you? They include loading/saving all RNGs state (and separately in distributed training) as you suggested.

[stas00]

Everything looks great, @sgugger.

Given the fragile nature of this added feature it'd be relatively easily to break it, so may I suggest that at least some rudimentary test would go a long way to prevent this?

(fragile, not because how it was coded, but because it's just fragile in its own random nature)

perhaps duplicating test_trainer's resume test and generating a random number on each np, pt, py before and after resume?

[stas00]

Added a new doc section in a separate PR #11588

[dlwh]

i'm very late here, but shouldn't this use self.args.seed?

[sgugger]

No, this is copied from PyTorch actually. Since torch has been seeded, this will be deterministic.