Saving Models, GPUs & the PyTorch Toolkit in Review

The final appendix-A topics: making trained weights survive the Python process, and making training fast with GPUs — each one the small-scale original of a pattern the book used at GPT scale.

A.8 Saving and loading models

</> Persist and restore a model

torch.save(model.state_dict(), "model.pth")  #A

model = NeuralNetwork(2, 2)  #B
model.load_state_dict(torch.load("model.pth"))

#A model.state_dict() is a Python dictionary mapping each layer to its trainable parameters (weights and biases) — the portable form of everything the model learned.
#B The architecture must EXACTLY match the saved model: the file carries only named tensors, the class lives in your code. ch5 extended this same pattern with the optimizer's state_dict for resumable training.

A.9 Optimizing training performance with GPUs

Tensors (and whole models) move between devices with .to(...); operations require all operands on the same device:

</> The same computation, CPU then GPU

# CPU
tensor_1 = torch.tensor([1., 2., 3.])
tensor_2 = torch.tensor([4., 5., 6.])
print(tensor_1 + tensor_2)
# tensor([5., 7., 9.])

# GPU
tensor_1 = tensor_1.to("cuda")  #A
tensor_2 = tensor_2.to("cuda")
print(tensor_1 + tensor_2)
# tensor([5., 7., 9.], device='cuda:0')  #B

# Compare matmul speed (the book uses %timeit in a notebook):
a = torch.rand(100, 200)
b = torch.rand(200, 300)
# %timeit a @ b            — CPU timing
a, b = a.to("cuda"), b.to("cuda")
# %timeit a @ b            — GPU timing #C

#A One call moves the data; the same .to(device) works on entire models (ch5: model.to(device)) — which is why calc_loss_batch also moved every BATCH to that device.
#B Same values, new home: device="cuda:0". Mixing devices in one operation raises an error — placement is always explicit.
#C The book leaves the concrete timings to your hardware; matrix multiplication is where GPUs shine, and transformers are almost nothing but matmuls.

Multiple GPUs — distributed training

Distributed training divides model training across multiple GPUs and machines (PyTorch’s DistributedDataParallel). Two practical caveats from the book: DDP does not function properly inside interactive environments like Jupyter notebooks (they don’t handle multiprocessing like a standalone script), and you select which GPUs a script may use via an environment variable — CUDA_VISIBLE_DEVICES=0,2 python some_script.py exposes only the first and third of four GPUs.

A.10 Summary — and where the book used all of this

PyTorch = a tensor library (NumPy-like, GPU-ready) + automatic differentiation + deep-learning utilities. Every piece of this appendix reappeared, scaled up, in the main chapters:

Appendix A concept⇅	Where the book used it⇅
Tensors & shapes	Everywhere — (batch, tokens, emb_dim) from ch2 on	+
Autograd / backward()	ch5–ch7 training loops	+
nn.Module + Sequential	ch3 SelfAttention → ch4 GPTModel	+
Dataset / DataLoader	ch2 windows, ch6 spam, ch7 instructions	+
The training loop	ch5 pretraining, ch6/ch7 fine-tuning	+
state_dict save/load	ch5 checkpoints + OpenAI weights	+
Devices & no_grad	ch5–ch7 throughout	+

The appendix-A toolkit, mapped to its appearances in the book. Click rows for the connection.

Appendix A, closed

Nothing in the main chapters required PyTorch knowledge beyond this appendix — a deliberate design of the book. With the foundations covered, two specialized appendices remain: training-loop refinements (appendix D) and parameter-efficient fine-tuning with LoRA (appendix E).

📝 Check yourself: persistence & GPUs

0 / 4

1.Why must you construct NeuralNetwork(2, 2) BEFORE calling load_state_dict?
2.tensor_1 lives on "cuda" and tensor_2 on the CPU. What happens with tensor_1 + tensor_2?
3.What is distributed training?
4.Your machine has 4 GPUs but you want a script to use only the first and third. How?