Training & Debugging: Gradient Clipping, LR Scheduling, NaN Detection

🎯 Learning Objectives

• Diagnose and fix common training failures (NaN, vanishing gradients, loss plateaus)
• Implement gradient clipping and learning rate scheduling
• Use loss analysis to identify overfitting, underfitting, and data issues
• Debug model training systematically

📋 Prerequisites

• PyTorch basics (Week 1): Training loops, optimizers
• Data Pipelines (Week 2): Data loading

1. 📖 Core Content

1.1 The Training Debugging Mindset

1.2 Gradient Clipping

python
Run Inline
# runnable
import torch
import torch.nn as nn
model = nn.Linear(10, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
loss_fn = nn.MSELoss()
# Standard training step with gradient clipping
inputs = torch.randn(32, 10)
targets = torch.randn(32, 1)
outputs = model(inputs)
loss = loss_fn(outputs, targets)
loss.backward()
# Clip gradients: scale down if norm > max_norm
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
# Check gradient norm
total_norm = 0
for p in model.parameters():
    if p.grad is not None:
        param_norm = p.grad.data.norm(2)
        total_norm += param_norm.item() ** 2
total_norm = total_norm ** 0.5
print(f"Gradient norm after clipping: {total_norm:.4f}")
optimizer.step()
optimizer.zero_grad()

1.3 Learning Rate Scheduling

python
Run Inline
# runnable
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
model = nn.Linear(10, 1)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
# Common schedulers
schedulers = {
    'StepLR': torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1),
    'CosineAnnealingLR': torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100),
    'ReduceLROnPlateau': torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode='min', factor=0.5, patience=10
    ),
    'OneCycleLR': torch.optim.lr_scheduler.OneCycleLR(
        optimizer, max_lr=0.1, total_steps=100
    )
}
# Simulate training to see LR schedules
lrs = []
sch = schedulers['CosineAnnealingLR']
for epoch in range(100):
    lrs.append(optimizer.param_groups[0]['lr'])
    # Simulated loss
    loss = torch.tensor(0.5 * (0.95 ** epoch) + 0.1 * torch.randn(1))
    loss.backward()
    optimizer.step()
    sch.step()  # Update LR each epoch
# Plot would show cosine decay from 0.01 to nearly 0
print(f"Starting LR: {lrs[0]:.6f}")
print(f"Ending LR: {lrs[-1]:.6f}")

1.4 NaN Detection and Handling

python
Run Inline
# runnable
import torch
import torch.nn as nn
def detect_nan(model, loss):
    """Check for NaN/Inf in loss and gradients."""
    if torch.isnan(loss) or torch.isinf(loss):
        print(f"🚨 Loss is {loss.item()}!")
        return True
    for name, param in model.named_parameters():
        if param.grad is not None:
            if torch.isnan(param.grad).any():
                print(f"🚨 NaN gradient in {name}")
                return True
            if torch.isinf(param.grad).any():
                print(f"🚨 Inf gradient in {name}")
                return True
    return False
# Safety wrapper for training step
def safe_training_step(model, inputs, targets, optimizer, loss_fn, clip_norm=1.0):
    model.train()
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = loss_fn(outputs, targets)
    # NaN check before backward
    if torch.isnan(loss) or torch.isinf(loss):
        print("Skipping batch: NaN/Inf in loss")
        optimizer.zero_grad()
        return None
    loss.backward()
    # Clip gradients
    torch.nn.utils.clip_grad_norm_(model.parameters(), clip_norm)
    # NaN check after clipping
    if detect_nan(model, loss):
        optimizer.zero_grad()
        return None
    optimizer.step()
    return loss.item()

1.5 Common Debugging Workflow

python
Run Inline
# runnable
def diagnose_training(model, train_loader, val_loader, config):
    """Systematic training diagnosis."""
    print("=" * 50)
    print("TRAINING DIAGNOSIS")
    print("=" * 50)
    # 1. Verify data pipeline
    sample_batch, _ = next(iter(train_loader))
    print(f"✅ Data loaded: batch shape {sample_batch.shape}")
    print(f"   Data range: [{sample_batch.min():.3f}, {sample_batch.max():.3f}]")
    # 2. Check for NaN/Inf in data
    if torch.isnan(sample_batch).any():
        print("❌ NaN detected in input data!")
    # 3. Forward pass test
    try:
        output = model(sample_batch[:4])
        print(f"✅ Forward pass OK: output shape {output.shape}")
    except Exception as e:
        print(f"❌ Forward pass failed: {e}")
        return
    # 4. Backward pass test
    loss = output.sum()
    loss.backward()
    grad_norms = [p.grad.norm().item() for p in model.parameters() if p.grad is not None]
    print(f"✅ Backward pass OK: gradient norms [{min(grad_norms):.6f}, {max(grad_norms):.6f}]")
    # 5. Check for vanishing/exploding gradients
    if max(grad_norms) > 100:
        print("⚠️ Exploding gradients detected! Consider gradient clipping.")
    if max(grad_norms) < 1e-6:
        print("⚠️ Vanishing gradients detected! Check architecture.")
    print("=" * 50)
    print("Diagnosis complete. Ready for training.")

1.6 Why This Matters

2. 📐 Key Formulas / Concepts

3. ⚠️ Common Pitfalls

Pitfall 1: Not Zeroing Gradients

Pitfall 2: Evaluating Model in train() Mode

Pitfall 3: Wrong Loss Reduction

4. 📝 Practice Questions

Q1: Your training loss is consistently 8.35 after 100 epochs and won't decrease. What's likely wrong?

Possible causes:

1. Learning rate too low: Loss stuck at a high value suggests the optimizer can't escape a local minimum. Try LR=0.001 or OneCycleLR.
2. Wrong loss function: If you're using MSE for classification, the loss floor is higher. Verify the loss matches the task.
3. Data normalization: If inputs range [0, 255] (images) or have very different scales (age=30, income=100000), the model may struggle. Normalize inputs to mean=0, std=1.
4. Model capacity: The model may be too small. Add more layers or units.

Debug: Plot the LR, log the gradient norms, verify input scaling. Q2: Write a training loop that detects and logs NaN losses, skips bad batches, and saves the best model checkpoint.

python
Run Inline
best_val_loss = float('inf')
for epoch in range(num_epochs):
    for batch_idx, (inputs, targets) in enumerate(train_loader):
        loss = safe_training_step(model, inputs, targets, optimizer, loss_fn)

        if loss is None:  # NaN detected
            print(f"Epoch {epoch}, Batch {batch_idx}: Skipped (NaN)")
            continue

        if batch_idx % 100 == 0:
            print(f"Epoch {epoch}, Batch {batch_idx}: Loss = {loss:.4f}")

    # Validation
    model.eval()
    val_loss = 0
    with torch.no_grad():
        for inputs, targets in val_loader:
            outputs = model(inputs)
            val_loss += loss_fn(outputs, targets).item()
    val_loss /= len(val_loader)

    # Save best model
    if val_loss < best_val_loss:
        best_val_loss = val_loss
        torch.save(model.state_dict(), 'best_model.pth')
        print(f"✅ Saved best model (val_loss: {val_loss:.4f})")

5. 🔗 Cross-References

• Previous: Data Loading & Pipelines (Week 2)
• Next: Distributed Training (Week 5)
• Related: Profiling & Optimization (Week 7)
• External: PyTorch Lightning — Automated training with built-in debugging Join Discord PreviousPandas for DLNextDistributed Training