Save random generator states and resume training for DDP

username: Jiaji_Huang
cooked: <p>I’m running DDP training on a cluster with time limit. When the time limit hits, I have to checkpoint model, optimizer etc.'s states, and resubmit a job by loading these states.</p>
<p>Other than these regular states, if I want the training curve to be exactly the same as if there were no time limit. I have to save random generator’s states as well. I tried the following (the actual code differs, but all key components are included below):</p>
<p>Saving:</p>
<pre><code class="lang-auto"># Each rank does the following
rng_state_dict = {
    'cpu_rng_state': torch.get_rng_state(),
    'gpu_rng_state': torch.cuda.get_rng_state(),
    'numpy_rng_state': numpy.random.get_state(),
    'py_rng_state': random.getstate()
}
torch.save(rng_state_dict, f'rng_state_{rank}.ckpt')
</code></pre>
<p>Loading:</p>
<pre><code class="lang-auto"># Assume main.py already knows its local_rank and (global) rank
# At the very beginning of main.py: each rank does the following
torch.cuda.set_device(local_rank)
rng_state_dict = torch.load(f'rng_state_{rank}.ckpt', map_location='cpu')
torch.set_rng_state(rng_state_dict['cpu_rng_state'])
torch.cuda.set_rng_state(rng_state_dict['gpu_rng_state'])
numpy.random.set_state(rng_state_dict['numpy_rng_state'])
random.setstate(rng_state_dict['py_rng_state'])
</code></pre>
<p>I use the above strategy to save and resubmit job to continue training. However, the training curve differs from the case where there is no time limit. I wonder what piece is missing in my code?</p>

username: agu
cooked: <p>Are you using a PyTorch <code>DataLoader</code>? If so, are you using <code>shuffle=True</code>, and are there any usages of the RNG state before the training loop begins?</p>
<p>My hypothesis is as follows:<br>
If you save the RNG state <code>s_t</code> at the end of an epoch, then if you did not checkpoint and continued training in the next epoch, the <code>DataLoader</code> with <code>shuffle=True</code> consumes the next RNG state <code>s_{t+1}</code>. Now, if you did checkpoint and reload, if you have <em>any</em> usages of the RNG state before the <code>DataLoader</code> consumes its RNG state, then that RNG state for the <code>DataLoader</code> will no longer be <code>s_{t+1}</code> but rather <code>s_{t+k}</code> for <code>k &gt; 1</code>.</p>
<p>If you do not want to trace every possible internal usage of RNG state, can you try setting the seed via <code>torch.manual_seed(epoch)</code> at the beginning of each epoch before iterating over the data loader? Even if my hypothesis is wrong, doing this may still fix your issue. Let me know if you still see differing training curves.</p>

username: Jiaji_Huang
cooked: <p>Thank you for your helpful thoughts! I’m using torch’s dataLoader with <code>shuffle=True</code>, and call <code>dataLoader.set_epoch(epoch)</code> at beginning of each epoch. I believe in dataLoader, there is a “local” random number generator that is fully controlled by <code>seed=base_seed+epoch</code>, see <a href="https://github.com/pytorch/pytorch/blob/master/torch/utils/data/distributed.py#L98" rel="noopener nofollow ugc">line 98-100 of pytorch’s distributedSampler.py</a><br>
So anything consuming the “global” rng states outside dataLoader shouldn’t matter?</p>