Resume Training from Checkpoints

You've poured hours of compute and money into a fine-tuning run, and just when the loss curve looked beautiful, your GPU died, your cloud instance was preempted, or your session timed out. Without a plan, you'd be forced to restart from scratch, burning days of progress and budget. In this lesson, you'll master the art of resuming training from a checkpoint safely — not just hitting "resume" and hoping for the best, but doing it in a way that guarantees consistency, avoids silent corruption, and gets you back on track faster than ever.

The problem this lesson solves

Training a large language model (LLM) is a marathon, not a sprint. A single fine-tuning run can take hours or even days, using expensive GPUs and consuming significant energy. The harsh reality is that interruptions are inevitable: hardware failures, spot instance terminations, power outages, or even a careless Ctrl+C during a lunch break.

If you've ever had to restart a training run from epoch zero because you didn't have a robust checkpointing strategy, you know the pain. The loss curve resets, the learning rate schedule restarts, and you watch the same progress slowly climb again — a frustrating and costly déjà vu.

The core problem is state management. A training run isn't just the model weights; it includes the optimizer state (like Adam's momentum and variance), the learning rate scheduler, the random number generator state, the current global step, and your training/evaluation data loaders. If you only save the model weights, you can't truly resume training — you can only start a new run from those weights, losing the subtle dynamics that were built up over thousands of steps.

Pro tip: Always plan for checkpointing before you start training, not as an afterthought. The best time to design a recovery strategy is when everything is going well.

Core concept / mental model

Think of a training run as a long, complex journey. The model weights are your destination on a map; the optimizer state is your vehicle's fuel and engine tuning; and the scheduler state is your GPS, telling you how fast to go and when to stop for refueling. A checkpoint is a snapshot of the entire journey at a specific point in time.

To resume training from a checkpoint safely, you need to capture and restore the complete state of the run, not just the destination. A safe resume is like pressing "play" on a paused video — you expect to continue exactly where you left off, not rewind to the beginning.

Here's a mental model with the key components:

• Model weights (model_state_dict): The learned parameters of your LLM.
• Optimizer state (optimizer_state_dict): Internal variables like momentum and variance for adaptive optimizers like AdamW. Without this, you lose the "memory" of past gradients.
• Scheduler state (scheduler_state_dict): The current learning rate and position in the decay schedule.
• Global step and epoch: The position in your training loop.
• RNG states: For Python's random, NumPy's np.random, and PyTorch's torch.random. Critical for reproducibility, especially with data shuffling and augmentation.
• Data loader state: Which batch you were on, especially important if you use distributed sampling or custom iterators.

A simple mental model: A checkpoint is a container that holds a full snapshot of your training universe. A safe resume restores the entire snapshot, not just a fragment.

How it works step by step

Safely resuming training from a checkpoint follows a systematic process. Let's break it down:

1. Create regular checkpoints: During training, save the complete state every N steps or epochs. Don't rely on a single checkpoint — keep a few recent ones (e.g., best model and latest step) to hedge against corruption.

2. Design a checkpoint format: Use a structured format like a directory with separate files (e.g., model.pt, optimizer.pt, scheduler.pt, metadata.json) or a single tar archive. Ensure atomic writes by saving to a temporary file and then renaming, to avoid partial checkpoints.

3. Record all essential state: At each checkpoint, include the model, optimizer, scheduler, step, epoch, RNG states, and any custom training loop variables (e.g., best validation loss, patience counters).

4. Detect an interruption: When training stops (due to error or user interruption), catch the signal or exception and log that you're stopping. This is where you'll later decide to resume.

5. Load the checkpoint: When starting a new training script, check for the existence of a checkpoint. If found, load all components back into memory.

6. Validate the loaded state: After loading, perform sanity checks — e.g., verify the model weights match the expected shape, the optimizer state is consistent, and the current step is what you expect.

7. Restore the data loader state: Skip the number of samples already processed. This is often the trickiest part, especially with custom data pipelines. Use torch.utils.data.DataLoader with a sampler that you can set to a specific index.

8. Resume the loop: Continue the training loop from the saved step/epoch, not from scratch.

The key principle is atomicity and completeness. A safe resume is impossible if you miss any part of the state.

Remember: A checkpoint is not just a model file. It's a complete snapshot of your training run.

Hands-on walkthrough

Let's get our hands dirty with a concrete example using PyTorch and Hugging Face Transformers. We'll build a simple training script that can be safely interrupted and resumed.

Saving a checkpoint

First, let's define a helper function that saves the full training state. We'll use a temporary file and rename it for atomicity.

import os
import torch
import json
from pathlib import Path

def save_checkpoint(state, filename):
    """Save a checkpoint atomically."""
    checkpoint_dir = Path(filename).parent
    checkpoint_dir.mkdir(parents=True, exist_ok=True)
    tmp_path = filename + ".tmp"
    torch.save(state, tmp_path)
    os.replace(tmp_path, filename)  # atomic rename
    print(f"Checkpoint saved to {filename}")

def save_training_state(model, optimizer, scheduler, step, epoch, best_loss, rng_states):
    """Assemble and save the full training state."""
    state = {
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "scheduler_state_dict": scheduler.state_dict() if scheduler else None,
        "step": step,
        "epoch": epoch,
        "best_loss": best_loss,
        "rng_states": rng_states,
    }
    save_checkpoint(state, f"checkpoints/checkpoint_step_{step}.pt")
    # Also save a 'latest' checkpoint for easy resume
    save_checkpoint(state, f"checkpoints/latest.pt")

Capturing RNG states

Before each training step, you can capture the current RNG states. This is useful for resuming with exact reproducibility.

def get_rng_states():
    """Capture RNG states for Python, NumPy, and PyTorch."""
    return {
        "python": random.getstate(),
        "numpy": np.random.get_state(),
        "torch": torch.random.get_rng_state(),
        "cuda": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
    }

Loading and resuming

Now, the resume function:

def load_checkpoint(config, model, optimizer, scheduler):
    """Load a checkpoint and restore the full training state."""
    checkpoint_path = config.get("resume_from", "checkpoints/latest.pt")
    if not os.path.exists(checkpoint_path):
        return model, optimizer, scheduler, 0, 0, float('inf'), None

    print(f"Loading checkpoint from {checkpoint_path}")
    state = torch.load(checkpoint_path, map_location=config["device"])

    model.load_state_dict(state["model_state_dict"])
    if optimizer and state["optimizer_state_dict"]:
        optimizer.load_state_dict(state["optimizer_state_dict"])
    if scheduler and state["scheduler_state_dict"]:
        scheduler.load_state_dict(state["scheduler_state_dict"])

    rng_states = state.get("rng_states")
    if rng_states:
        random.setstate(rng_states["python"])
        np.random.set_state(rng_states["numpy"])
        torch.random.set_rng_state(rng_states["torch"])
        if rng_states["cuda"] and torch.cuda.is_available():
            torch.cuda.set_rng_state_all(rng_states["cuda"])

    return model, optimizer, scheduler, state["step"], state["epoch"], state["best_loss"], state

Putting it together

Here's a minimal but complete training loop that uses these functions:

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_scheduler

def train(config):
    # Setup model, optimizer, scheduler
    model = AutoModelForSequenceClassification.from_pretrained(config["model_name"])
    optimizer = torch.optim.AdamW(model.parameters(), lr=config["lr"])
    num_training_steps = config["num_steps"]
    scheduler = get_scheduler("linear", optimizer=optimizer, num_warmup_steps=100, num_training_steps=num_training_steps)

    start_step, start_epoch, best_loss, _ = load_checkpoint(config, model, optimizer, scheduler)

    # Dummy dataloader for example
    train_loader = create_dataloader(config)
    # Restore data loader position if needed
    train_loader = restore_dataloader(train_loader, start_step, config)

    model.train()
    global_step = start_step
    for epoch in range(start_epoch, config["num_epochs"]):
        for batch in train_loader:
            # forward/backward
            outputs = model(**batch)
            loss = outputs.loss
            loss.backward()
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()
            global_step += 1

            # Save checkpoint every N steps
            if global_step % config["save_every"] == 0:
                rng_states = get_rng_states()
                save_training_state(model, optimizer, scheduler, global_step, epoch, loss.item(), rng_states)

        # Optionally save at end of epoch
    print("Training complete!")

Note: In a real project, you'd also handle KeyboardInterrupt to save a final checkpoint, and use distributed training utilities like torch.distributed with proper checkpoint sharding.

Compare options / when to choose what

There are several ways to implement checkpointing. Here's a comparison of common approaches:

For most real-world LLM fine-tuning, full state checkpointing is the gold standard. It gives you the flexibility to resume from arbitrarily long interruptions without losing progress.

Pro tip: Always save at least two checkpoints: the latest one and the one with the best validation metric. If the latest is corrupted (e.g., due to a crash during save), you can fall back to the best.

Troubleshooting & edge cases

Even with a solid plan, things can go wrong. Here are common issues and how to fix them:

1. Mismatched model architecture: If you change the model configuration (e.g., hidden size) between runs, loading the state dict will fail. Fix: Always load with strict=False and handle missing/unexpected keys, or better, don't change architecture when resuming.

2. Corrupted checkpoint: A partial write (e.g., power loss during save) can make the file unreadable. Fix: Use atomic writes (like in our example) and verify file integrity with checksums or just try/except when loading.

3. Learning rate schedule restarts: If you forget to load scheduler_state_dict, the LR will reset to the initial high value, potentially causing loss spikes. Fix: Always save and load the scheduler state.

4. Data loader resets to beginning: If you don't restore the data loader state, you'll repeat previously seen data, overfitting and corrupting the training progress. Fix: Implement a sampler that can be advanced to a specific index.

5. RNG state mismatch: Different RNG states lead to different data orders, which may not be a real bug but can affect reproducibility. Fix: Restore RNG states as shown.

6. CUDA out-of-memory when loading: The checkpoint may contain tensors on GPU. Fix: Use map_location='cpu' when loading, then move to GPU after.

Remember: The two most common silent bugs are forgetting to restore the scheduler and not resuming the data loader correctly. Always verify that your training loss curve continues smoothly after a resume.

What you learned & what's next

Congratulations! You've learned the fundamental principles of resuming training from a checkpoint safely. You can now:

• Explain the importance of saving the complete training state, not just model weights.
• Implement full-state checkpointing in a PyTorch training loop, including optimizer, scheduler, RNG states, and step counts.
• Load a checkpoint and resume training exactly where you left off, avoiding wasted compute.
• Troubleshoot common edge cases like corrupted checkpoints and scheduler resets.

These skills are essential for any serious LLM fine-tuning project, saving you time, money, and frustration.

Next in your learning path, you'll explore evaluating your fine-tuned model. You'll learn how to set up proper evaluation metrics, compare different checkpoints, and decide when your model is ready for deployment. With robust checkpointing under your belt, you'll be able to iterate confidently through multiple training cycles, always knowing you can recover from any interruption.