PyTorch From First Principles on AIBussin — AI applications, systems and books

What Are We Actually Doing?

Sat, 08 Aug 2026 12:40:00 +0100

It has never been easier to obtain working PyTorch code. Describe a model to an assistant, paste the result into a file, and a training loop will usually appear and usually run. For a large fraction of everyday work this is fine, and this book does not ask you to pretend otherwise.

The difficulty arrives afterwards.

The code runs, but the loss sits at exactly the same value for two hundred steps. Or one tensor is [32, 128] where the next operation wanted [128, 32]. Or the gradient of a weight is None and nothing says why. Or the model trained yesterday and today it does not, with no change you can point to. None of these are syntax errors. They do not produce a helpful traceback, and often they produce no traceback at all. They are questions about what the program is doing, and to answer them you need somewhere to look.

The Tensor: What Is Actually Flowing Through the Loop?

Sat, 08 Aug 2026 12:45:00 +0100

The last program in Chapter 1 trained a two-parameter model on four examples. It worked, and we read it as a story about gradients: predict, measure, differentiate, step. But something in it went unexamined. w was a single number and x held four, and w * x + b produced four predictions without anyone specifying how a scalar and a four-element vector should combine. PyTorch had a rule. We never stated it.

Autograd: What Did PyTorch Record, and Where Does the Gradient Stop?

Sat, 08 Aug 2026 12:55:00 +0100

Here is a training loop with two parameters. It runs, it does not warn, and the loss goes to zero.

import torch

x = torch.tensor([1.0, 2.0, 3.0, 4.0])
target = torch.tensor([3.0, 6.0, 9.0, 12.0])

w1 = torch.tensor(0.5, requires_grad=True)
w2 = torch.tensor(1.0, requires_grad=True)
params = [w1, w2]

feature_log = []

def stage_one(t):
 f = w1 * t
 detached = f.detach() # keep the values for later inspection
 feature_log.append(detached)
 return detached

learning_rate = 0.02

for step in range(400):
 f = stage_one(x)
 prediction = w2 * f
 loss = ((prediction - target) ** 2).mean()

 loss.backward()

 with torch.no_grad():
 for p in params:
 if p.grad is not None:
 p -= learning_rate * p.grad
 p.grad.zero_()

 if step in (0, 9, 99, 399):
 print(f"step={step:3d} loss={loss.item():.6f} "
 f"w1={w1.item():.4f} w2={w2.item():.4f}")

step= 0 loss=46.875000 w1=0.5000 w2=1.3750
step= 9 loss=11.521060 w1=0.5000 w2=3.7071
step= 99 loss=0.000009 w1=0.5000 w2=5.9979
step=399 loss=0.000000 w1=0.5000 w2=6.0000

Read the columns the way Chapter 1 taught. The loss falls to zero, which is the outcome we asked for. And w1 never moves. Not slowly, not by a little: it holds 0.5000 for four hundred steps while w2 climbs to exactly 6.0, which is the value that makes w1 * w2 equal to the 3 we were trying to learn. Half the model was frozen, the other half absorbed the entire job, and the objective was satisfied anyway.

The Network: What Is It Without nn.Module?

Sat, 08 Aug 2026 12:59:00 +0100

Here is a two-layer neural network, written with nothing but the machinery of the last three chapters. Four trainable tensors, a forward pass, a loss, backward(), an update. It runs without warnings, and the loss falls by a third.

import math
import torch

torch.manual_seed(0)

n = 2000
X = torch.randn(n, 2)
y = ((X[:, 0] > 0) ^ (X[:, 1] > 0)).long()

perm = torch.randperm(n)
n_train = int(0.8 * n)
X_train, y_train = X[perm[:n_train]], y[perm[:n_train]]
X_val, y_val = X[perm[n_train:]], y[perm[n_train:]]

D_in, H, C = 2, 16, 2

torch.manual_seed(1)
W1 = (torch.randn(D_in, H) / math.sqrt(D_in)).requires_grad_()
b1 = torch.zeros(H, requires_grad=True)
W2 = (torch.randn(H, C) / math.sqrt(H)).requires_grad_()
b2 = torch.zeros(C, requires_grad=True)
parameters = [W1, b1, W2, b2]

