How to save the random state of Dropout for resuming training?

meteorlin: <p>Hello everyone, I am trying to use PyTorch to save model checkpoints, optimizer states, and random states for ‘resume training’. However, I found that if there is a dropout in the model, it cannot reproduce the same output as the original model after loading checkpoints, random seeds, and random states.<br>
Below is a minimum test unit. You can first set ‘load_checkpoint=1’ and run it once to save the checkpoint, then set it to 0 and run it again.<br>
As you can see, I have already set the same random seeds (including torch, torch.cuda, numpy, and random) and optimizer states before starting the experiment.<br>
In theory, if the states are consistent, the output values of the model should be the same in both normal training and checkpoint loading modes (refer to the <code>print()</code> on line 118). That’s exactly the case, in the parameter <code>use_dropout=0</code> (i.e. the model does not use Dropout), the final output values are the same. But when <code>use_dropout=1</code>, the output values in the two modes are different.<br>
The reason for this issue is that most models have Dropouts (such as BERT). However, due to limited resources, we may not be able to complete the training at once. And due to the dropout, the performance changing is astonishing between normal training and resuming training modes (in the task I am currently working on).<br>
How should this be resolved? One solution I can think of is to save the random state of the dropout, but it seems infeasible.<br>
My current version of Pytorch is <code>1.13.0+cu116 </code>.</p>
<pre><code class="lang-python">import os, argparse, pickle, sys
import random, torch
import numpy as np

import torch.nn as nn
from tqdm import tqdm

SEED = 1

parser = argparse.ArgumentParser(description='KGTL')
parser.add_argument('--cuda', nargs='+', type=int, default=[0,1,2,3], help="cuda")
parser.add_argument('--device', type=int, default=0, help='-1 is CPU')
parser.add_argument('--load_checkpoint', type=int, default=1, help="whether to load checkpoint or not, 0 is not")
parser.add_argument('--use_dropout', type=int, default=1, help="whether to use Dropout on model or not, 0 is not")
params = parser.parse_args()

os.environ['CUDA_VISIBLE_DEVICES'] = ','.join([str(i) for i in params.cuda])
os.environ['WORLD_SIZE'] = str(len(params.cuda))
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
# set the state of random
torch.manual_seed(SEED)
torch.cuda.manual_seed_all(SEED)
np.random.seed(SEED)
random.seed(SEED)
torch.backends.cudnn.deterministic = True

class M(nn.Module):
    def __init__(self) -&gt; None:
        super().__init__()
        self.E = nn.Embedding(5, 128)
        self.linear = nn.Linear(128, 256)
        self.tanh = nn.Tanh()
        self.dropout = nn.Dropout(0.2)
    
    def forward(self, data):
        if params.use_dropout != 0:
            return self.dropout(
                self.tanh(self.linear(self.E(data)))
            )
        else:
            return self.tanh(self.linear(self.E(data)))

batch_size = 3
save_at_epoch = 6
device = torch.device(f"cuda:{params.device}" if params.device != -1 else "cpu")
load_checkpoint = params.load_checkpoint == 1
dir_path = "./test_checkpoints"
os.makedirs(dir_path, exist_ok=True)
ELEs = [
    "checkpoint.pt", "", 
    "random_state.pkl", "torch_random_state.pkl", "torch_cuda_random_state.pkl", "numpy_random_state.pkl",
    "optimizer.pkl"
]

if load_checkpoint:
    with open(os.path.join(dir_path, ELEs[2]), 'rb') as handle:
        random.setstate(pickle.load(handle))
    with open(os.path.join(dir_path, ELEs[3]), 'rb') as handle:
        torch.set_rng_state(pickle.load(handle))
    with open(os.path.join(dir_path, ELEs[4]), 'rb') as handle:
        torch.cuda.set_rng_state_all(pickle.load(handle))
    with open(os.path.join(dir_path, ELEs[5]), 'rb') as handle:
        np.random.set_state(pickle.load(handle))

model = M().to(device)
opt = torch.optim.Adam(model.parameters(), lr=2e-4)
loss_func = nn.CrossEntropyLoss()

if load_checkpoint:
    print("loading the checkpoint and states...")
    model.load_state_dict(
        torch.load(os.path.join(dir_path, ELEs[0]), map_location=device)
    )
    with open(os.path.join(dir_path, ELEs[6]), 'rb') as handle:
        opt.load_state_dict(pickle.load(handle))

training_data = [1,2,3,3,2,1,1,2,3,3,2,1,1,2,3,3,2,1,1,2,3,3,2,1]
training_labels = [1 for _ in range(len(training_data))]

for epoch in tqdm(range(8)):
    if epoch == save_at_epoch:
        if not load_checkpoint:
            print("saving the checkpoint and states...")
            # save the random state of random
            with open(os.path.join(dir_path, ELEs[2]), 'wb') as handle:
                pickle.dump(random.getstate(), handle, protocol=pickle.HIGHEST_PROTOCOL)
            # save the random state of torch
            with open(os.path.join(dir_path, ELEs[3]), 'wb') as handle:
                pickle.dump(torch.get_rng_state(), handle, protocol=pickle.HIGHEST_PROTOCOL)
            # save the random state of cuda
            with open(os.path.join(dir_path, ELEs[4]), 'wb') as handle:
                pickle.dump(torch.cuda.get_rng_state_all(), handle, protocol=pickle.HIGHEST_PROTOCOL)
            # save the random state of numpy
            with open(os.path.join(dir_path, ELEs[5]), 'wb') as handle:
                pickle.dump(np.random.get_state(), handle, protocol=pickle.HIGHEST_PROTOCOL)
            
            # save the state of optimizer
            with open(os.path.join(dir_path, ELEs[6]), 'wb') as handle:
                pickle.dump(opt.state_dict(), handle, protocol=pickle.HIGHEST_PROTOCOL)

            # save the checkpoint of model
            torch.save(model.state_dict(), os.path.join(dir_path, ELEs[0]))            
    else:
        if load_checkpoint:
            continue
        
    for index in range(0, len(training_data), batch_size):
        batch = training_data[index:index+batch_size]
        labels = training_labels[index:index+batch_size]
        
        opt.zero_grad()
        loss = loss_func(
            model(torch.LongTensor(batch).to(device)), 
            torch.LongTensor(labels).to(device)
        )

        if epoch == save_at_epoch:
            print(loss.item())
            sys.exit()

        loss.backward()
        opt.step()
</code></pre>

ptrblck: <p>Your code to restore the PRNG states looks wrong since you are using a different order of calls into the number generator.<br>
In the first run you are:</p>
<ul>
<li>seeding the code</li>
<li>initializing the model, which calls into the PRNG to sample the parameters and buffers</li>
<li>save the states</li>
<li>train the model.</li>
</ul>
<p>If you restore the states you are:</p>
<ul>
<li>seeding the code</li>
<li>loading the states</li>
<li>initializing the model (this will change the states)</li>
<li>load the parameters of the model and optimizer</li>
<li>train the model.</li>
</ul>
<p>Handling the PRNG with seeds is not trivial since you have to guarantee to keep the same order of calls, which is not the case here.</p>