§5.4–5.6Saving, Loading & Borrowing Weights

Part II pp. 89–90 · ~6 min read

  • state_dict

Training costs real time and money — its results shouldn’t evaporate when the Python process exits. This final unit of chapter 5 persists what we’ve built: first checkpointing our own model (with everything needed to resume training), then the chapter’s endgame — filling our architecture with OpenAI’s pretrained GPT-2 weights.

5.4 Saving and loading model weights

The right unit of persistence is the state_dict : a dictionary mapping each parameter’s name (tok_emb.weight, trf_blocks.0.att.W_query.weight, …) to its tensor. One subtlety matters: LLMs are trained with adaptive optimizers like AdamW, which store additional state for every model weight (moment estimates). To continue pretraining later, that state must survive too:

</> Checkpointing model AND optimizer 
torch.save({
  "model_state_dict": model.state_dict(),        #A
  "optimizer_state_dict": optimizer.state_dict(),  #B
  },
  "model_and_optimizer.pth"
)
  • #A All 163M learned parameters, by name. Enough on its own for inference-only use.
  • #B AdamW's per-parameter statistics — without these, resumed training starts the optimizer cold and behaves differently from an uninterrupted run.
</> Restoring the checkpoint and resuming 
checkpoint = torch.load("model_and_optimizer.pth", weights_only=True)  #A

model = GPTModel(GPT_CONFIG_124M)  #B
model.load_state_dict(checkpoint["model_state_dict"])

optimizer = torch.optim.AdamW(model.parameters(), lr=0.0005, weight_decay=0.1)
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])  #C
model.train()  #D
  • #A weights_only=True restricts deserialization to plain tensors — the safe way to load files, refusing arbitrary pickled code.
  • #B The state_dict carries only numbers: you must first re-create the architecture (the class + config), then pour the weights in. Names and shapes must match exactly.
  • #C The optimizer is rebuilt over the fresh model's parameters, then its saved statistics are restored.
  • #D Loading doesn't set the mode — switch to train() to continue pretraining (or eval() if you're here to generate).
model.state_dict() 163M named tensors optimizer.state_dict() AdamW moments per weight torch.save model_and_optimizer.pth one file, both dictionaries torch.load 1. rebuild: GPTModel(GPT_CONFIG_124M) the class + config live in CODE, not the file 2. load_state_dict × 2 names & shapes must match exactly 3. model.train() — or .eval() mode matches what happens next result: training resumes as if never interrupted — same weights, same optimizer momentum
1

Two dictionaries to preserve

The model's learned parameters AND the optimizer's adaptive state. Saving only the first is fine for inference; resuming TRAINING needs both.

1 / 5

The checkpoint round-trip: what goes into the .pth file and what has to happen — in order — to come back from it.

5.5 Loading pretrained weights from OpenAI

Pretraining an LLM on a large corpus is time- and resource-intensive — chapter 1 put GPT-scale runs in the millions of dollars. The alternative: load OpenAI’s openly published GPT-2 weights into the very architecture we built. One catch: the published checkpoint (original TensorFlow format) uses different tensor names and structure than our GPTModel, so loading requires a custom mapping — matching each of OpenAI’s tensors to the corresponding module of ours:

OpenAI GPT-2 checkpointwte (token embedding)wpe (positional embedding)h.0.attn.c_attn (fused QKV!)h.0.mlp.c_fc / c_proj …TensorFlow naming & layout,output head tied to wtecustom mappingname → module, split/reshapeour GPTModeltok_emb / pos_embtrf_blocks[i].att.W_query/K/Vtrf_blocks[i].ff / normsfinal_norm · out_headsame architecture — which iswhy the mapping works at all
Same machine, different labels: OpenAI’s checkpoint stores the QKV projections fused and the output head tied to the token embedding (the ch4 weight-tying story), so a custom routine matches, splits, and copies each tensor into place.

Thin ice, marked

The source notes cover §5.5 in a single sentence — “custom loading is required because the data-structure names differ” — so this page keeps to that scope. The full mapping code lives in the book’s repository; what matters for the journey: after loading, our hand-built model generates coherent English, and the fine-tuning chapters start from exactly these weights.

5.6 Summary

  • LLMs generate one token at a time; by default the highest-probability token wins — greedy decoding.
  • Probabilistic sampling and temperature scaling let us influence the diversity and coherence of the output; top-k restricts the pool.
  • Training and validation losses gauge generated-text quality during training; pretraining minimizes the training loss with a conventional cross-entropy objective and the AdamW optimizer.
  • Pretraining a large model from scratch is expensive — loading openly available weights is the practical alternative.

Stage 2 complete

The model is built (ch2–4), trainable, measurable, and — via OpenAI’s weights — genuinely fluent. Everything from here on is specialization: chapter 6 fine-tunes this foundation into a spam classifier; chapter 7 teaches it to follow instructions.

📝 Check yourself: saving, loading & borrowed weights

0 / 5

  1. 1.Why does the checkpoint save the OPTIMIZER's state_dict alongside the model's?

  2. 2.What exactly is a state_dict?

  3. 3.After loading a checkpoint to CONTINUE pretraining, the code calls model.train(). Why?

  4. 4.Why can't OpenAI's published GPT-2 weights be loaded with a plain load_state_dict call?

  5. 5.What's the payoff of loading OpenAI's pretrained weights instead of pretraining ourselves?