Resume-from-checkpoint: restoring optimizer state, not just weights

The problem, from first principles

Large-scale LLM training runs for days or weeks, and interruptions (hardware failures, scheduled maintenance, deliberately pausing to change something) are a genuine, routine part of the process, not a rare edge case. Resuming correctly means far more than just reloading the model's WEIGHTS: [03-dl-training/01-optimizers/04-adam-full-update]'s Adam optimizer maintains its OWN persistent per-parameter state across steps, the running first-moment estimate m and second-moment estimate v, PLUS a step counter t used for bias correction. If a checkpoint saves only the weights and discards this optimizer state, resuming training effectively means restarting Adam's internal history from scratch, with m/v reset to zero and t reset to 1, even though the WEIGHTS themselves are already well into training. This produces a measurably DIFFERENT training trajectory than an uninterrupted run would have taken, exactly the kind of subtle bug that's easy to introduce (checkpointing code that "just" saves model.state_dict() and forgets optimizer.state_dict()) and easy to miss in short experiments, since the damage only compounds over many resumed steps.

From theory to code

Implement save_checkpoint(params, m_list, v_list, t) (saving everything needed to resume identically), train_n_steps (running [04-adam-full-update]'s adam_step repeatedly), resume_from_full_checkpoint (the CORRECT way), and resume_from_weights_only (the WRONG way, included specifically to demonstrate the difference).

Constraints

• save_checkpoint copies (never merely references) params, m_list, and v_list, so later in-place mutation of the originals cannot silently corrupt an already-saved checkpoint.
• resume_from_full_checkpoint continues training using the checkpoint's OWN m_list, v_list, and t, exactly where they left off.
• resume_from_weights_only reinitializes m_list/v_list to zero-filled arrays and t to 1, discarding all prior optimizer history, deliberately reproducing the buggy behavior this question is built to expose.
• Given the SAME sequence of gradients, resume_from_full_checkpoint must reproduce an uninterrupted run's final weights EXACTLY (up to floating-point precision); resume_from_weights_only must NOT.

Hints

return dict(
    params=[p.copy() for p in params],
    m_list=[m.copy() for m in m_list],
    v_list=[v.copy() for v in v_list],
    t=t,
)

def resume_from_full_checkpoint(checkpoint, grads_sequence, lr):
    return train_n_steps(checkpoint["params"], checkpoint["m_list"], checkpoint["v_list"], checkpoint["t"], grads_sequence, lr)

def resume_from_weights_only(params, grads_sequence, lr):
    m_list = [np.zeros_like(p) for p in params]
    v_list = [np.zeros_like(p) for p in params]
    return train_n_steps(params, m_list, v_list, 1, grads_sequence, lr)   # t reset to 1!