def forward(x):
 z1 = x @ W1 + b1
 return z1 @ W2 + b2

def cross_entropy(logits, targets):
 log_probs = torch.log_softmax(logits, dim=1)
 return -log_probs[torch.arange(targets.shape[0]), targets].mean()

def accuracy(logits, targets):
 return (logits.argmax(dim=1) == targets).float().mean().item()

learning_rate = 0.5

# First verify that one training step is mechanically healthy.
names = ["W1", "b1", "W2", "b2"]
initial_values = [p.detach().clone() for p in parameters]

loss = cross_entropy(forward(X_train), y_train)
before = [p.detach().clone() for p in parameters]

loss.backward()

with torch.no_grad():
 for p in parameters:
 p -= learning_rate * p.grad

for name, p, old in zip(names, parameters, before):
 print(f"{name:3s} leaf={p.is_leaf} grad_none={p.grad is None} "
 f"grad_norm={p.grad.norm().item():8.4f} "
 f"finite={torch.isfinite(p.grad).all().item()} "
 f"moved={(p.detach() - old).abs().max().item():.4f}")

W1 leaf=True grad_none=False grad_norm= 0.4516 finite=True moved=0.1048
b1 leaf=True grad_none=False grad_norm= 0.0192 finite=True moved=0.0053
W2 leaf=True grad_none=False grad_norm= 1.1058 finite=True moved=0.1918
b2 leaf=True grad_none=False grad_norm= 0.0175 finite=True moved=0.0062

Every parameter is a leaf. Every parameter requires gradients. Every parameter receives a finite gradient, and every parameter changes by a measurable amount when the update runs. The basic training machinery is working.

nn.Module: What Does PyTorch Think Belongs to Your Model?

Sat, 08 Aug 2026 13:09:00 +0100

Here is a classifier for the XOR-shaped dataset from Chapter 4, written the way most people write PyTorch. It trains. Validation accuracy climbs from chance to 95.7%. It saves a checkpoint, the checkpoint loads without a warning, and then the reloaded model performs at chance.

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)

n = 2000
X = torch.randn(n, 2)
y = ((X[:, 0] > 0) ^ (X[:, 1] > 0)).long()

perm = torch.randperm(n)
X_train, y_train = X[perm[:1600]], y[perm[:1600]]
X_val, y_val = X[perm[1600:]], y[perm[1600:]]

class Encoder(nn.Module):
 def __init__(self, width=16, depth=3):
 super().__init__()
 self.blocks = [nn.Linear(width, width) for _ in range(depth)]

 def forward(self, x):
 for block in self.blocks:
 x = torch.relu(block(x))
 return x

class Classifier(nn.Module):
 def __init__(self, width=16, depth=3):
 super().__init__()
 self.stem = nn.Linear(2, width)
 self.encoder = Encoder(width, depth)
 self.head = nn.Linear(width, 2)

 def forward(self, x):
 x = torch.relu(self.stem(x))
 x = self.encoder(x)
 return self.head(x)

model = Classifier()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)

for epoch in range(200):
 logits = model(X_train)
 loss = F.cross_entropy(logits, y_train)

 optimizer.zero_grad()
 loss.backward()
 optimizer.step()

 if epoch % 50 == 0 or epoch == 199:
 with torch.no_grad():
 val_acc = (model(X_val).argmax(1) == y_val).float().mean()
 print(f"epoch={epoch:3d} loss={loss.item():.4f} "
 f"val_acc={val_acc.item():.3f}")

torch.save(model.state_dict(), "classifier.pt")

reloaded = Classifier()
print(reloaded.load_state_dict(torch.load("classifier.pt", weights_only=True)))

with torch.no_grad():
 acc = (reloaded(X_val).argmax(1) == y_val).float().mean()
print(f"reloaded val_acc={acc.item():.3f}")

epoch= 0 loss=0.6964 val_acc=0.512
epoch= 50 loss=0.5930 val_acc=0.810
epoch=100 loss=0.3856 val_acc=0.873
epoch=150 loss=0.2550 val_acc=0.925
epoch=199 loss=0.1908 val_acc=0.957
<All keys matched successfully>
reloaded val_acc=0.512

