This article is published in English.
Debugging a Small GPT in PyTorch: Tests That Isolate Each Failure
A stage-by-stage workflow for debugging a character-level GPT in PyTorch, from token IDs and target shifting to gradients, NaN losses and checkpoints.
Evaluation metrics such as loss, perplexity and repetition rate tell you that a small GPT model is misbehaving, but rarely why. The tempting response is to tweak the learning rate, add a layer and hope. This guide replaces that guessing with a repeatable workflow for a compact, character-level GPT (Mini-GPT) trained on WikiText-2: map each symptom to probable causes, run a targeted check for every pipeline stage, and fix problems in the order they occur. You end up with a set of assertions and one diagnostic script to run before every expensive training job.
Why a GPT pipeline can be wrong without crashing
Training a language model chains many transformations, each consuming the output of the previous one:
Raw dataset
→ cleaned text
→ tokenizer
→ token IDs
→ training batches
→ embeddings
→ Transformer blocks
→ vocabulary logits
→ cross-entropy loss
→ gradients
→ optimizer
→ checkpoints
→ generation
A defect anywhere contaminates everything downstream. Suppose the targets are not offset from the inputs by one position:
Input: The cat
Target: The cat
The network is now rewarded for reproducing the token it already sees instead of predicting the next one. Nothing throws, and the loss may still fall, because copying is easy. The model is simply optimizing the wrong objective. That is the key difference from debugging a web service, where a wrong return value usually breaks a test or a page:
Code that runs to completion can still train a broken model.
Work from the start of the pipeline toward the end
Verify stages in the order data flows through them:
1. Environment
2. Files
3. Tokenizer
4. Token IDs
5. Training batches
6. Model shapes
7. Initial loss
8. Gradients
9. Optimizer
10. Validation behavior
11. Checkpoints
12. Generation
Judging generation quality before the data pipeline is confirmed wastes time, because a bad sample could originate in any of the eleven earlier stages. At each stage, ask one narrow question and answer it with a check that passes or fails.
Four families of failure
Nearly every problem with a model like this falls into one of four groups, and knowing the group narrows the search.
- Correctness failures: the code is logically wrong. Targets are not shifted, the causal mask lets positions see later tokens, the loss uses tensors with the wrong layout, or the tokenizer's IDs disagree with the vocabulary the model was built for.
- Numerical failures: the math turns unstable. The loss becomes
NaN, gradients explode, logits overflow to infinity, or softmax receives a row with no valid entries. - Optimization failures: the implementation is correct but learning is ineffective, because the learning rate is too high or too low, the model is too small, or the run is too short.
- Generalization and generation failures: training works but the model does not. Validation loss rises, samples loop, output ignores the prompt, or the model reproduces training passages.
Start with a tiny debug configuration
Debugging on a full run turns every hypothesis into a long wait. Define a small model that can memorize a few examples within seconds:
debug_config = MiniGPTConfig(
vocab_size=tokenizer.vocab_size,
block_size=32,
embedding_dim=64,
num_heads=4,
num_layers=2,
expansion_factor=4,
dropout=0.0,
)
Pair it with a small batch:
debug_batch_size = 8
and a short run:
debug_steps = 200
Dropout is deliberately off:
dropout = 0.0
Dropout zeroes random activations, so identical runs diverge. Removing it (together with a fixed seed) makes each test reproducible. Restore the production settings once the pipeline passes.
Confirm the runtime environment
Before touching the model, print the Python and PyTorch versions and whether CUDA or Apple's Metal backend (MPS) is usable:
import platform
import torch
print("Python:", platform.python_version())
print("PyTorch:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())
if hasattr(torch.backends, "mps"):
print(
"MPS available:",
torch.backends.mps.is_available(),
)
A helper picks the best device, preferring CUDA, then MPS, then the CPU. The hasattr guard keeps it working on older builds without an MPS backend:
def get_device():
if torch.cuda.is_available():
return torch.device("cuda")
if (
hasattr(torch.backends, "mps")
and torch.backends.mps.is_available()
):
return torch.device("mps")
return torch.device("cpu")
Call it once and log the result:
device = get_device()
print("Selected device:", device)
When training is mysteriously slow, this line is often the explanation: the job expected a GPU but fell back to the CPU because of a driver or installation issue.
Make sure every input file is present
Confirm that the tokenizer definition and the encoded train, validation and test splits exist, failing early with a clear FileNotFoundError instead of a confusing error deep in the training loop:
from pathlib import Path
required_paths = [
Path("tokenizer/char_tokenizer.json"),
Path("data/encoded/train_ids.pt"),
Path("data/encoded/val_ids.pt"),
Path("data/encoded/test_ids.pt"),
]
for path in required_paths:
if not path.exists():
raise FileNotFoundError(
f"Required file not found: {path}"
)
print("Found:", path)
Print the sizes too:
for path in required_paths:
print(
path,
path.stat().st_size,
"bytes",
)
An empty or unusually small file usually means a preprocessing job was interrupted and left a truncated artifact.
Test the tokenizer in isolation
Load the character tokenizer:
tokenizer = CharTokenizer.from_file(
"tokenizer/char_tokenizer.json"
)
Check the vocabulary size and both ends of the character list for anything missing or garbled:
print("Vocabulary size:", tokenizer.vocab_size)
print("First tokens:", tokenizer.chars[:20])
print("Last tokens:", tokenizer.chars[-20:])
The round-trip property
A lossless tokenizer returns the exact input after encoding and decoding. Printing with repr exposes invisible characters such as trailing spaces:
sample = "The history of"
encoded = tokenizer.encode(sample)
decoded = tokenizer.decode(encoded)
print("Encoded:", encoded)
print("Decoded:", repr(decoded))
assert decoded == sample
The invariant being checked:
decode(encode(text)) = text
A failure means at least one character cannot be represented faithfully, typically a symbol missing from the vocabulary. The encode method in the full script raises a KeyError in that case rather than silently dropping it, which is what you want.
Catch tokenizer and checkpoint mismatches
Equal vocabulary sizes are necessary but not sufficient. Two vocabularies can both hold 100 characters and still assign different IDs:
Tokenizer A: "a" → 10
Tokenizer B: "a" → 24
A model trained with one mapping and served with the other emits nonsense while every tensor shape looks right. Persist the vocabulary with the checkpoint, or at least a fingerprint of it. A SHA-256 hash of the serialized character list works; fixing separators and ensure_ascii guarantees the same list always serializes to the same bytes:
import hashlib
import json
def tokenizer_fingerprint(chars):
payload = json.dumps(
chars,
ensure_ascii=False,
separators=(",", ":"),
).encode("utf-8")
return hashlib.sha256(
payload
).hexdigest()
Compute it for the loaded tokenizer:
fingerprint = tokenizer_fingerprint(
tokenizer.chars
)
print("Tokenizer fingerprint:", fingerprint)
Store it in the checkpoint dictionary when saving:
checkpoint[
"tokenizer_fingerprint"
] = fingerprint
On load, compare and refuse to continue on a mismatch. The is not None condition still lets older checkpoints without a fingerprint load:
saved_fingerprint = checkpoint.get(
"tokenizer_fingerprint"
)
if (
saved_fingerprint is not None
and saved_fingerprint != fingerprint
):
raise ValueError(
"Checkpoint and tokenizer do not match"
)
Inspect the encoded token IDs
Load the training split onto the CPU as 64-bit integers, the type embeddings and cross-entropy expect:
train_ids = torch.load(
"data/encoded/train_ids.pt",
map_location="cpu",
).long()
Print shape, dtype and range:
print("Shape:", train_ids.shape)
print("Dtype:", train_ids.dtype)
print("Minimum ID:", train_ids.min().item())
print("Maximum ID:", train_ids.max().item())
Every ID must lie inside the vocabulary:
0 ≤ token ID < vocabulary size
As assertions:
assert train_ids.min().item() >= 0
assert (
train_ids.max().item()
< tokenizer.vocab_size
)
An out-of-range ID breaks the embedding lookup, and on a GPU the error can surface as an opaque device-side assertion far from its cause. Typical reasons are a wrong tokenizer file, corrupt encoded files, a vocabulary rebuilt after encoding, or inconsistent special-token handling.
Read the stored data back as text
In-range numbers can still encode the wrong text, so decode a few hundred IDs and read them:
sample_ids = train_ids[:500]
sample_text = tokenizer.decode(
sample_ids.tolist()
)
print(sample_text)
You should see readable WikiText with its usual formatting, line breaks, headings and punctuation, and no runs of repeated or corrupted characters. If the sample looks wrong, stop: no model change compensates for a broken tokenizer or dataset.
Verify the one-token shift between inputs and targets
Build one example by hand, with the target window starting one position later:
block_size = 32
start = 100
inputs = train_ids[
start:
start + block_size
]
targets = train_ids[
start + 1:
start + block_size + 1
]
Decode both to compare them:
input_text = tokenizer.decode(
inputs.tolist()
)
target_text = tokenizer.decode(
targets.tolist()
)
print("Input: ", repr(input_text))
print("Target:", repr(target_text))
The target should read like the input with its first character dropped and one new character appended. Then assert the relationship on the tensors:
assert torch.equal(
inputs[1:],
targets[:-1],
)
Few checks in the project catch more serious bugs. The invariant:
inputs[1:] == targets[:-1]
Each target position holds the token that follows the matching input position, which is exactly what next-token prediction needs.
The identical-slice bug
The classic mistake uses the same bounds for both slices:
inputs = data[
start:
start + block_size
]
targets = data[
start:
start + block_size
]
The model then learns an identity mapping:
Current token → current token
Starting the target slice one token later fixes it:
targets = data[
start + 1:
start + block_size + 1
]
and restores the intended task:
Current context → next token
A telltale sign of this bug is a loss that drops suspiciously fast early on.
Check batch shapes, dtypes and devices
Sample a real batch:
inputs, targets = get_batch(
data=train_ids,
batch_size=8,
block_size=32,
device=device,
)
Print everything that could be wrong:
print("Input shape:", inputs.shape)
print("Target shape:", targets.shape)
print("Input dtype:", inputs.dtype)
print("Target dtype:", targets.dtype)
print("Input device:", inputs.device)
print("Target device:", targets.device)
For batch size 8 and block size 32, expect:
Input shape: [8, 32]
Target shape: [8, 32]
Dtype: torch.int64
Device: same as model
Make the expectations permanent:
assert inputs.shape == targets.shape
assert inputs.dtype == torch.long
assert targets.dtype == torch.long
assert inputs.device == device
assert targets.device == device
Fix device mismatches
This error appears constantly in PyTorch work:
Expected all tensors to be on the same device
One operation received tensors on different devices, such as CPU and GPU. Print where the parameters and the batch live:
model_device = next(
model.parameters()
).device
print("Model device:", model_device)
print("Input device:", inputs.device)
Move the model and every batch explicitly:
model = model.to(device)
inputs = inputs.to(device)
targets = targets.to(device)
A subtler source hides inside the model: torch.arange defaults to the CPU, so take the device from the incoming tokens:
positions = torch.arange(
sequence_length,
device=token_ids.device,
)
Otherwise, adding position embeddings to token embeddings on CUDA or MPS fails. Tying the device to the input also keeps the model portable.
Validate the forward pass
Run one batch with targets so the model returns logits and loss:
logits, loss = model(
inputs,
targets,
)
Inspect them:
print("Logits shape:", logits.shape)
print("Loss shape:", loss.shape)
print("Loss value:", loss.item())
Logits need one score per vocabulary entry for every position, and the loss must be a scalar:
assert logits.shape == (
inputs.size(0),
inputs.size(1),
tokenizer.vocab_size,
)
assert loss.ndim == 0
Logits shaped [B, V, T] mean a transpose or reshape put dimensions in the wrong order. Because F.cross_entropy accepts class scores in the second dimension, a misordered tensor can sometimes reach the loss without an error and compute something meaningless.
Compare the initial loss with the random baseline
A freshly initialized model with small weights predicts an almost uniform distribution, and cross-entropy against a uniform distribution over V classes is log(V):
import math
expected_loss = math.log(
tokenizer.vocab_size
)
print("Expected loss:", expected_loss)
print("Actual loss:", loss.item())
A small deviation is normal; a large one is a clue. A starting loss far above the baseline suggests extreme logits, unstable initialization, invalid token IDs, targets that do not match the vocabulary, or a wrong output shape. A starting loss far below it, which a model that knows nothing cannot achieve honestly, suggests data leakage, trained weights loaded by accident, targets equal to inputs, visible future tokens, or an unintended checkpoint resume.
Prove that the causal mask works
A GPT must predict each position from earlier tokens only. Build two sequences with a shared prefix and different suffixes; if the model is causal, prefix logits must match. Evaluation mode disables dropout so randomness does not interfere:
model.eval()
prefix_length = 8
sequence_length = 16
sequence_a = torch.randint(
0,
tokenizer.vocab_size,
(1, sequence_length),
device=device,
)
sequence_b = sequence_a.clone()
sequence_b[
:,
prefix_length:
] = torch.randint(
0,
tokenizer.vocab_size,
(
1,
sequence_length - prefix_length,
),
device=device,
)
Run both without gradients:
with torch.no_grad():
logits_a, _ = model(sequence_a)
logits_b, _ = model(sequence_b)
Measure the largest prefix difference:
prefix_difference = (
logits_a[:, :prefix_length, :]
- logits_b[:, :prefix_length, :]
).abs().max().item()
print(
"Maximum prefix difference:",
prefix_difference,
)
Assert equality within a small tolerance that absorbs floating-point noise:
assert torch.allclose(
logits_a[:, :prefix_length, :],
logits_b[:, :prefix_length, :],
atol=1e-5,
)
A failure means later positions leak into earlier ones, usually because the mask is missing, applied to the wrong dimension, or built from the wrong triangle. Testing behavior end to end is stronger than inspecting the mask tensor.
Overfit a single batch
If you adopt one technique from this guide, choose this one:
A model with enough capacity should be able to memorize one small batch.
It exercises data, model, loss, backpropagation and optimizer together. Fix one batch that every step reuses:
fixed_inputs, fixed_targets = get_batch(
data=train_ids,
batch_size=8,
block_size=32,
device=device,
)
Build the small, dropout-free model:
debug_config = MiniGPTConfig(
vocab_size=tokenizer.vocab_size,
block_size=32,
embedding_dim=64,
num_heads=4,
num_layers=2,
expansion_factor=4,
dropout=0.0,
)
debug_model = MiniGPT(
debug_config
).to(device)
Train on that batch repeatedly, logging every 50 steps:
optimizer = torch.optim.AdamW(
debug_model.parameters(),
lr=1e-3,
)
for step in range(500):
optimizer.zero_grad(
set_to_none=True
)
_, debug_loss = debug_model(
fixed_inputs,
fixed_targets,
)
debug_loss.backward()
optimizer.step()
if step % 50 == 0:
print(
step,
debug_loss.item(),
)
The loss should fall far below the baseline. If it does not, suspect a broken loss, gradients that never reach some parameters, unshifted targets, a model too small even for this, a badly chosen learning rate, a faulty causal mask, or an optimizer that is not updating anything. Treat the test as a gate for any full run.
Confirm that gradients reach every parameter
Run one forward and backward pass:
optimizer.zero_grad(
set_to_none=True
)
_, loss = model(
inputs,
targets,
)
loss.backward()
Report every trainable parameter whose .grad is still None, and the norm for the rest:
for name, parameter in (
model.named_parameters()
):
if not parameter.requires_grad:
continue
if parameter.grad is None:
print(
"NO GRADIENT:",
name,
)
else:
print(
name,
parameter.grad.norm().item(),
)
A missing gradient usually means a layer defined in __init__ but unused in forward, an accidental .detach(), a forward branch that skips a component, a loss computed from a disconnected tensor, or requires_grad=False.
Track the global gradient norm
Per-parameter norms find dead layers; one aggregate number tracks stability over time. This function combines all gradient L2 norms, the same quantity clipping uses:
def calculate_gradient_norm(model):
squared_norm = 0.0
for parameter in model.parameters():
if parameter.grad is None:
continue
parameter_norm = (
parameter.grad
.detach()
.norm(2)
.item()
)
squared_norm += (
parameter_norm ** 2
)
return squared_norm ** 0.5
Log it after each backward pass:
gradient_norm = (
calculate_gradient_norm(model)
)
print("Gradient norm:", gradient_norm)
Watch for norms that are exactly zero, huge or spiking, NaN, or inf.
Detect NaN and infinity early
A helper raises as soon as a tensor contains a nonfinite value, naming the tensor:
def assert_finite_tensor(
tensor,
name,
):
if not torch.isfinite(
tensor
).all():
raise FloatingPointError(
f"{name} contains NaN or infinity"
)
Apply it to logits and loss:
assert_finite_tensor(
logits,
"logits",
)
assert_finite_tensor(
loss,
"loss",
)
and to every gradient after backward():
for name, parameter in (
model.named_parameters()
):
if parameter.grad is not None:
assert_finite_tensor(
parameter.grad,
f"gradient for {name}",
)
Checking several points reveals the first place invalid numbers appear, which is far more useful than spotting a NaN loss hundreds of steps later.
Why the loss turns into NaN
Frequent causes: a learning rate that is too high, exploding gradients, an attention row with every position masked, invalid softmax input, mixed-precision overflow, already corrupted parameters, division by zero, a logarithm of zero or a negative value, and infinite logits. When it happens:
- Stop the run.
- Find the last step with a finite loss.
- Lower the learning rate.
- Enable gradient clipping.
- Re-check the causal mask.
- Turn off mixed precision.
- Scan parameters and gradients for nonfinite values.
Never keep stepping once parameters contain NaN; each update spreads the corruption, so resume from the last good checkpoint instead.
Use gradient clipping as a guardrail
Clipping rescales gradients whose combined norm exceeds a threshold, so it belongs between backward() and optimizer.step():
loss.backward()
gradient_norm = (
torch.nn.utils.clip_grad_norm_(
model.parameters(),
max_norm=1.0,
)
)
optimizer.step()
clip_grad_norm_ returns the norm measured before clipping, which doubles as monitoring:
print(
"Gradient norm before clipping:",
float(gradient_norm),
)
If the threshold is exceeded almost every step, clipping is masking a problem such as an excessive learning rate or numerical instability. It protects against an occasional bad batch; it does not replace a sensible learning rate.
Verify that the optimizer changes the weights
Copy one parameter before an update. The .clone() matters: without it, before shares storage with the parameter and changes too:
parameter_name, parameter = next(
model.named_parameters()
)
before = parameter.detach().clone()
Run one training step:
optimizer.zero_grad(
set_to_none=True
)
_, loss = model(
inputs,
targets,
)
loss.backward()
optimizer.step()
Measure the change:
after = parameter.detach()
maximum_change = (
after - before
).abs().max().item()
print(
"Maximum parameter change:",
maximum_change,
)
and require one:
assert maximum_change > 0
Unchanged weights point to a zero learning rate, an optimizer built without the model's parameters (for instance before the model was replaced), missing gradients, a missing optimizer.step(), or frozen parameters.
Tune the learning rate in both directions
A rate that is too high shows up as loss climbing quickly or swinging wildly, very large gradient norms, a NaN loss, and samples that never improve. A first fix is lowering the peak, for example to:
max_learning_rate = 1e-4
rather than:
max_learning_rate = 1e-3
Log the learning rate beside the loss. With warmup, instability often begins exactly at the peak, which a loss plot alone hides.
A rate that is too low looks different: loss falls very slowly although gradients exist, parameters barely move, and even the single-batch test takes many steps. Then raise it, for example to:
max_learning_rate = 3e-4
rather than:
max_learning_rate = 1e-5
No value is universally right; it varies with the size of the model and batch, the optimizer and the dataset. Run short experiments where only the learning rate changes.
A checklist for a loss that will not go down
Work through these questions in order. Data:
Are targets shifted by one token?
Are token IDs within range?
Does decoded input look correct?
Model:
Are logits shaped [B, T, V]?
Is the causal mask valid?
Are positions on the correct device?
Loss:
Does cross-entropy receive raw logits?
Are logits and targets flattened correctly?
Feeding softmax probabilities into cross_entropy is a classic mistake, since the function applies log-softmax itself. Gradients:
Do all important parameters receive gradients?
Are gradient norms finite and nonzero?
Optimizer:
Is the learning rate positive?
Does optimizer.step() run?
Do parameters change?
Capacity:
Can the model overfit one batch?
This order avoids random trial and error.
This order eliminates one class of cause at a time instead of relying on trial and error.
Recognize overfitting and underfitting
Overfitting shows up as diverging curves:
Training loss: continues decreasing
Validation loss: stops decreasing or increases
Quantify the gap:
generalization_gap = (
validation_loss
- training_loss
)
Remedies include keeping the best-validation checkpoint, raising dropout or weight decay, shrinking the model, adding more or more varied data, and stopping earlier. Pick the stopping point with the validation split; using the test split leaks information and inflates the final score.
Underfitting shows up as two curves that stay high together:
Training loss: remains high
Validation loss: remains similarly high
Likely reasons are too little capacity, a run that is too short, a low learning rate, a short context window, data too hard for the architecture, or tokenization that wastes context. Options include more steps, a larger embedding dimension, more Transformer layers, a longer context, byte pair encoding (BPE) instead of characters, and re-tuning the learning rate. First rerun the single-batch test: if the model cannot memorize one batch, the problem is correctness or optimization, not capacity.
Diagnose repetitive generation
Repetition looks like this:
the the the the
or, with WikiText heading markers:
= = = = = = =
Causes include greedy decoding, very low temperature, a tiny top-k, an undertrained or overfit model, repeated structures in the data, and a short context window. Try more balanced sampling:
temperature = 0.8
top_k = 20
top_p = 0.9
A repetition penalty can help if it stays gentle:
repetition_penalty = 1.05
In character models a strong penalty discourages reusing letters and quickly ruins spelling. If every decoding setup still loops, the model is the weak point, not the sampler. For how these settings interact, see our guide to temperature, top-k and top-p.
Diagnose chaotic generation
The opposite failure produces stray symbols, broken words, excess punctuation, abrupt topic jumps and unreadable strings. Probable causes are a high temperature, no top-k or top-p filtering, a mismatched tokenizer, the wrong checkpoint, an undertrained model with high validation loss, or weights that never loaded. Try tighter sampling:
temperature = 0.6
top_k = 10
top_p = 0.9
Confirm the weights really came from the checkpoint:
model.load_state_dict(
checkpoint["model_state_dict"]
)
and that dropout is off during sampling:
model.eval()
When output ignores the prompt
The prompt may be very short or unlike the training data; the model may be small, undertrained, weak at long-range dependencies, or limited by a short context; and character tokens make semantic patterns harder to learn. Test longer, WikiText-style prompts. Compare a minimal one:
"The "
with a richer one:
"The history of the city began"
The second gives the model much more to condition on. If continuations still drift, examine validation loss and the attention implementation.
Checkpoints that refuse to load
The errors are familiar:
Missing key(s) in state_dict
Unexpected key(s) in state_dict
Size mismatch
They mean the model you built is not the one you saved: its configuration, layer count, embedding dimension, vocabulary size or weight tying changed, classes or attributes were renamed, or an optimizer state from another architecture is being loaded. Inspect the stored configuration:
print(
checkpoint["config"]
)
Build the model from it instead of from current defaults:
config = MiniGPTConfig(
**checkpoint["config"]
)
model = MiniGPT(config)
Then load the state dictionary. Constructing a model from a new configuration and expecting old weights to fit causes most of these errors.
Listing missing and unexpected keys
For diagnosis only, load non-strictly and print the mismatches:
load_result = model.load_state_dict(
checkpoint["model_state_dict"],
strict=False,
)
print(
"Missing keys:",
load_result.missing_keys,
)
print(
"Unexpected keys:",
load_result.unexpected_keys,
)
The lists usually reveal the cause, such as a renamed submodule. For inference or resumed training keep strict loading, so an incompatible checkpoint fails loudly instead of leaving layers at random initial values.
Move optimizer state to the right device
After restoring an optimizer, its internal tensors (such as AdamW's moment estimates) may sit on a different device than the model. This helper moves every tensor in the state:
def move_optimizer_to_device(
optimizer,
device,
):
for state in optimizer.state.values():
for key, value in state.items():
if torch.is_tensor(value):
state[key] = value.to(
device
)
Call it right after loading:
optimizer.load_state_dict(
checkpoint[
"optimizer_state_dict"
]
)
move_optimizer_to_device(
optimizer,
device,
)
It matters most when saving on one machine and resuming on another, for example CUDA to MPS or CPU.
Bundle the key checks into one health function
Collect the most important assertions in one function: token dtype and range, target shifting, logits shape, a finite loss, and a comparison with the random baseline:
def run_model_health_checks(
model,
tokenizer,
train_ids,
device,
):
model.eval()
assert train_ids.dtype == torch.long
assert train_ids.min().item() >= 0
assert (
train_ids.max().item()
< tokenizer.vocab_size
)
batch_size = 4
block_size = min(
32,
model.config.block_size,
)
inputs, targets = get_batch(
data=train_ids,
batch_size=batch_size,
block_size=block_size,
device=device,
)
assert inputs.shape == targets.shape
assert torch.equal(
inputs[:, 1:],
targets[:, :-1],
)
with torch.no_grad():
logits, loss = model(
inputs,
targets,
)
assert logits.shape == (
batch_size,
block_size,
tokenizer.vocab_size,
)
assert torch.isfinite(loss)
expected_loss = math.log(
tokenizer.vocab_size
)
print("Current loss:", loss.item())
print(
"Random baseline:",
expected_loss,
)
print("Model health checks passed.")
For a trained checkpoint the loss should be clearly below the baseline; otherwise the weights did not load or the tokenizer does not match.
Locate numerical problems with forward hooks
When a NaN appears somewhere unknown, forward hooks inspect each module's output during the pass. This hook handles single tensors and tuples and raises with the module's class name on the first nonfinite value:
def finite_output_hook(
module,
inputs,
output,
):
tensors = []
if torch.is_tensor(output):
tensors = [output]
elif isinstance(output, tuple):
tensors = [
item
for item in output
if torch.is_tensor(item)
]
for tensor in tensors:
if not torch.isfinite(
tensor
).all():
raise FloatingPointError(
"Nonfinite output detected in "
f"{module.__class__.__name__}"
)
Attach it to every linear, layer-norm and embedding module, keeping the handles:
hooks = []
for module in model.modules():
if isinstance(
module,
(
torch.nn.Linear,
torch.nn.LayerNorm,
torch.nn.Embedding,
),
):
hooks.append(
module.register_forward_hook(
finite_output_hook
)
)
Run one forward pass; since layers execute in order, the first exception names the first failing layer type. Then remove the hooks:
for hook in hooks:
hook.remove()
Hooks run on every forward call and slow the model, so use them only while hunting a bug. For the exact module path, record names from named_modules() when registering.
A complete diagnostic script
All checks combine into one command-line tool. Save it as:
debug_mini_gpt.py
The script defines a minimal CharTokenizer, device selection, the fingerprint, a batching helper and the numerical utilities. It builds the model from the checkpoint's own configuration, rejects a tokenizer whose vocabulary size differs, then runs eight numbered tests: tokenizer round trip, token range, batch shifting, forward pass, causal independence, gradients, optimizer update and, optionally, single-batch overfitting on a fresh debug model, all under a fixed seed. Two details are worth noticing: the tokenizer test decodes stored IDs and re-encodes them, validating the real dataset, and the optimizer test uses a fresh AdamW instance so stale state cannot interfere.
import argparse
import hashlib
import json
import math
from pathlib import Path
import torch
from mini_gpt import MiniGPT
from mini_gpt import MiniGPTConfig
class CharTokenizer:
def __init__(self, chars):
self.chars = chars
self.vocab_size = len(chars)
self.stoi = {
char: index
for index, char in enumerate(chars)
}
self.itos = {
index: char
for index, char in enumerate(chars)
}
@classmethod
def from_file(cls, path):
with open(
path,
"r",
encoding="utf-8",
) as file:
data = json.load(file)
return cls(data["chars"])
def encode(self, text):
return [
self.stoi[char]
for char in text
]
def decode(self, token_ids):
return "".join(
self.itos[int(token_id)]
for token_id in token_ids
)
def get_device():
if torch.cuda.is_available():
return torch.device("cuda")
if (
hasattr(torch.backends, "mps")
and torch.backends.mps.is_available()
):
return torch.device("mps")
return torch.device("cpu")
def tokenizer_fingerprint(chars):
payload = json.dumps(
chars,
ensure_ascii=False,
separators=(",", ":"),
).encode("utf-8")
return hashlib.sha256(
payload
).hexdigest()
def get_batch(
data,
batch_size,
block_size,
device,
):
start_positions = torch.randint(
low=0,
high=len(data) - block_size,
size=(batch_size,),
)
inputs = torch.stack([
data[
position:
position + block_size
]
for position in start_positions
])
targets = torch.stack([
data[
position + 1:
position + block_size + 1
]
for position in start_positions
])
return (
inputs.to(device),
targets.to(device),
)
def assert_finite_tensor(
tensor,
name,
):
if not torch.isfinite(
tensor
).all():
raise FloatingPointError(
f"{name} contains NaN or infinity"
)
def calculate_gradient_norm(model):
squared_norm = 0.0
for parameter in model.parameters():
if parameter.grad is None:
continue
norm = (
parameter.grad
.detach()
.norm(2)
.item()
)
squared_norm += norm ** 2
return squared_norm ** 0.5
def load_model(
checkpoint_path,
device,
):
checkpoint = torch.load(
checkpoint_path,
map_location=device,
)
config = MiniGPTConfig(
**checkpoint["config"]
)
model = MiniGPT(config)
model.load_state_dict(
checkpoint["model_state_dict"]
)
model = model.to(device)
return model, checkpoint
def test_tokenizer(
tokenizer,
train_ids,
):
print("\n1. Testing tokenizer")
print(
"Vocabulary size:",
tokenizer.vocab_size,
)
print(
"Tokenizer fingerprint:",
tokenizer_fingerprint(
tokenizer.chars
),
)
sample_ids = train_ids[:300]
sample_text = tokenizer.decode(
sample_ids.tolist()
)
round_trip_ids = tokenizer.encode(
sample_text
)
assert round_trip_ids == (
sample_ids.tolist()
)
print("Decoded sample:")
print(repr(sample_text))
print(
"Tokenizer round-trip test passed."
)
def test_token_ids(
tokenizer,
train_ids,
):
print("\n2. Testing token IDs")
print("Shape:", train_ids.shape)
print("Dtype:", train_ids.dtype)
minimum_id = train_ids.min().item()
maximum_id = train_ids.max().item()
print("Minimum ID:", minimum_id)
print("Maximum ID:", maximum_id)
assert minimum_id >= 0
assert maximum_id < (
tokenizer.vocab_size
)
print("Token ID test passed.")
def test_batch(
tokenizer,
train_ids,
block_size,
device,
):
print("\n3. Testing batches")
inputs, targets = get_batch(
data=train_ids,
batch_size=4,
block_size=block_size,
device=device,
)
print("Input shape:", inputs.shape)
print("Target shape:", targets.shape)
assert inputs.shape == targets.shape
assert inputs.dtype == torch.long
assert targets.dtype == torch.long
assert torch.equal(
inputs[:, 1:],
targets[:, :-1],
)
input_text = tokenizer.decode(
inputs[0].cpu().tolist()
)
target_text = tokenizer.decode(
targets[0].cpu().tolist()
)
print("Input sample:")
print(repr(input_text))
print("Target sample:")
print(repr(target_text))
print("Batch shift test passed.")
return inputs, targets
def test_forward_pass(
model,
tokenizer,
inputs,
targets,
):
print("\n4. Testing forward pass")
model.eval()
with torch.no_grad():
logits, loss = model(
inputs,
targets,
)
print("Logits shape:", logits.shape)
print("Loss:", loss.item())
assert logits.shape == (
inputs.size(0),
inputs.size(1),
tokenizer.vocab_size,
)
assert_finite_tensor(
logits,
"logits",
)
assert_finite_tensor(
loss,
"loss",
)
print(
"Random baseline:",
math.log(tokenizer.vocab_size),
)
print("Forward-pass test passed.")
def test_future_independence(
model,
tokenizer,
device,
):
print(
"\n5. Testing causal independence"
)
model.eval()
sequence_length = min(
16,
model.config.block_size,
)
prefix_length = (
sequence_length // 2
)
sequence_a = torch.randint(
0,
tokenizer.vocab_size,
(1, sequence_length),
device=device,
)
sequence_b = sequence_a.clone()
sequence_b[
:,
prefix_length:
] = torch.randint(
0,
tokenizer.vocab_size,
(
1,
sequence_length
- prefix_length,
),
device=device,
)
with torch.no_grad():
logits_a, _ = model(sequence_a)
logits_b, _ = model(sequence_b)
difference = (
logits_a[
:,
:prefix_length,
:,
]
- logits_b[
:,
:prefix_length,
:,
]
).abs().max().item()
print(
"Maximum shared-prefix difference:",
difference,
)
assert torch.allclose(
logits_a[
:,
:prefix_length,
:,
],
logits_b[
:,
:prefix_length,
:,
],
atol=1e-5,
)
print(
"Causal independence test passed."
)
def test_gradients(
model,
inputs,
targets,
):
print("\n6. Testing gradients")
model.train()
model.zero_grad(
set_to_none=True
)
_, loss = model(
inputs,
targets,
)
loss.backward()
missing_gradients = []
nonfinite_gradients = []
for name, parameter in (
model.named_parameters()
):
if not parameter.requires_grad:
continue
if parameter.grad is None:
missing_gradients.append(name)
continue
if not torch.isfinite(
parameter.grad
).all():
nonfinite_gradients.append(
name
)
print(
"Gradient norm:",
calculate_gradient_norm(model),
)
if missing_gradients:
print(
"Missing gradients:",
missing_gradients,
)
if nonfinite_gradients:
print(
"Nonfinite gradients:",
nonfinite_gradients,
)
assert not missing_gradients
assert not nonfinite_gradients
print("Gradient test passed.")
def test_optimizer_update(
model,
inputs,
targets,
):
print("\n7. Testing optimizer update")
optimizer = torch.optim.AdamW(
model.parameters(),
lr=1e-3,
)
name, parameter = next(
model.named_parameters()
)
before = parameter.detach().clone()
optimizer.zero_grad(
set_to_none=True
)
_, loss = model(
inputs,
targets,
)
loss.backward()
torch.nn.utils.clip_grad_norm_(
model.parameters(),
max_norm=1.0,
)
optimizer.step()
maximum_change = (
parameter.detach() - before
).abs().max().item()
print("Tracked parameter:", name)
print(
"Maximum parameter change:",
maximum_change,
)
assert maximum_change > 0
print(
"Optimizer update test passed."
)
def run_single_batch_overfit(
tokenizer,
train_ids,
device,
steps,
):
print(
"\n8. Running single-batch "
"overfitting test"
)
config = MiniGPTConfig(
vocab_size=tokenizer.vocab_size,
block_size=32,
embedding_dim=64,
num_heads=4,
num_layers=2,
expansion_factor=4,
dropout=0.0,
)
model = MiniGPT(config).to(device)
inputs, targets = get_batch(
data=train_ids,
batch_size=8,
block_size=config.block_size,
device=device,
)
optimizer = torch.optim.AdamW(
model.parameters(),
lr=1e-3,
)
initial_loss = None
final_loss = None
for step in range(steps):
optimizer.zero_grad(
set_to_none=True
)
_, loss = model(
inputs,
targets,
)
if initial_loss is None:
initial_loss = loss.item()
assert_finite_tensor(
loss,
"single-batch loss",
)
loss.backward()
torch.nn.utils.clip_grad_norm_(
model.parameters(),
max_norm=1.0,
)
optimizer.step()
final_loss = loss.item()
if (
step % 50 == 0
or step == steps - 1
):
print(
f"Step {step:4d}: "
f"loss {final_loss:.4f}"
)
print(
"Initial loss:",
initial_loss,
)
print(
"Final loss:",
final_loss,
)
assert final_loss < initial_loss
print(
"Single-batch overfitting "
"test passed."
)
def parse_args():
parser = argparse.ArgumentParser(
description=(
"Run Mini-GPT diagnostic tests"
)
)
parser.add_argument(
"--checkpoint",
type=str,
default=(
"checkpoints/mini_gpt_best.pt"
),
)
parser.add_argument(
"--tokenizer",
type=str,
default=(
"tokenizer/char_tokenizer.json"
),
)
parser.add_argument(
"--train-data",
type=str,
default=(
"data/encoded/train_ids.pt"
),
)
parser.add_argument(
"--overfit-steps",
type=int,
default=300,
)
parser.add_argument(
"--skip-overfit",
action="store_true",
)
return parser.parse_args()
def main():
args = parse_args()
torch.manual_seed(42)
device = get_device()
print("Using device:", device)
tokenizer = CharTokenizer.from_file(
args.tokenizer
)
train_ids = torch.load(
args.train_data,
map_location="cpu",
).long()
checkpoint_path = Path(
args.checkpoint
)
if not checkpoint_path.exists():
raise FileNotFoundError(
f"Checkpoint not found: "
f"{checkpoint_path}"
)
model, checkpoint = load_model(
checkpoint_path=checkpoint_path,
device=device,
)
if (
tokenizer.vocab_size
!= model.config.vocab_size
):
raise ValueError(
"Tokenizer and model vocabulary "
"sizes do not match"
)
print(
"Checkpoint step:",
checkpoint.get("step"),
)
test_tokenizer(
tokenizer,
train_ids,
)
test_token_ids(
tokenizer,
train_ids,
)
test_block_size = min(
32,
model.config.block_size,
)
inputs, targets = test_batch(
tokenizer=tokenizer,
train_ids=train_ids,
block_size=test_block_size,
device=device,
)
test_forward_pass(
model=model,
tokenizer=tokenizer,
inputs=inputs,
targets=targets,
)
test_future_independence(
model=model,
tokenizer=tokenizer,
device=device,
)
test_gradients(
model=model,
inputs=inputs,
targets=targets,
)
test_optimizer_update(
model=model,
inputs=inputs,
targets=targets,
)
if not args.skip_overfit:
run_single_batch_overfit(
tokenizer=tokenizer,
train_ids=train_ids,
device=device,
steps=args.overfit_steps,
)
print(
"\nAll requested diagnostics passed."
)
if __name__ == "__main__":
main()
Recent PyTorch releases changed the default behavior of torch.load toward loading weights only, so depending on your version and checkpoint contents you may need to set weights_only explicitly; check the current documentation.
Running the diagnostics
Run the full suite with default paths:
python debug_mini_gpt.py
Skip the overfitting stage for a quick check. The flag is --skip-overfit, with two leading hyphens:
python debug_mini_gpt.py - skip-overfit
Use another checkpoint:
python debug_mini_gpt.py \
--checkpoint checkpoints/mini_gpt_latest.pt
Give the overfitting test more steps:
python debug_mini_gpt.py \
--overfit-steps 500
The trailing backslash continues a command in Unix-style shells; if yours does not support it, put the command on one line.
The debugging order in practice
When the model misbehaves, walk these steps without skipping ahead.
Steps 1 to 5: data and shapes
Read the decoded data:
Does the tokenized dataset decode correctly?
Check target shifting:
Does inputs[:, 1:] equal targets[:, :-1]?
Check the token range:
Are all IDs between 0 and vocab_size - 1?
Check tensor shapes:
Inputs: [B, T]
Targets: [B, T]
Logits: [B, T, V]
Check the initial loss:
Is it near log(vocab_size) for a new model?
Steps 6 to 10: model behavior and training
Check causal independence:
Can changing the future affect prefix logits?
The only acceptable answer is no. Check gradients:
Are gradients present, finite, and nonzero?
Check parameter updates:
Does optimizer.step() change weights?
Memorize one batch:
Can the model memorize a tiny fixed batch?
Only then launch full training:
Only after all earlier tests pass should you invest in a long training run.
Checklists for each phase
Before training:
□ Dataset files exist
□ Tokenizer round-trip works
□ Token IDs are within vocabulary range
□ Decoded data looks correct
□ Inputs and targets are shifted by one
□ Batch tensors use torch.long
□ Model and batch use the same device
□ Logits have shape [B, T, V]
□ Initial loss is near log(V)
□ Future-independence test passes
□ All important parameters receive gradients
□ Optimizer changes parameters
□ Model can overfit one batch
During training:
□ Loss remains finite
□ Gradient norms remain finite
□ Learning rate follows the intended schedule
□ Training loss decreases
□ Validation loss is evaluated in eval mode
□ Best checkpoint updates when validation improves
□ Samples become more structured
During generation:
□ Best checkpoint is loaded
□ Matching tokenizer is loaded
□ Model is in evaluation mode
□ Context is cropped to block size
□ Only final-position logits are sampled
□ Temperature is positive
□ Top-k does not exceed vocabulary size
□ Repetition is measured, not only observed
Habits that make debugging harder
Changing many settings at once
If one experiment alters all of these together, you cannot attribute the result to any of them:
Learning rate
Batch size
Dropout
Model size
Context length
Change one major variable per experiment.
Judging only by samples
Weak text can come from undertraining, poor decoding, the wrong checkpoint or tokenizer, overfitting or underfitting, and samples cannot tell these apart. Look at metrics and pipeline tests first.
Silencing warnings
Warnings that mention reshaped tensors, a fallback to another device, NaN or infinite values, or unmatched checkpoint keys frequently signal genuine bugs; understand them before suppressing them.
Skipping smoke tests
Before a long run, start small:
Tiny model
Tiny batch
Short context
Few training steps
Then scale up gradually.
Treating clipping as a cure
Clipping can absorb one oversized update, but constant clipping means something deeper needs attention: the learning rate, initialization, loss scaling, numerical precision or anomalies in the data.
Exercises: break the pipeline on purpose
You trust a test more after watching it fail, so each exercise plants a known bug.
Remove the target shift
Make inputs and targets identical, confirm the batch test fails, then restore the shift.
Inject an out-of-range token
Set one token ID to:
tokenizer.vocab_size
The range check must fail, because the highest valid ID is:
vocab_size - 1
Disable the causal mask
Remove the mask temporarily and run the future-independence test; changing only the suffix should now change the prefix logits.
Use an absurd learning rate
Set:
learning_rate = 0.1
Track loss, gradient norm, parameter values and the nonfinite checks, and keep the run short.
Freeze the model
Apply the following and see how the gradient and optimizer tests report it:
for parameter in model.parameters():
parameter.requires_grad = False
Load into the wrong architecture
Load a checkpoint into a model that differs in one of these, then read the missing-key, unexpected-key and size-mismatch errors:
Vocabulary size
Embedding dimension
Number of layers
Compare dropout settings
Run the single-batch test with both values and compare how fast each memorizes:
dropout = 0.0
dropout = 0.2
Write a debug report
Save these results as JSON so runs can be compared and failures reproduced:
Tokenizer fingerprint
Vocabulary size
Token range
Batch shape
Initial loss
Expected baseline
Gradient norm
Missing gradients
Parameter update size
Single-batch final loss
Key takeaways
- A job that finishes without errors can still learn the wrong task.
- Debug in data-flow order; the round-trip test, range check and one-token shift assertion catch most data bugs.
- Logits must be
[B, T, V], and a fresh model should start nearlog(vocab_size). - The shared-prefix test proves causality through behavior.
- Gradient checks and a before-and-after parameter comparison confirm learning can happen; overfitting one batch confirms the whole loop.
- Nonfinite checks and hooks locate numerical failures; clipping only contains them.
- Train and validation curves separate overfitting from underfitting, and both model quality and decoding shape generated text.
- Rebuild models from the checkpoint's configuration and verify the tokenizer by fingerprint.
The full flow:
Environment
→ files
→ tokenizer
→ token IDs
→ batch shifting
→ shapes
→ initial loss
→ causal independence
→ gradients
→ optimizer updates
→ one-batch overfitting
→ full training
→ evaluation
→ generation
Behind all of it is one rule: do not debug by intuition; write a test that isolates one assumption, confirm it, and move on.
A natural next step is better data representation. Character tokens create long sequences, while word vocabularies grow huge; byte pair encoding learns merges for frequent character sequences, shortening sequences so the same context window holds more text. Adopting it means training merges, re-encoding WikiText-2 and resizing the model's vocabulary, and every check here applies unchanged:
Character tokens
→ learned subword merges
→ shorter sequences
→ better use of the context window