PyTorch Model Not Learning? A Systematic Debugging Guide
PyTorch: Zero to Hero — Step 08
Your model runs.
The loss is finite.
Nothing crashes.
And it still does not learn.
This is one of the most frustrating states in machine learning because there is no stack trace telling you what is wrong.
The program is valid Python.
The tensors have legal shapes.
The GPU is busy.
The optimizer is stepping.
And the model is useless.
This post is a systematic way to debug that situation.
Why so much debugging in this series?
There is a deliberate reason this series keeps returning to debugging.
In our view, more and more of the code you write is going to be generated by an LLM.
The model can generate the nn.Module.
It can generate the training loop.
It can generate the optimizer setup, DataLoader, checkpointing code, transforms, attention blocks and evaluation loop.
And when all of that works, that is useful.
But when it does not work, you eventually reach the point where asking for another generated version of the same code stops helping.
That is where you need to get underneath the generated code.
You need to inspect:
inputs
activations
loss
computation graph
gradients
parameters
optimizer state
updates
validation behaviour
An LLM can suggest where to look.
But the debugger, the profiler, the assertions and the numbers running on your machine are the ground truth.
That is why this series contains so much debugging.
If AI writes most of the boilerplate, then understanding how to interrogate a broken system becomes more valuable, not less.
The debugging ladder
When a model is not learning, do not immediately change the architecture.
Work downward through the system.
flowchart TD
A[Model not learning] --> B[1. Is the data correct?]
B -->|Yes| C[2. Are targets correct?]
B -->|No| BFIX[Fix data pipeline]
C -->|Yes| D[3. Do outputs look sane?]
C -->|No| CFIX[Fix labels / alignment]
D -->|Yes| E[4. Is the loss correct?]
D -->|No| DFIX[Fix model output scale]
E -->|Yes| F[5. Is autograd connected?]
E -->|No| EFIX[Fix loss function]
F -->|Yes| G[6. Are gradients finite and non-zero?]
F -->|No| GFIX[Find graph break]
G -->|Yes| H[7. Optimizer holds right params?]
G -->|No| HFIX[Gradient problem]
H -->|Yes| I[8. Does optimizer.step change params?]
H -->|No| IFIX[Fix optimizer construction]
I -->|Yes| J[9. Can model overfit tiny batch?]
I -->|No| JFIX[Update rule broken]
J -->|Yes| K[10. Architecture / hyperparameters]
J -->|No| KFIX[Training system suspect]
That ordering matters.
A surprising amount of wasted ML debugging comes from tuning learning rates against a broken target pipeline.
A small model we can deliberately break
We will use a simple binary classification problem.
import torch
import torch.nn as nn
import torch.nn.functional as F
torch.manual_seed(42)
def make_data(n=2048):
x = torch.randn(n, 2)
y = ((x[:, 0] + 0.75 * x[:, 1]) > 0).long()
return x, y
class TinyClassifier(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(2, 32),
nn.ReLU(),
nn.Linear(32, 32),
nn.ReLU(),
nn.Linear(32, 2),
)
def forward(self, x):
return self.net(x)
A minimal training loop:
x, y = make_data()
model = TinyClassifier()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
for step in range(500):
optimizer.zero_grad(set_to_none=True)
logits = model(x)
loss = F.cross_entropy(logits, y)
loss.backward()
optimizer.step()
if step % 50 == 0:
accuracy = (logits.argmax(dim=-1) == y).float().mean()
print(
f"step={step:04d} "
f"loss={loss.item():.4f} "
f"acc={accuracy.item():.4f}"
)
This should learn quickly.
Now we can use it as a controlled system for debugging.
1. Inspect the data before the model
The model cannot recover from bad inputs.
Start by printing what is actually entering the network.
def inspect_batch(x, y):
print("x.shape:", x.shape)
print("x.dtype:", x.dtype)
print("x.device:", x.device)
print("x.min:", x.min().item())
print("x.max:", x.max().item())
print("x.mean:", x.mean().item())
print("x.std:", x.std().item())
print("y.shape:", y.shape)
print("y.dtype:", y.dtype)
print("y.device:", y.device)
print("y.unique:", torch.unique(y, return_counts=True))
inspect_batch(x, y)
For image data you might want:
print(images.shape)
print(images.dtype)
print(images.min().item(), images.max().item())
print(images.mean().item(), images.std().item())
For token data:
print(input_ids.shape)
print(input_ids.min().item())
print(input_ids.max().item())
print((input_ids == pad_token_id).float().mean().item())
A common failure is perfectly valid tensors containing nonsense.
Examples:
images scaled to 0..255 when the model expects 0..1
labels shifted by one class
token IDs outside the expected vocabulary mapping
all labels accidentally equal to zero
training and target arrays shuffled independently
normalization applied twice
padding dominating the sequence
The model will happily optimize against bad data.
Assert your data contract
Do not rely on visual inspection alone.
def assert_classification_batch(x, y, num_classes):
assert x.ndim >= 2
assert y.ndim == 1
assert x.shape[0] == y.shape[0]
assert torch.isfinite(x).all(), "non-finite input"
assert y.dtype == torch.long
assert y.min().item() >= 0
assert y.max().item() < num_classes
Use it:
assert_classification_batch(x, y, num_classes=2)
Generated code benefits enormously from explicit contracts.
The assertion is not glamorous.
It is also much more trustworthy than assuming the generated data pipeline is correct.
2. Establish a stupid baseline
Before asking whether the neural network is good, ask whether it is better than doing nothing.
For classification:
counts = torch.bincount(y)
majority_accuracy = counts.max().float() / counts.sum()
print("majority baseline:", majority_accuracy.item())
For regression:
mean_target = y.float().mean()
baseline_mse = ((y.float() - mean_target) ** 2).mean()
print("mean baseline MSE:", baseline_mse.item())
If your network achieves 90% accuracy on a dataset where one class is 90% of the examples, you have not necessarily learned anything useful.
3. Inspect the raw model outputs
Before the loss, inspect the logits.
with torch.no_grad():
logits = model(x[:16])
print(logits)
print("shape:", logits.shape)
print("min:", logits.min().item())
print("max:", logits.max().item())
print("mean:", logits.mean().item())
print("std:", logits.std().item())
A useful helper:
def tensor_stats(name, t):
t = t.detach()
print(
f"{name:20s} "
f"shape={tuple(t.shape)} "
f"mean={t.float().mean().item():+.4e} "
f"std={t.float().std().item():.4e} "
f"min={t.min().item():+.4e} "
f"max={t.max().item():+.4e} "
f"finite={torch.isfinite(t).all().item()}"
)
Then:
logits = model(x[:64])
tensor_stats("logits", logits)
Warning signs:
all outputs identical
all outputs exactly zero
magnitudes exploding immediately
NaN or Inf
one logit already enormous before training
4. Verify the loss independently
A loss function is code too.
It can be wrong.
For classification, remember that cross_entropy expects logits, not probabilities.
Correct:
logits = model(x)
loss = F.cross_entropy(logits, y)
Usually wrong:
probs = logits.softmax(dim=-1)
loss = F.cross_entropy(probs, y)
F.cross_entropy already incorporates the appropriate log-softmax operation internally.
Check one example manually
logits = model(x[:1])
target = y[:1]
loss = F.cross_entropy(logits, target)
log_probs = F.log_softmax(logits, dim=-1)
manual = -log_probs[0, target.item()]
print(loss.item())
print(manual.item())
print(torch.allclose(loss, manual))
When debugging, reducing an abstraction to one example is incredibly powerful.
5. Check whether the loss is connected to the graph
logits = model(x[:32])
loss = F.cross_entropy(logits, y[:32])
print("loss.requires_grad:", loss.requires_grad)
print("loss.grad_fn:", loss.grad_fn)
Expected:
loss.requires_grad: True
loss.grad_fn: <...>
If requires_grad is false, something broke the graph.
Common graph breaks include:
x = x.detach()
value = tensor.item()
array = tensor.cpu().numpy()
new_tensor = torch.tensor(existing_tensor)
Not all of those are wrong in general.
They are wrong if you expected gradients to flow through them.
6. Audit parameter gradients
After backward:
optimizer.zero_grad(set_to_none=True)
logits = model(x[:128])
loss = F.cross_entropy(logits, y[:128])
loss.backward()
Inspect every parameter:
def gradient_report(model):
for name, p in model.named_parameters():
if p.grad is None:
print(f"{name:30s} grad=None")
continue
g = p.grad.detach()
print(
f"{name:30s} "
f"shape={tuple(p.shape)!s:18s} "
f"norm={g.norm().item():.4e} "
f"mean={g.mean().item():+.4e} "
f"max={g.abs().max().item():.4e} "
f"finite={torch.isfinite(g).all().item()}"
)
gradient_report(model)
This distinguishes several very different problems.
grad is None
The parameter is not connected to the loss, is frozen, or never participated in the forward path.
Gradient norm is exactly zero
The graph exists, but the local derivative may be zero.
Gradient is tiny everywhere
Possible vanishing-gradient or scale problem.
Gradient is enormous
Possible exploding-gradient problem.
Gradient contains NaN/Inf
Numerical instability has already occurred.
7. Calculate the total gradient norm
def total_grad_norm(model):
grads = [
p.grad.detach().norm(2)
for p in model.parameters()
if p.grad is not None
]
if not grads:
return torch.tensor(0.0)
return torch.stack(grads).norm(2)
print("grad norm:", total_grad_norm(model).item())
Track it during training:
for step in range(100):
optimizer.zero_grad(set_to_none=True)
logits = model(x)
loss = F.cross_entropy(logits, y)
loss.backward()
grad_norm = total_grad_norm(model)
optimizer.step()
print(
f"step={step:03d} "
f"loss={loss.item():.4f} "
f"grad_norm={grad_norm.item():.4e}"
)
The absolute value is architecture-dependent.
The pattern matters.
8. Prove that optimizer.step() changes the model
Do not assume it does.
Snapshot the parameters.
def snapshot_parameters(model):
return {
name: p.detach().clone()
for name, p in model.named_parameters()
}
def parameter_delta_report(model, before):
for name, p in model.named_parameters():
delta = (p.detach() - before[name]).norm().item()
print(f"{name:30s} delta={delta:.6e}")
Use it around a single step:
optimizer.zero_grad(set_to_none=True)
logits = model(x[:128])
loss = F.cross_entropy(logits, y[:128])
loss.backward()
before = snapshot_parameters(model)
optimizer.step()
parameter_delta_report(model, before)
If every delta is zero, stop tuning the learning rate.
Your optimizer is not updating the model.
Verify optimizer membership
def optimizer_parameter_ids(optimizer):
return {
id(p)
for group in optimizer.param_groups
for p in group["params"]
}
def report_optimizer_membership(model, optimizer):
optimizer_ids = optimizer_parameter_ids(optimizer)
for name, p in model.named_parameters():
print(
f"{name:30s} "
f"requires_grad={str(p.requires_grad):5s} "
f"in_optimizer={id(p) in optimizer_ids}"
)
A parameter can have gradients and still not be in the optimizer.
That bug is particularly nasty because backward appears completely healthy.
9. The tiny-batch overfit test
This is one of the highest-value debugging techniques in deep learning.
Take a tiny number of examples.
Try to memorize them.
x_small = x[:32]
y_small = y[:32]
model = TinyClassifier()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-2)
for step in range(2000):
optimizer.zero_grad(set_to_none=True)
logits = model(x_small)
loss = F.cross_entropy(logits, y_small)
loss.backward()
optimizer.step()
if step % 100 == 0:
acc = (logits.argmax(dim=-1) == y_small).float().mean()
print(step, loss.item(), acc.item())
For a sufficiently expressive model, you should usually be able to drive training accuracy on a tiny fixed batch extremely high.
If you cannot memorize 32 examples, generalization is not your immediate problem.
The training system itself is suspect.
You can visualise the overfit test clearly:
import matplotlib.pyplot as plt
x_small = x[:32]
y_small = y[:32]
model = TinyClassifier()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-2)
losses = []
for step in range(2000):
optimizer.zero_grad(set_to_none=True)
logits = model(x_small)
loss = F.cross_entropy(logits, y_small)
loss.backward()
optimizer.step()
losses.append(loss.item())
plt.plot(losses)
plt.xlabel('Step')
plt.ylabel('Loss')
plt.title('Tiny-batch overfit test (should plummet)')
plt.grid(True)
plt.show()
A loss curve that refuses to drop is a loud signal that something fundamental is broken.
Why the tiny-batch test is so useful
It removes a huge number of variables.
You do not care about:
generalization
validation splits
dataset size
data-loader throughput
regularization
augmentation quality
You are asking one simpler question:
Can this model and this optimizer reduce this loss on these exact examples?
If no, debug locally.
10. Disable regularization while debugging
Regularization makes sense once the core system learns.
During debugging it can hide the signal.
Temporarily simplify:
drop dropout
remove augmentation
reduce weight decay
remove label smoothing
remove stochastic depth
remove complex schedulers
For example:
optimizer = torch.optim.AdamW(
model.parameters(),
lr=1e-3,
weight_decay=0.0,
)
The goal is not to produce the final training recipe.
The goal is to create the smallest system that should obviously learn.
11. Check train/eval mode
model.train() and model.eval() change the behavior of modules such as dropout and batch normalization.
They do not enable or disable gradients.
A clean training/evaluation structure:
model.train()
for x_batch, y_batch in train_loader:
optimizer.zero_grad(set_to_none=True)
logits = model(x_batch)
loss = F.cross_entropy(logits, y_batch)
loss.backward()
optimizer.step()
model.eval()
with torch.no_grad():
for x_batch, y_batch in val_loader:
logits = model(x_batch)
Do not confuse:
model.eval()
with:
with torch.no_grad():
They solve different problems.
12. Inspect activation statistics layer by layer
A model may receive gradients and still have pathological activations.
Forward hooks are useful here.
def activation_hook(name):
def hook(module, inputs, output):
if not torch.is_tensor(output):
return
out = output.detach()
print(
f"{name:30s} "
f"shape={tuple(out.shape)!s:18s} "
f"mean={out.float().mean().item():+.3e} "
f"std={out.float().std().item():.3e} "
f"min={out.min().item():+.3e} "
f"max={out.max().item():+.3e} "
f"finite={torch.isfinite(out).all().item()}"
)
return hook
handles = []
for name, module in model.named_modules():
if isinstance(module, (nn.Linear, nn.Conv2d, nn.ReLU)):
handles.append(
module.register_forward_hook(activation_hook(name))
)
_ = model(x[:16])
for handle in handles:
handle.remove()
Potential problems:
activations explode with depth
activations collapse to nearly zero
ReLUs are almost always zero
non-finite values first appear in one layer
variance grows dramatically every layer
13. Measure dead ReLUs
def zero_fraction(t):
return (t == 0).float().mean().item()
A hook:
def relu_hook(name):
def hook(module, inputs, output):
print(name, "zero_fraction=", zero_fraction(output.detach()))
return hook
If nearly every activation is zero, the downstream model receives little useful signal.
14. Detect NaN and Inf at the source
Do not merely notice that the final loss is NaN.
Find the first operation that becomes invalid.
Start with explicit checks:
def assert_finite(name, t):
if not torch.isfinite(t).all():
raise RuntimeError(f"{name} contains NaN or Inf")
Use them between stages:
h = model.net[0](x)
assert_finite("linear_0", h)
h = F.relu(h)
assert_finite("relu_0", h)
For difficult backward problems, PyTorch anomaly detection can provide the traceback of the forward operation associated with a failing backward computation.
with torch.autograd.detect_anomaly():
logits = model(x)
loss = F.cross_entropy(logits, y)
loss.backward()
Do not leave anomaly detection enabled for normal training.
It is a debugging tool and adds overhead.
15. Common sources of NaN loss
Log of zero
Bad:
loss = -torch.log(probability)
If probability == 0, the result is infinite.
Safer patterns often work in logit space instead.
Division by zero
normalized = x / x.norm(dim=-1, keepdim=True)
Safer:
normalized = x / x.norm(dim=-1, keepdim=True).clamp_min(1e-8)
Exponentials of large values
weights = torch.exp(scores)
Use numerically stable formulations when possible.
For classification, prefer:
F.cross_entropy(logits, targets)
over implementing softmax and log manually.
Learning rate too high
One step may send parameters into an unstable region.
Track both loss and gradient norm.
16. Gradient clipping: useful, but not a diagnosis
PyTorch provides:
torch.nn.utils.clip_grad_norm_(
model.parameters(),
max_norm=1.0,
)
A standard pattern:
optimizer.zero_grad(set_to_none=True)
logits = model(x_batch)
loss = loss_fn(logits, y_batch)
loss.backward()
grad_norm = torch.nn.utils.clip_grad_norm_(
model.parameters(),
max_norm=1.0,
error_if_nonfinite=True,
)
optimizer.step()
But clipping can conceal a deeper problem.
If the gradient norm jumps from 1 to 1e12, do not merely celebrate that clipping prevented the crash.
Find out why it happened.
17. Learning-rate sanity test
When the training system is known to work, learning rate becomes a meaningful suspect.
Try an order-of-magnitude sweep.
learning_rates = [
1e-5,
3e-5,
1e-4,
3e-4,
1e-3,
3e-3,
1e-2,
]
For each, train from the same initial state for a small number of steps.
import copy
base_model = TinyClassifier()
base_state = copy.deepcopy(base_model.state_dict())
for lr in learning_rates:
model = TinyClassifier()
model.load_state_dict(base_state)
optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
for _ in range(100):
optimizer.zero_grad(set_to_none=True)
logits = model(x)
loss = F.cross_entropy(logits, y)
loss.backward()
optimizer.step()
print(lr, loss.item())
Do not compare different random initializations and call it a learning-rate comparison.
18. Check for accidental target leakage — and accidental target destruction
Both can happen.
Leakage example:
features = torch.cat([features, target.float().unsqueeze(-1)], dim=-1)
The model appears brilliant because the answer is in the input.
Target destruction example:
x = x[perm]
# forgot: y = y[perm]
Now inputs and targets no longer correspond.
A model may converge to noise-level performance forever.
19. Verify input/target alignment explicitly
For synthetic data this is easy.
expected = ((x[:, 0] + 0.75 * x[:, 1]) > 0).long()
print((expected == y).float().mean())
For real datasets, build a debugging view that prints or visualizes a few (input, target) pairs.
This is especially important for generated data pipelines.
A syntactically correct transform can silently misalign your labels.
20. Track training and validation separately
Different patterns imply different problems.
train loss flat, val loss flat
-> optimization/system problem
train loss falls, val loss flat
-> generalization/distribution problem
train loss falls, val loss rises
-> overfitting
train and val both NaN
-> numerical instability
train accuracy high, val accuracy suspiciously high
-> possible leakage
A simple evaluation function:
@torch.no_grad()
def evaluate(model, x, y):
model.eval()
logits = model(x)
loss = F.cross_entropy(logits, y)
accuracy = (logits.argmax(dim=-1) == y).float().mean()
return loss.item(), accuracy.item()
21. Reproducibility is a debugging tool
When debugging, you want failures to be repeatable.
At minimum:
import random
import numpy as np
import torch
seed = 1234
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
PyTorch does not guarantee identical results across every release, platform, device and execution mode.
But reducing nondeterminism on one environment can make development and regression debugging dramatically easier.
Save the failing batch
If a failure happens only occasionally, preserve the evidence.
if not torch.isfinite(loss):
torch.save(
{
"x": x_batch.detach().cpu(),
"y": y_batch.detach().cpu(),
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
"step": step,
},
"failing_batch.pt",
)
raise RuntimeError("non-finite loss")
Now the failure can become a deterministic test case.
That is much better than trying to reproduce it from memory.
22. Debug mixed precision separately
Mixed precision adds another layer of state.
If training is unstable under AMP, first establish whether full precision works.
A typical CUDA mixed-precision loop looks like:
scaler = torch.amp.GradScaler("cuda")
for x_batch, y_batch in train_loader:
optimizer.zero_grad(set_to_none=True)
with torch.autocast(
device_type="cuda",
dtype=torch.float16,
):
logits = model(x_batch)
loss = F.cross_entropy(logits, y_batch)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
If you need to inspect or clip the real gradients, unscale first.
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
grad_norm = torch.nn.utils.clip_grad_norm_(
model.parameters(),
max_norm=1.0,
)
scaler.step(optimizer)
scaler.update()
Debugging scaled gradients as if they were ordinary gradients will mislead you.
23. Check whether the scheduler is destroying the learning rate
Print it.
print(optimizer.param_groups[0]["lr"])
During training:
for step in range(num_steps):
# train...
lr = optimizer.param_groups[0]["lr"]
print(step, lr)
Common mistakes include:
scheduler stepped at the wrong frequency
scheduler stepped before optimizer unexpectedly
warmup never leaves near-zero learning rates
learning rate decays to almost zero too early
checkpoint restores optimizer but not scheduler state
Again: inspect the number rather than reasoning from the configuration file.
24. Audit parameter scale
def parameter_report(model):
for name, p in model.named_parameters():
d = p.detach()
print(
f"{name:30s} "
f"norm={d.norm().item():.4e} "
f"mean={d.mean().item():+.4e} "
f"std={d.std().item():.4e} "
f"max={d.abs().max().item():.4e}"
)
Look for:
one parameter tensor exploding
weights collapsing to zero
unexpectedly enormous embeddings
bias terms dwarfing weight matrices
Parameter statistics are especially useful when comparing a healthy checkpoint with a broken one.
25. Compare two runs programmatically
Suppose one configuration works and one does not.
Do not eyeball two logs.
Record structured metrics.
history.append(
{
"step": step,
"loss": loss.item(),
"grad_norm": grad_norm.item(),
"lr": optimizer.param_groups[0]["lr"],
"accuracy": accuracy.item(),
}
)
Then compare:
when did the trajectories diverge?
which metric diverged first?
did gradient norm diverge before loss?
did learning rate diverge before gradient norm?
The first divergence is usually more informative than the eventual crash.
26. A reusable training-step debugger
Here is a compact helper you can drop into a project.
class TrainingStepDebugger:
def __init__(self, model, optimizer):
self.model = model
self.optimizer = optimizer
def optimizer_ids(self):
return {
id(p)
for group in self.optimizer.param_groups
for p in group["params"]
}
def inspect_parameters(self):
optimizer_ids = self.optimizer_ids()
for name, p in self.model.named_parameters():
grad = p.grad
if grad is None:
grad_norm = None
finite = None
else:
grad_norm = grad.detach().norm().item()
finite = torch.isfinite(grad).all().item()
print(
name,
{
"shape": tuple(p.shape),
"requires_grad": p.requires_grad,
"in_optimizer": id(p) in optimizer_ids,
"grad_norm": grad_norm,
"grad_finite": finite,
},
)
def snapshot(self):
return {
name: p.detach().clone()
for name, p in self.model.named_parameters()
}
def report_updates(self, before):
for name, p in self.model.named_parameters():
delta = (p.detach() - before[name]).norm().item()
print(name, "delta=", delta)
Use it:
debugger = TrainingStepDebugger(model, optimizer)
optimizer.zero_grad(set_to_none=True)
logits = model(x_batch)
loss = loss_fn(logits, y_batch)
loss.backward()
debugger.inspect_parameters()
before = debugger.snapshot()
optimizer.step()
debugger.report_updates(before)
That gives you a direct answer to:
which parameters exist?
which require gradients?
which got gradients?
are those gradients finite?
which parameters are owned by the optimizer?
which actually changed?
27. Build a fail-fast training loop
A production-ish training loop should catch corruption early.
def train_step(model, optimizer, x, y):
model.train()
optimizer.zero_grad(set_to_none=True)
assert torch.isfinite(x).all()
logits = model(x)
assert torch.isfinite(logits).all()
loss = F.cross_entropy(logits, y)
if not torch.isfinite(loss):
raise RuntimeError(f"non-finite loss: {loss.item()}")
loss.backward()
grad_norm = torch.nn.utils.clip_grad_norm_(
model.parameters(),
max_norm=10.0,
error_if_nonfinite=True,
)
optimizer.step()
return {
"loss": loss.detach().item(),
"grad_norm": grad_norm.detach().item(),
"lr": optimizer.param_groups[0]["lr"],
}
This does not guarantee correctness.
It converts several silent failures into immediate failures.
That is a major improvement.
28. The one-batch forensic script
When a generated training pipeline is behaving strangely, reduce it to one batch.
x_batch, y_batch = next(iter(train_loader))
print("=== INPUT ===")
inspect_batch(x_batch, y_batch)
print("\n=== PARAMETERS BEFORE ===")
parameter_report(model)
optimizer.zero_grad(set_to_none=True)
print("\n=== FORWARD ===")
logits = model(x_batch)
tensor_stats("logits", logits)
loss = loss_fn(logits, y_batch)
print("loss:", loss.item())
print("loss.requires_grad:", loss.requires_grad)
print("loss.grad_fn:", loss.grad_fn)
print("\n=== BACKWARD ===")
loss.backward()
gradient_report(model)
print("\n=== STEP ===")
before = snapshot_parameters(model)
optimizer.step()
parameter_delta_report(model, before)
This is the ML equivalent of reducing a failing program to the smallest reproducible example.
29. Questions to ask in order
When the model does not learn, use this list.
DATA
[ ] Are the input values sane?
[ ] Are the labels sane?
[ ] Are input and target still aligned?
[ ] Is class balance what I think it is?
FORWARD
[ ] Are output shapes correct?
[ ] Are outputs finite?
[ ] Do activations have reasonable scale?
[ ] Are activations collapsing or exploding?
LOSS
[ ] Is the loss appropriate for the task?
[ ] Am I passing logits/probabilities in the expected form?
[ ] Can I reproduce the loss manually for one example?
AUTOGRAD
[ ] Does loss.requires_grad == True?
[ ] Does loss have a grad_fn?
[ ] Do expected parameters receive gradients?
[ ] Are gradients finite?
[ ] Are gradient norms plausible?
OPTIMIZER
[ ] Are all trainable parameters in the optimizer?
[ ] Does optimizer.step() change them?
[ ] What is the actual current learning rate?
[ ] Did I replace parameters after constructing the optimizer?
TRAINING
[ ] Can the model overfit 32 examples?
[ ] Does the loss decrease without augmentation/regularization?
[ ] Does full precision work before AMP?
[ ] Can I reproduce the failure with fixed seeds?
We can turn this checklist into a decision flow:
flowchart TD
A[Model not learning] --> DATA
subgraph DATA
D1[Input values sane?] --> D2[Labels sane?] --> D3[Input/target aligned?]
end
DATA --> FORWARD
subgraph FORWARD
F1[Output shapes correct?] --> F2[Outputs finite?] --> F3[Activation scale reasonable?]
end
FORWARD --> LOSS
subgraph LOSS
L1[Loss appropriate?] --> L2[Logits/probs correct?] --> L3[Loss correct per example?]
end
LOSS --> AUTOGRAD
subgraph AUTOGRAD
A1[Loss requires_grad?] --> A2[Grad_fn exists?] --> A3[Grads present/finite/plausible?]
end
AUTOGRAD --> OPTIMIZER
subgraph OPTIMIZER
O1[Params in optimizer?] --> O2[step changes them?] --> O3[Actual LR as expected?]
end
OPTIMIZER --> TRAINING
subgraph TRAINING
T1[Overfit tiny batch?] --> T2[Learn without regularization?] --> T3[Full precision works?]
end
TRAINING --> FIX[Diagnose found issue]
Do not skip directly to architecture changes until you can answer those questions.
30. Deliberately break the model
The best way to learn this is to create failures on purpose.
Break 1: remove parameters from the optimizer
optimizer = torch.optim.AdamW(
model.net[-1].parameters(),
lr=1e-3,
)
What happens to the earlier layers?
Can your optimizer-membership report detect it?
Break 2: detach the hidden representation
class BrokenClassifier(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(2, 32)
self.fc2 = nn.Linear(32, 2)
def forward(self, x):
h = F.relu(self.fc1(x))
h = h.detach()
return self.fc2(h)
Which parameters stop receiving gradients?
Break 3: destroy target alignment
perm = torch.randperm(len(x))
x_bad = x[perm]
y_bad = y
The code runs.
The shapes match.
The learning problem has been destroyed.
Break 4: use a ridiculous learning rate
optimizer = torch.optim.AdamW(model.parameters(), lr=10.0)
Track:
loss
gradient norm
parameter norm
first non-finite tensor
Break 5: apply softmax before cross entropy
probs = model(x).softmax(dim=-1)
loss = F.cross_entropy(probs, y)
Compare learning speed and gradients with the correct logit-based version.
31. What an LLM should do for you here
Use the LLM to accelerate the investigation.
Give it evidence.
Bad debugging prompt:
My PyTorch model does not learn. Fix it.
Better:
Training loss stays at 0.693.
The dataset is balanced.
The model can not overfit 32 examples.
The final layer has grad norm 2.1e-2.
The first layer has grad=None.
All parameters appear in model.named_parameters().
Here is the forward method and optimizer construction.
Explain the most likely graph break.
The difference is enormous.
The LLM becomes far more useful once you have turned the failure into observations.
The machine can propose hypotheses.
Your runtime decides which one is true.
32. What we are really learning
At this point the course is no longer primarily about remembering PyTorch syntax.
You already know how to write:
loss.backward()
optimizer.step()
The important skill is knowing what those lines imply about the system around them.
You should be able to ask:
What graph did backward traverse?
Which leaves received gradients?
Which parameters does the optimizer own?
Which values changed?
Where did non-finite values first appear?
Can this network memorize one batch?
That is the level where generated code becomes inspectable rather than mysterious.
A compact debugging harness
Here is a final version worth keeping around.
import torch
def debug_training_step(model, optimizer, loss_fn, x, y):
model.train()
print("INPUT")
print("x", x.shape, x.dtype, x.device)
print("y", y.shape, y.dtype, y.device)
print("x finite", torch.isfinite(x).all().item())
optimizer.zero_grad(set_to_none=True)
logits = model(x)
loss = loss_fn(logits, y)
print("\nFORWARD")
print("logits", logits.shape)
print("logits finite", torch.isfinite(logits).all().item())
print("loss", loss.item())
print("loss requires_grad", loss.requires_grad)
print("loss grad_fn", loss.grad_fn)
if not torch.isfinite(loss):
raise RuntimeError("loss is non-finite")
loss.backward()
optimizer_ids = {
id(p)
for group in optimizer.param_groups
for p in group["params"]
}
print("\nGRADIENTS")
for name, p in model.named_parameters():
if p.grad is None:
grad_info = "None"
else:
grad_info = (
f"norm={p.grad.norm().item():.3e} "
f"finite={torch.isfinite(p.grad).all().item()}"
)
print(
name,
"requires_grad=", p.requires_grad,
"optimizer=", id(p) in optimizer_ids,
"grad=", grad_info,
)
before = {
name: p.detach().clone()
for name, p in model.named_parameters()
}
optimizer.step()
print("\nUPDATES")
for name, p in model.named_parameters():
delta = (p.detach() - before[name]).norm().item()
print(name, f"delta={delta:.3e}")
return loss.detach()
If you can run this against one batch, you can answer a remarkable number of “model not learning” questions in a few minutes.
Where we go next
We have now built enough understanding to ask a different question.
Suppose the model does learn.
But it is slow.
Or it runs out of GPU memory.
Or torch.compile makes it slower rather than faster.
Or mixed precision changes the result.
Or the GPU spends half its time idle.
That is Step 09:
PyTorch Performance: CUDA Memory, Profiling, Mixed Precision and
torch.compile
After that, Step 10 will combine the entire series into the capstone:
Build a Small Language Model From Scratch in PyTorch.
