How to Debug Neural Networks by Inspecting Gradients, Detecting NaNs, and Analyzing Loss Curves

Debug neural networks effectively by monitoring gradient norms, automatically detecting NaN or Inf values during backpropagation, and persisting step-wise metrics to CSV for post-hoc visualization of loss curves.

Debugging unstable training runs requires real-time visibility into the optimization process. The rohitg00/ai-engineering-from-scratch repository provides a production-ready toolkit in Phase 19, Capstone Project 45 – “Gradient Clipping & AMP” that demonstrates how to debug neural networks through systematic gradient inspection, finite-value validation, and structured logging.

Detecting NaNs and Infinite Values in Gradients

The first pillar of robust debugging is immediate detection of non-finite gradients. Any NaN or Inf value in a gradient signals a pathological training step caused by exploding activations, division-by-zero errors, or numerical overflow.

In phases/19-capstone-projects/45-gradient-clipping-amp/code/main.py at line 76, the has_non_finite_grad function scans every parameter's .grad tensor:

def has_non_finite_grad(model: nn.Module) -> bool:
    """Return True if any parameter gradient contains NaN or Inf."""
    for p in model.parameters():
        if p.grad is not None and not torch.all(torch.isfinite(p.grad)):
            return True
    return False

When integrated into the training loop, this utility enables fail-fast behavior: if a gradient contains non-finite values, the step is aborted and the reason is recorded as "non_finite_grad" before the optimizer updates weights.

Monitoring Gradient Magnitudes to Prevent Explosions

The second pillar involves quantifying gradient health through L2 norm calculation. Vanishing gradients (norms approaching zero) or exploding gradients (norms exceeding thresholds) indicate architectural or hyperparameter issues.

The repository provides two complementary utilities in main.py:

• compute_global_l2_norm (line 88): Concatenates all gradients (ignoring None entries) and computes the Euclidean norm of the resulting vector.
• clip_global_l2_norm (line 100): Clips the global gradient norm to a configurable max_norm, returning both pre-clip and post-clip values for diagnostic logging.

def compute_global_l2_norm(model: nn.Module) -> float:
    """Compute the L2 norm of all gradients combined."""
    grads = [p.grad.detach() for p in model.parameters() if p.grad is not None]
    if not grads:
        return 0.0
    return torch.cat([g.flatten() for g in grads]).norm(2).item()

Tracking grad_l2_pre_clip across epochs exposes whether your learning rate is too aggressive or if specific layers are experiencing disproportionate gradient growth.

Orchestrating Diagnostics with AmpTrainState

The AmpTrainState class (line 125 in main.py) orchestrates the complete debugging workflow. Its step method implements a defensive training protocol:

1. Runs forward passes under torch.amp.autocast for automatic mixed precision
2. Aborts backward computation if the loss is non-finite
3. Validates gradient finiteness using has_non_finite_grad
4. Applies global L2 clipping when norms exceed the threshold
5. Records a structured StepLog entry with metrics including loss, learning rate, gradient norms, and skip reasons

This automation ensures that corrupted steps never update model parameters, preventing the propagation of NaN values through subsequent training iterations.

Logging Step-by-Step Metrics for Loss Curve Analysis

The third pillar requires persistent storage of diagnostic data for offline analysis. The repository defines a StepLog dataclass (line 39) that captures:

• Step number and learning rate
• Pre-clip and post-clip gradient L2 norms
• Loss value and skip flags
• Skip reason ("non_finite_loss" or "non_finite_grad")
• Scaler scale for mixed-precision tracking

The write_step_log_csv helper exports these records to a CSV file, enabling you to analyze loss curves and gradient trajectories using standard data science tools.

Practical Implementation Example

The following snippet demonstrates how to integrate these utilities into a custom PyTorch training loop using the repository's debugging toolkit:

import torch
from pathlib import Path
from main import AmpTrainState, write_step_log_csv, has_non_finite_grad

# Build a tiny model and dummy data

model, inputs, targets = build_toy_model()  # Returns nn.Module, Tensor, Tensor