Read the last two lines together. PyTorch reports that every key in the checkpoint matched every key the model expected, and the model that came back is worthless. Nothing raised. Nothing warned. load_state_dict was as satisfied as it is capable of being.

DataLoader: Where Is the Training Loop Actually Waiting?

Sat, 08 Aug 2026 13:21:00 +0100

Here are two controlled training pipelines. They use the same batch size, the same machine, and the same synthetic post-batch workload. Each sample also carries the same nominal two-millisecond cost: in one pipeline that cost is waiting, while in the other it is fixed CPU work.

The tensor construction around that controlled cost is the same in both cases. What changes is the resource those two milliseconds consume.

Both are given the same treatment — raise num_workers from 0 to 8 — and measured the same way.

Transforms: What Does the Model Actually See?

Sat, 29 Aug 2026 12:00:00 +0100

Here are two training runs. Same model, same data, same optimizer, same seed, same number of epochs. The preprocessing pipelines differ by one keyword argument, and neither one raises.

scale=True range=[ -1.09, 1.45] first_epoch_loss= 0.5222 last=0.0032 val_acc=1.000
scale=False range=[ 259.89, 886.36] first_epoch_loss=28.5556 last=0.6938 val_acc=0.400

The second pipeline produced tensors with the right shape, the right dtype, the right device, and no non-finite values. Every check from Chapter 2 passes. Every check from Chapter 5 passes. The batch arrived on time, which is everything Chapter 6 taught us to verify. And the model went from perfect to worse than guessing.

CNN Geometry: What Shape Reaches the Next Layer?

Fri, 28 Aug 2026 10:15:00 +0100

Here is a batch of 32 RGB images and an ordinary first convolution.

x = torch.randn(32, 224, 224, 3)
conv = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3)
conv(x)

RuntimeError: Given groups=1, weight of size [16, 3, 3, 3],
expected input[32, 224, 224, 3] to have 3 channels, but got 224 channels instead

The images have three channels. PyTorch says there are 224. The message names a number that appears nowhere in the model definition, so the obvious reading is that the layer was declared wrong, and the obvious repair is to declare it right:

Feature Space: What Does a Linear Model Actually See?

Fri, 28 Aug 2026 15:30:00 +0100

Here is a sequence classifier. Eight examples, 128 positions each, 768 features per position, and a linear head that produces one score.

x = torch.randn(8, 128, 768)
classifier = nn.Linear(768, 1)

scores = classifier(x)

input (8, 128, 768)
Linear(768,1) -> (8, 128, 1)

One thousand and twenty-four scores where eight were expected. Nothing raised, nothing is non-finite, and the shape is entirely predictable once you know the rule. The layer did exactly what the tensor asked of it: nn.Linear transforms the last axis and preserves every axis before it, so it produced one score for every (example, position) pair — 128 scores per example, computed independently.

Attention: Which Position Is Comparing With Which?

Sat, 29 Aug 2026 10:00:00 +0100

Here is a tensor of sequence representations and the line that splits it into heads.

B, T, E, Nh = 1, 4, 8, 2
Dh = E // Nh # 4

x = torch.arange(B * T * E).reshape(B, T, E)
heads = x.reshape(B, Nh, T, Dh)

x: (1, 4, 8)
heads: (1, 2, 4, 4)

That is exactly the shape multi-head attention wants: batch, heads, positions, head dimension. Nothing raised. Every shape assertion passes.

Now the same split written the other way:

Training: Which Link in the Learning Chain Is Broken?

Sat, 29 Aug 2026 14:00:00 +0100

Here is an ordinary piece of transfer-learning code. A small model, a frozen feature extractor, an optimizer, and a new classification head for the new task.

model = TinyClassifier()

for p in model.features.parameters():
 p.requires_grad_(False)

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-2, weight_decay=0.0)

model.head = nn.Linear(32, 2) # adaptation, performed later

Train it for two hundred steps on a balanced, perfectly aligned, entirely learnable binary task:

