15. PyTorch mental model

A PyTorch tensor is a NumPy-style array that also knows its device and can record the operations that produced it. Training is three moves: run the forward computation, let autograd walk that record backward to get gradients, and let the optimizer update the parameters.

Tensors: array, plus device, plus autograd

import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
x = torch.randn(32, 128, device=device)                  # float32
w = torch.randn(128, 10, device=device, requires_grad=True)
y = x @ w                                                # (32, 10)
y.shape, y.dtype, y.device, y.requires_grad

• Shapes, strides, views and broadcasting work as in NumPy (Part 13). PyTorch says dim where NumPy says axis.
• Every tensor in an operation must live on the same device. “Expected all tensors to be on the same device” means something was not moved.
• torch.from_numpy(a) shares memory with a. Going back requires t.detach().cpu().numpy().
• In-place operations end in an underscore (add_, zero_). They can break autograd if they overwrite a value that backward still needs.

Autograd from first principles

Each operation on a tensor that requires gradients records a node: its inputs, the operation, and how to compute its local derivatives. The result’s grad_fn points to that node. loss.backward() walks the graph in reverse, applies the chain rule, and accumulates gradients into the .grad of leaf tensors (your parameters).

a = torch.tensor(2.0, requires_grad=True)
b = torch.tensor(3.0, requires_grad=True)
c = a * b + b ** 2       # dc/da = b = 3;  dc/db = a + 2b = 8
c.backward()
a.grad, b.grad           # (tensor(3.), tensor(8.))

• Gradients accumulate with +=. That is why you call optimizer.zero_grad() every step, and why gradient accumulation over micro-batches works.
• The graph is rebuilt on every forward pass (“define-by-run”), so ordinary Python if and for inside forward just work.
• The graph holds intermediate activations until backward runs. That is where training memory goes.
• .detach() cuts a tensor out of the graph; .item() returns a Python number and forces a GPU sync.
• torch.no_grad() and torch.inference_mode() record no graph, so evaluation and inference use less memory and run faster.

Building a tiny autograd engine yourself (Part 19) removes the magic for good.

nn.Module

• A module is a container of parameters and submodules with a forward method. Part 7 explains why model(x) works and why super().__init__() and nn.ModuleList matter.
• model.parameters() and model.named_parameters() yield every trainable tensor, recursively.
• model.state_dict() is an ordered dict of names to tensors. Save that, not the whole model object.
• model.train() and model.eval() flip a flag that changes dropout and batch-norm behavior. That is separate from no_grad, and evaluation usually needs both.
• register_buffer stores state that moves with .to() and is saved, but is not trained, such as running statistics or attention masks.

The training loop, annotated

model = MLP(784, 256, 10).to(device)
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
loss_fn = torch.nn.CrossEntropyLoss()        # expects raw logits and integer class labels

for epoch in range(epochs):
    model.train()
    for xb, yb in train_loader:              # DataLoader: batching, shuffling, workers
        xb, yb = xb.to(device), yb.to(device)
        logits = model(xb)                   # forward pass builds the graph
        loss = loss_fn(logits, yb)
        opt.zero_grad(set_to_none=True)      # clear the previous step's gradients
        loss.backward()                      # fill .grad on every parameter
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()                           # update the parameters

    model.eval()
    correct = total = 0
    with torch.no_grad():
        for xb, yb in val_loader:
            xb, yb = xb.to(device), yb.to(device)
            pred = model(xb).argmax(dim=-1)
            correct += (pred == yb).sum().item()
            total += yb.numel()
    print(f"epoch {epoch}: val acc {correct / total:.3f}")

Dataset and DataLoader

• A map-style Dataset implements __len__ and __getitem__(i), returning one example.
• A DataLoader samples indices, calls __getitem__, and collates the examples into batch tensors. For variable-length text, a custom collate_fn pads sequences and builds attention masks.
• num_workers runs loading in separate processes; pin_memory=True speeds up copies to the GPU.
• Iterable-style datasets handle streams too large to index.

Precision, speed and memory

• On modern GPUs, wrap the forward pass and loss in torch.autocast(device_type="cuda", dtype=torch.bfloat16). float16 additionally needs a gradient scaler.
• torch.compile(model) fuses operations into faster kernels. The first call is slow because it compiles.
• Memory levers: smaller batches with gradient accumulation, activation checkpointing, lower precision, frozen layers, and adapters such as LoRA.

Bugs and their signatures

Reading transformer code

Open any Hugging Face model file with this lens: a config class, a module tree (embeddings, then N blocks of attention plus MLP, then an output head), and forward methods with shape comments. Trace one tensor’s shape through one block; everything after that is repetition.

Questions to ask:

• What device and dtype is every tensor on this line?
• Is a graph being built where none is needed?
• Is the model in train mode or eval mode?
• Does this loss function expect logits or probabilities, class indices or one-hot vectors?