# Initialize the AMP-aware trainer with gradient clipping

trainer = AmpTrainState(
    model=model,
    lr=1e-2,
    max_norm=1.0,          # Clip gradients greater than 1.0

    device_type="cpu",     # Use "cuda" for GPU training

)

# Run training steps with optional gradient corruption for testing

for epoch in range(5):
    # Inject NaN on epoch 3 to demonstrate skip logic

    corruptor = (lambda m: m.linear.weight.grad.__setitem__(slice(None), float("nan"))) if epoch == 3 else None
    log = trainer.step(inputs, targets, gradient_corruptor=corruptor)
    
    print(f"Epoch {epoch} → loss={log.loss:.4f}, grad_norm={log.grad_l2_pre_clip:.4f}"
          f"{' [SKIPPED]' if log.skipped else ''}")

# Export the full log for offline analysis

log_path = Path("training_log.csv")
write_step_log_csv(trainer.log, log_path)
print(f"Full training log written to {log_path}")

Expected output:

Epoch 0 → loss=0.0372, grad_norm=0.1245
Epoch 1 → loss=0.0321, grad_norm=0.0987
Epoch 2 → loss=0.0289, grad_norm=0.0823
Epoch 3 → loss=nan, grad_norm=nan [SKIPPED]
Epoch 4 → loss=0.0254, grad_norm=0.0679
Full training log written to training_log.csv

Visualizing Training Logs to Identify Failure Points

Once exported, the CSV enables correlation analysis between loss spikes and gradient explosions:

import pandas as pd
import matplotlib.pyplot as plt

df = pd.read_csv("training_log.csv")
plt.plot(df["step"], df["loss"], label="Loss")
plt.plot(df["step"], df["grad_l2_pre_clip"], label="Grad L2 (pre-clip)")
plt.legend()
plt.xlabel("Step")
plt.title("Training Debugging Curve")
plt.show()

Plotting grad_l2_pre_clip against loss reveals whether spikes in gradient magnitude precede divergence events, confirming the effectiveness of your clipping strategy.

Summary

• Detect non-finite values using has_non_finite_grad in main.py to catch numerical instabilities immediately and prevent parameter corruption.
• Quantify gradient health with compute_global_l2_norm to identify vanishing or exploding gradient regimes before they destabilize training.
• Stabilize training via clip_global_l2_norm when global norms exceed your configured max_norm threshold.
• Automate defensive training using AmpTrainState to orchestrate loss validation, gradient scanning, and clipping in a single step method.
• Enable post-hoc analysis by exporting structured logs with write_step_log_csv to visualize loss curves and diagnose failure points offline.

Frequently Asked Questions

Why do gradients become NaN during neural network training?

Gradients become NaN due to exploding activations in forward passes, division-by-zero operations in loss calculations, or numerical overflow in mixed-precision training. According to the ai-engineering-from-scratch implementation, scanning parameter gradients with has_non_finite_grad immediately after loss.backward() catches these instabilities before they propagate to model weights.

What is the ideal threshold for gradient clipping?

While the repository uses max_norm=1.0 as a default in clip_global_l2_norm, optimal thresholds typically range between 0.5 and 5.0 depending on model architecture and initialization. You should tune this value by monitoring the grad_l2_pre_clip values logged during initial epochs; if pre-clip norms consistently exceed your threshold by orders of magnitude, consider reducing learning rate or increasing the clipping bound.

How can I visualize loss curves from PyTorch training logs?

Export step-wise metrics using write_step_log_csv to generate a structured CSV file containing loss values and gradient norms. Load this data with pandas and plot using matplotlib to visualize loss trajectories alongside gradient magnitude correlations, as demonstrated in Phase 19, Capstone Project 45 of the repository.

Where does the training loop check for non-finite loss values?

Within AmpTrainState.step in main.py, the implementation validates loss finiteness immediately after the backward pass using torch.isfinite(). If the loss contains NaN or Inf, the method sets skipped=True and records skip_reason="non_finite_loss", preventing the optimizer from updating parameters with corrupted gradients.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}