step= 0 loss=0.683606 acc=0.7144
step= 50 loss=0.683606 acc=0.7144
step= 100 loss=0.683606 acc=0.7144
step= 150 loss=0.683606 acc=0.7144
step= 200 loss=0.683606 acc=0.7144

Not approximately constant. Identical to six decimal places, two hundred times.

Performance: What Is the Machine Waiting For?

Sat, 29 Aug 2026 16:00:00 +0100

Here is a benchmark. A model, a batch, a loop, a timer. Nothing exotic.

model.train()
t0 = time.perf_counter()
for _ in range(20):
 optimizer.zero_grad(set_to_none=True)
 loss = F.cross_entropy(model(x), y) # MLP, batch 256
 loss.backward()
 optimizer.step()
ms_per_step = (time.perf_counter() - t0) / 20 * 1000

naive loop of 20: 9.491 ms/step reported + 132.0 ms still queued at the final sync -> really 16.093 ms/step
naive loop of 100: 14.692 ms/step reported + 138.8 ms still queued at the final sync -> really 16.080 ms/step
naive loop of 400: 15.741 ms/step reported + 132.1 ms still queued at the final sync -> really 16.071 ms/step

synchronized + warmed, median of 400 : 16.430 ms/step <- the real number
one step, host-visible return only : 4.866 ms/step
first training step ever : 20.890 ms

Read the first line again. The loop ran twenty iterations and reported 9.491 ms per step. Then a single torch.cuda.synchronize() immediately afterward blocked for 132 more milliseconds — roughly eight more steps’ worth of work that the GPU had not finished when the timer stopped. Loop over a hundred iterations instead and the same code reports 14.7. Over four hundred, 15.7. Time one step in isolation and it looks like 4.9.

Compilation: Which Assumption Stopped Holding?

Sat, 29 Aug 2026 17:00:00 +0100

Here is an inference service. It loads a model, compiles it once, and serves requests whose sequence length varies from request to request — the ordinary situation for anything that takes text.

model = load_model().eval()
compiled = torch.compile(model, dynamic=False) # "dynamic shapes are slow, force static"

for request in stream:
 with torch.no_grad():
 output = compiled(request.tokens) # [1, T], T varies per request

Thirteen requests. Every compiled output agrees with the eager model to a maximum absolute difference of 7e-7. Nothing raises.

Regressions: Did the Model Change, or the Measurement?

Sat, 29 Aug 2026 18:00:00 +0100

Here is a refactor. The old training script set the seed at the top; the new one moves that line below model construction, where it reads more naturally, and draws the batch-shuffle permutation from a local torch.Generator() instead of the global RNG, so the shuffling is “self-contained.” The training procedure is mathematically the same and the model architecture is unchanged. But the concrete initialized parameters are no longer coupled to the old run, because model construction now consumes randomness before manual_seed is applied. Run it once:

Assembly: A Language Model You Can Interrogate

Sun, 30 Aug 2026 10:00:00 +0100

Here are two training runs of the same small language model on the same data. The only difference is one line in the batching function.

CORRECT y = data[i+1 : i+block+1] the next-token shift
MISALIGNED y = data[i : i+block] no shift

step correct val misaligned val
 0 3.6695 2.6469
 250 1.1038 0.0007
 500 0.4730 0.0004
 750 0.3940 0.0003
1000 0.3690 0.0002
1250 0.3581 0.0002
1500 0.3558 0.0001

The misaligned run’s loss falls to 0.0001 — three thousand times lower than the correct run’s 0.36. By every number on the dashboard it is the best training run you have ever seen. Nothing raised. The shapes are identical: x is [B, T], y is [B, T], the loss is a finite scalar with a grad_fn.

Appendix A: PyTorch Diagnostic Field Guide

Sun, 30 Aug 2026 12:00:00 +0100

This appendix introduces no new PyTorch mechanism.

Everything here was earned earlier in the book by building something small, breaking it deliberately, measuring what changed, and locating the first place where reality stopped matching the intended computation.

The purpose of this appendix is different.

When a real model is failing, you usually do not need another explanation of autograd, broadcasting, attention or CUDA. You need to answer a narrower question: