Resume Training Exactly

Save a training run and resume it so that the result is the same, bit for bit, as a run that never stopped.

Long training jobs stop: a machine is pre-empted, a job hits its time limit. A checkpoint must hold everything that affects the next step, not only the weights:

model.state_dict()        weights and buffers
optimizer.state_dict()    Adam's moment estimates and step counts
scheduler.state_dict()    where the learning rate schedule is
torch.get_rng_state()     the random state used by dropout and shuffling
step                      how far the run got

Write save_checkpoint to put these in one file with torch.save, and load_checkpoint to restore them into fresh objects and return the step.

Function signature

def save_checkpoint(path, model, optimizer, scheduler, step):

def load_checkpoint(path, model, optimizer, scheduler):

Example

Input

run A: train 12 steps
run B: train 6 steps, save_checkpoint(path, model, opt, sched, 6)
       build a new model, optimizer and scheduler with another seed
       step = load_checkpoint(path, model, opt, sched); train 6 more steps

Output

step  ->  6
the weights of run A and run B are exactly equal

Reasoning

After loading, run B has the same weights, the same Adam moments, the same learning rate and the same random state as run A had at step 6. Every later step then does exactly the same arithmetic.

Constraints and notes

• The model uses dropout and the loop draws random batches, so the random state matters.
• load_checkpoint gets objects that already exist. Load state into them with load_state_dict. Do not replace them.
• The file must load with torch.load(path, weights_only=True). Save state dicts and tensors, not whole modules or optimizers, which need pickle.