Model Checkpointing

Why Checkpoint?

• Resume training after interruption
• Save best model based on validation
• Enable distributed training recovery
• Version control for experiments

What to Save

Full Checkpoint

checkpoint = {
    "model_state_dict": model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
    "scheduler_state_dict": scheduler.state_dict(),
    "epoch": epoch,
    "step": global_step,
    "best_val_loss": best_val_loss,
    "config": model_config,
}
torch.save(checkpoint, "checkpoint.pt")

Model Only (for inference)

torch.save(model.state_dict(), "model.pt")

Loading Checkpoints

Resume Training

checkpoint = torch.load("checkpoint.pt")
model.load_state_dict(checkpoint["model_state_dict"])
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
start_epoch = checkpoint["epoch"] + 1

Load for Inference

model.load_state_dict(torch.load("model.pt"))
model.eval()

Hugging Face Checkpointing

Save

model.save_pretrained("./my_model")
tokenizer.save_pretrained("./my_model")

# Or with Trainer
trainer.save_model("./my_model")

Load

model = AutoModelForCausalLM.from_pretrained("./my_model")
tokenizer = AutoTokenizer.from_pretrained("./my_model")

Best Practices

Checkpointing Strategy

Example: Keep Best + Last 3

import os
import glob

def save_checkpoint(model, optimizer, step, val_loss, save_dir, keep_last=3):
    path = f"{save_dir}/checkpoint-{step}.pt"
    torch.save({...}, path)

    # Remove old checkpoints
    checkpoints = sorted(glob.glob(f"{save_dir}/checkpoint-*.pt"))
    for old in checkpoints[:-keep_last]:
        if "best" not in old:
            os.remove(old)

    # Save best separately
    if val_loss < best_val_loss:
        torch.save({...}, f"{save_dir}/best_model.pt")

Checkpoint Size

Use safetensors for faster saving/loading.