Debugging a Training Run

The Reality of Deep Learning Bugs

Deep learning bugs are uniquely challenging because models rarely crash with clear error stack traces.

Instead, a buggy model will silently train, consuming thousands of GPU dollars while producing suboptimal accuracy or stalling midway through training.

  SYSTEMATIC DEBUGGING PROTOCOL:
  [ Step 1: Overfit 1 Batch ] ──► [ Step 2: Step 0 Loss Check ] ──► [ Step 3: Gradient Norms ] ──► [ Step 4: Data Inspection ]

Step 1: The Single-Batch Overfit Test (Sanity Check #1)

Before launching a full training run on 1,000,000 samples:

Train your model on a single micro-batch of 10 samples for 100 steps:

# Single Batch Overfit Test
inputs, targets = next(iter(train_loader)) # Grab 1 mini batch!
for step in range(200):
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()
    if step % 20 == 0:
        print(f"Step {step} Loss: {loss.item():.4f}")

• EXPECTED RESULT: Loss should drop smoothly to near 0.0 within 100 steps.
• FAILURE: If loss stays flat or fluctuates, you have a critical bug in target labels, loss calculation, or layer gradient flow!

Step 2: Step Zero Initial Loss Check (Sanity Check #2)

Before any weight updates occur at step = 0, check the initial loss value:

For a classification task with $C$ classes using Cross-Entropy Loss:

$$L_{\text{initial}} = -\log\left(\frac{1}{C}\right) = \log(C)$$

• For $C = 10$ classes $\implies L \approx \log(10) = 2.302$.
• For $C = 1000$ classes $\implies L \approx \log(1000) = 6.907$.

If your initial loss is $45.0$ or $0.01$, your final classification layer weights are mis-initialized or loss scaling is incorrect.

Step 3: Monitoring Gradients & Activation Statistics

Track gradient norms and layer activations using Weights & Biases (W&B) or TensorBoard:

┌──────────────────────────┬──────────────────────────┬──────────────────────────┐
│ ISSUE                    │ SYMPTOM                  │ CAUSE & FIX              │
├──────────────────────────┼──────────────────────────┼──────────────────────────┤
│ Vanishing Gradients      │ Grad norm -> 0 in early  │ Deep layers without      │
│                          │ layers.                  │ Residuals / Dying ReLU.  │
├──────────────────────────┼──────────────────────────┼──────────────────────────┤
│ Exploding Gradients      │ Grad norm -> ∞ / NaN     │ Learning rate too high.  │
│                          │ spikes.                  │ Enable Gradient Clipping!│
├──────────────────────────┼──────────────────────────┼──────────────────────────┤
│ Dead Activations         │ 90% layer outputs zero.  │ Negative biases in ReLU. │
│                          │                          │ Switch to LeakyReLU.     │
└──────────────────────────┴──────────────────────────┴──────────────────────────┘

Step 4: Inspecting Data Pipeline Batches

$80%$ of deep learning bugs stem from data pipeline errors:

1. Un-normalized Inputs: Forgetting to scale image pixels to $[0, 1]$ or standardizing tabular features ($\mu=0, \sigma=1$).
2. Label Flipping: Off-by-one errors in classification target indices.
3. Data Leakage: Applying data augmentation to validation splits or leaking target labels into input feature columns.

Say this out loud

Systematic training debugging starts by overfitting a single micro batch of 10 samples to zero loss to verify backpropagation logic. Step zero loss must match log C random initialization. Tracking gradient norms and activation statistics isolates vanishing or exploding gradients, while inspecting data pipelines catches un normalized inputs and label leakage.

Followups to expect

1. What is torch.autograd.set_detect_anomaly(True)? A PyTorch debugging context manager that tracks backward graph operations, throwing an error pinpointing the exact forward layer that created NaN or Inf tensors.
2. How do you debug loss spikes midway through LLM pretraining? Roll back to a checkpoint 500 steps prior, skip or re-weight the corrupt data batch that caused the spike, and resume training with a temporary lower learning rate.