Save & Load a Finetuned Checkpoint

You've spent hours curating data, wrangling GPU memory, and watching loss curves crawl downward. Then the training session crashes, your notebook kernel dies, or your cloud instance gets recycled — and everything vanishes. This is the exact pain this lesson solves: the ability to save and load a finetuned checkpoint so your hard-won model weights survive interruptions and can be reused or deployed later. Without this skill, every finetuning run is a one-shot gamble; with it, you get a safety net and the foundation for sharing and serving your model.

The problem this lesson solves

Finetuning a large language model is not a single, linear operation. It's a long, resource-intensive process that can take hours or days, and it's riddled with points of failure: power outages, spot-instance terminations, out-of-memory errors, or simply closing the wrong terminal tab. If you don't explicitly save your progress, all of that compute is wasted.

Beyond the crash scenario, think about the workflow: you want to try different hyperparameters, compare early checkpoints against final ones, or hand your model off to a teammate for evaluation. Each of these requires you to persist model state to disk and reload it on demand.

The core issue is that a finetuned LLM is a stateful artifact — it's not just code that runs the same every time; it carries learned weights, optimizer states, and tokenizer vocabularies. If you don't manage that state correctly, you can't reproduce results, resume training, or serve the model in production. This lesson gives you the tools to treat your finetuning runs as continuous, resumable experiments rather than black-box one-offs.

Core concept / mental model

Think of a checkpoint like a video-game save point. In a game, you save before a boss fight so you can retry without replaying the whole level. In finetuning, you save after a certain number of steps or epochs so you can:

• Resume if training crashes.
• Rollback to a model version with better validation loss.
• Deploy the model for inference without keeping the Trainer running.

Technically, a checkpoint bundles several files that together capture the full training state:

• Model weights — the learned parameters of the transformer (the state_dict).
• Optimizer state — momentum, variance, and learning-rate scheduler positions (needed to resume training exactly).
• Tokenizer — the vocabulary and tokenization rules that map text to input IDs.
• Config — model configuration (architecture, vocab size, hidden dimensions) and potentially the training arguments.

A useful analogy: the model weights are like a cooked dish, while the optimizer state is the recipe with notes on how far you got through the cooking process. To serve the dish, you only need the dish; to continue cooking seamlessly, you need the notes too.

How it works step by step

Saving and loading a finetuned checkpoint with Hugging Face Transformers follows a predictable pattern. Here's how it works under the hood:

1. Training with a trainer: You use Trainer from Transformers, which calls save_model() at the end of training and periodically via the save_strategy argument.
2. Saving the model: The trainer writes files to the output directory — typically pytorch_model.bin (or model.safetensors), config.json, the tokenizer files, and training_args.bin if you use save_only_model=False.
3. Loading for inference: Once you want to use the model (not continue training), you load it back with AutoModelForCausalLM.from_pretrained() and its tokenizer with AutoTokenizer.from_pretrained(). No optimizer state needed.
4. Resuming training: To pick up where you left off, you pass resume_from_checkpoint=True to Trainer.train(), and it loads the full snapshot including optimizer and scheduler.

💡 Pro tip: The save_strategy in TrainingArguments defaults to steps, so you can set save_steps (e.g., every 500 steps) to create regular snapshots. These become your safety net.

Hands-on walkthrough

Let's make this concrete with three complete examples.

Example 1: Save and load a model checkpoint for inference

First, finetune a small model and save the final checkpoint. We'll use GPT-2 as a placeholder; in a real project you'd swap in your own model and dataset.

from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
from datasets import Dataset

texts = ["Hello world", "Finetuning is fun", "GPT-2 saves checkpoints"]
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token

def tokenize(batch):
    return tokenizer(batch["text"], truncation=True, padding=True, max_length=64)

dataset = Dataset.from_dict({"text": texts}).map(tokenize, batched=True)

model = AutoModelForCausalLM.from_pretrained("gpt2")

training_args = TrainingArguments(
    output_dir="./gpt2-finetuned",
    save_strategy="epoch",          # Save a checkpoint at the end of each epoch
    save_total_limit=2,              # Keep only the last 2 checkpoints to save space
    num_train_epochs=1,
    per_device_train_batch_size=2,
    logging_steps=10,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
)
trainer.train()
trainer.save_model()  # Explicitly save the final model (also happens automatically)

# Expected folder contents after training:
# ./gpt2-finetuned/
#   config.json
#   model.safetensors
#   tokenizer.json
#   tokenizer_config.json
#   training_args.bin

Loading the saved model for inference:

from transformers import AutoModelForCausalLM, AutoTokenizer

model_dir = "./gpt2-finetuned"
model = AutoModelForCausalLM.from_pretrained(model_dir)
tokenizer = AutoTokenizer.from_pretrained(model_dir)

inputs = tokenizer("Finetuning is", return_tensors="pt")
outputs = model.generate(**inputs, max_new_tokens=10)
print(tokenizer.decode(outputs[0]))
# Expected output (may vary): "Finetuning is fun"

Example 2: Resume training from a checkpoint

Imagine you want to train for another epoch without losing progress. The Trainer makes this trivial.

from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir="./gpt2-finetuned",
    num_train_epochs=2,  # Now we train for 2 epochs total; first epoch already done
    save_strategy="steps",
    save_steps=500,
)

trainer = Trainer(
    model=AutoModelForCausalLM.from_pretrained("gpt2"),
    args=training_args,
    train_dataset=dataset,  # same dataset as before
)

trainer.train(resume_from_checkpoint="./gpt2-finetuned/checkpoint-XXX")  # replace with actual checkpoint dir

Example 3: Save only the weights (lightweight output)

If you only need the model weights (e.g., to export to ONNX or share elsewhere), use model.save_pretrained().

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

output_dir = "./my-lora-model"  # Example for a PEFT-merged model
model.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)
print("Saved to", output_dir)

Compare options / when to choose what

Rule of thumb: Use full checkpoints during training (for resumability); use the final model or a PEFT adapter for deployment or sharing. If you finetuned with LoRA, saving only the adapter dramatically reduces storage and lets you swap it onto different base models.

Troubleshooting & edge cases

1. OSError: Can't load model

This usually means the checkpoint directory is missing a config.json or the model architecture doesn't match. Fix: always save the config with the model (save_pretrained does this).

2. Tokenizer mismatch after loading

Loading a model with the wrong tokenizer leads to garbled inputs. Fix: always load both with the same checkpoint directory and avoid manually overriding the tokenizer.

3. Resume training changes results

If you resume from a checkpoint but change the data order, the randomness might differ. Fix: set the same seed in TrainingArguments for reproducibility.

4. Disk space fills with checkpoints

save_total_limit controls how many checkpoints are kept; older ones are automatically deleted. Set it to a reasonable number (e.g., 2–3) to balance safety and storage.

What you learned & what's next

You now understand how to save and load a finetuned checkpoint — from the mental model of checkpoints as save points, to the step-by-step mechanics of Trainer, to choosing the right saving strategy for your use case. You can save a final model for inference, resume training after a crash, and keep your work safe with save_strategy.

This skill connects directly to the next lesson in the LLM Finetuning track, where you'll likely evaluate your finetuned model. With checkpoints safely stored, you can now load them for metric computation, human evaluation, or further training without fear of losing progress. Master this, and your finetuning experiments become truly iterative.

Next, you'll learn how to evaluate the quality of a finetuned model — the natural next step after preserving your checkpoints.