PyTorch essentials for LLM work

Why this matters

Almost every modern LLM was trained in PyTorch (the rest mostly in JAX). HuggingFace Transformers, vLLM, TRL, PEFT, Unsloth, and torchtune are all PyTorch on top. If you cannot read or write nn.Module and a training loop, you cannot actually do this work — you can only call APIs.

This lesson gets you fluent with the PyTorch primitives you will use in every later lesson.

Learning objectives

1. Manipulate tensors with confidence (shapes, dtypes, devices).
2. Define modules with nn.Module and parameters.
3. Write a training loop that supports gradient accumulation, mixed precision, and gradient clipping.
4. Save/load checkpoints correctly.
5. Use DataLoader and Dataset for streaming token batches.

1. Tensors — the data structure

import torch

a = torch.zeros(2, 3)                  # shape (2,3), dtype float32, device cpu
b = torch.tensor([1, 2, 3])
c = torch.randn(4, 4, dtype=torch.float16, device="cuda")
print(c.shape, c.dtype, c.device)

Important attributes: .shape, .dtype, .device, .requires_grad.

Move tensors with .to("cuda") / .cpu(). Convert dtype with .float(), .half(), .bfloat16().

Reshaping ops

• x.view(...) — same memory, requires contiguous tensor.
• x.reshape(...) — view if possible, else copy.
• x.permute(2, 0, 1) — reorder axes (commonly used for (B, T, d) → (B, d, T)).
• x.transpose(-1, -2) — swap last two dims.
• x.unsqueeze(0) / .squeeze() — add/remove a size-1 axis.
• einops.rearrange(x, "b t (h k) -> b h t k", h=8) — the readable reshape; install einops.

Broadcasting

Trailing dimensions must match or be 1.

A = torch.randn(3, 4)
B = torch.randn(4)            # broadcast to (3,4)
A + B                         # works

Use this for masking, bias addition, and per-token scaling.

Common gotchas

• torch.tensor([1, 2]) makes int64 by default → breaks float matmul.
• x.view(-1) flattens; x.flatten(start_dim=1) keeps batch dim.
• x.cpu().numpy() requires requires_grad=False and detach().

2. nn.Module — the building block

Anything trainable is a subclass:

import torch.nn as nn

class TinyMLP(nn.Module):
    def __init__(self, d_in, d_hidden, d_out):
        super().__init__()
        self.fc1 = nn.Linear(d_in, d_hidden)
        self.act = nn.GELU()
        self.fc2 = nn.Linear(d_hidden, d_out)

    def forward(self, x):
        return self.fc2(self.act(self.fc1(x)))

model = TinyMLP(64, 256, 10)
out = model(torch.randn(32, 64))   # (32, 10)
print(sum(p.numel() for p in model.parameters()))

Rules:

• Always call super().__init__() first.
• Submodules assigned to self.xxx are automatically registered (their parameters appear in .parameters()).
• Use nn.ModuleList or nn.ModuleDict (NOT plain Python lists/dicts) for child modules.

Parameters vs buffers

self.weight = nn.Parameter(torch.randn(d, d))       # trainable
self.register_buffer("mask", torch.tril(...))       # NOT trainable but moved with .to(device)

Buffers are how you store causal masks, RoPE frequencies, running statistics — anything the model needs but should not learn.

3. The training loop (annotated)

The single most important code in deep learning. Every framework wraps this:

import torch
from torch.utils.data import DataLoader

model.cuda()
optim = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
scaler = torch.amp.GradScaler("cuda")        # for fp16; bf16 doesn't need a scaler
loader = DataLoader(train_ds, batch_size=32, shuffle=True, num_workers=4, pin_memory=True)

for epoch in range(3):
    for step, (x, y) in enumerate(loader):
        x, y = x.cuda(non_blocking=True), y.cuda(non_blocking=True)

        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
            logits = model(x)
            loss = F.cross_entropy(logits.view(-1, V), y.view(-1))

        scaler.scale(loss).backward()
        scaler.unscale_(optim)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(optim)
        scaler.update()
        optim.zero_grad(set_to_none=True)

        if step % 50 == 0:
            print(f"epoch {epoch} step {step} loss {loss.item():.4f}")

Anatomy

1. Forward under autocast for mixed precision (bf16 on Ampere+, fp16 otherwise).
2. Loss computed in higher precision (cross-entropy is sensitive).
3. Backward via scaler.scale(loss).backward() — the scaler avoids underflow in fp16.
4. Unscale, clip, step, update — clipping at max_norm=1.0 is standard for LLMs.
5. Zero grads with set_to_none=True — slightly faster than zeroing buffers.

Gradient accumulation (for big effective batches on small GPUs)

ACCUM = 8
for step, batch in enumerate(loader):
    loss = compute_loss(batch) / ACCUM
    loss.backward()
    if (step + 1) % ACCUM == 0:
        optim.step()
        optim.zero_grad(set_to_none=True)

This simulates an 8x larger batch with no extra memory cost (just slower wall-clock).

Mixed precision in 2026

• bf16 is the default everywhere training-side. Same exponent range as fp32, no scaler needed.
• fp16 only on T4 / V100 — use a scaler.
• fp8 (E4M3 / E5M2) is now production for H100 + Blackwell training (NVIDIA Transformer Engine, FP8 in DeepSpeed). You will not write it from scratch but you will see flags.

4. Datasets and DataLoaders

For LLM pre-training you stream from disk; for fine-tuning you usually load a HuggingFace dataset:

from datasets import load_dataset
ds = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft")
print(ds[0])                                  # one chat record

A Dataset class needs __len__ and __getitem__:

from torch.utils.data import Dataset

class TokenDataset(Dataset):
    def __init__(self, ids, block_size):
        self.ids = ids
        self.block = block_size
    def __len__(self):
        return (len(self.ids) - 1) // self.block
    def __getitem__(self, i):
        start = i * self.block
        x = self.ids[start : start + self.block]
        y = self.ids[start + 1 : start + 1 + self.block]
        return torch.tensor(x), torch.tensor(y)

DataLoader(dataset, ...) handles batching, shuffling, multi-process loading.

IterableDataset for huge corpora

When tokens won't fit in RAM (Phase 3 territory), implement IterableDataset and stream from sharded .bin files (memmap).

5. Saving and loading

torch.save(model.state_dict(), "model.pt")             # parameters only
model.load_state_dict(torch.load("model.pt", map_location="cpu"))

For training mid-run you also save optimiser and scheduler:

torch.save({
    "model": model.state_dict(),
    "optim": optim.state_dict(),
    "step":  step,
}, "ckpt.pt")

Use safetensors for production: faster, no arbitrary code execution, the format the whole HF ecosystem uses.

from safetensors.torch import save_file, load_file
save_file(model.state_dict(), "model.safetensors")

6. Devices and parallelism (preview)

• Single GPU: .to("cuda") and you are done.
• Multi-GPU on one node: torch.nn.parallel.DistributedDataParallel (DDP) — covered in Phase 3.
• Sharded across many GPUs: torch.distributed.fsdp.FullyShardedDataParallel (FSDP) — also Phase 3.
• Automatic mixed precision: torch.amp.autocast.
• torch.compile(model) (PyTorch 2.x): JIT-compiles forward into Triton kernels; usually 1.3-2x speedup for free.

model = torch.compile(model)

Hands-on lab (4 hours)

pytorch_drills.ipynb:

1. Build a 3-layer MLP that classifies MNIST. Get ≥97% test accuracy in ≤3 epochs.
2. Add torch.compile. Measure speedup.
3. Replace AdamW with torch.optim.SGD(momentum=0.9). Compare convergence curves.
4. Add gradient accumulation with ACCUM=4 and verify final accuracy is unchanged.
5. Add bf16 autocast. Verify the loss matches fp32 within 0.5%.
6. Save and reload the trained model; reproduce test accuracy.
7. Bonus: convert to safetensors format and load it back.

Common pitfalls

1. Forgetting optim.zero_grad() → gradients accumulate every step (not what you want unless intentional).
2. model.train() vs model.eval() — affects dropout, batchnorm (less relevant for LLMs but matters for CV).
3. Using a Python list of layers instead of nn.ModuleList → those layers' params won't appear in model.parameters().
4. Calling .to("cuda") after creating the optimiser — optim still references CPU tensors.
5. Mixing torch.from_numpy(arr) (shares memory) with torch.tensor(arr) (copies). Surprising bugs result.
6. Letting num_workers=0 and wondering why GPU is starved on data.

Self-check

1. Difference between a parameter and a buffer.
2. Why use set_to_none=True for gradient zeroing?
3. What does gradient clipping prevent?
4. When would you use bf16 vs fp16?
5. What does torch.compile do under the hood?

References

• PyTorch official tutorials — start with the 60-minute blitz.
• PyTorch internals (Edward Yang) — for the curious.
• einops docs — readable tensor reshaping.
• HuggingFace Datasets quickstart.
• PyTorch performance tuning guide.