PyTorch Debugging, Lab & Modern Performance

I used to write for loops. Then I discovered tensors, and now I judge everyone who still writes for loops.

Tensor, Tensor-Evangelizing AI Agent

This continuation of Section 0.3 picks up after you have the basic PyTorch training loop in hand. It covers the debugging tools that turn a non-training model into a training one, a hands-on FashionMNIST lab that exercises every concept from 0.3a, and the modern PyTorch features (torch.compile, AMP, FSDP) that move you from a working prototype to a fast one. Where 0.3a built the workbench, 0.3b is about using it productively: inspecting what your model is actually doing, fixing it when it silently produces wrong results, and pushing throughput on real hardware.

0.4.1 Debugging: Hooks, Gradient Inspection, and Profiling

When your model does not train, you need tools to look inside. PyTorch provides several mechanisms for introspection.

0.4.1.1 Inspecting Gradients

After a backward pass, you can iterate over named parameters to check gradient statistics for signs of vanishing or exploding gradients.

# Check gradients after a backward pass
for name, param in model.named_parameters():
    if param.grad is not None:
        print(f"{name:20s} grad mean={param.grad.mean():.6f} "
            f"std={param.grad.std():.6f}")

0.4.1.2 Forward and Backward Hooks

Hooks let you inspect (or modify) data flowing through a module without changing its code. This is invaluable for debugging and later for techniques like activation patching in interpretability research.

import torch
# Register a forward hook that prints the output shape
def print_shape_hook(module, input, output):
    print(f"{module.__class__.__name__:15s} output shape: {output.shape}")
    hooks = []
    for name, layer in model.named_children():
        h = layer.register_forward_hook(print_shape_hook)
        hooks.append(h)
        # Run one forward pass to see shapes
        dummy = torch.randn(1, 784).to(device)
        _ = model(dummy)
        # Clean up hooks when done
        for h in hooks:
            h.remove()

0.4.1.3 Profiling with torch.profiler

The built-in profiler measures CPU and GPU time per operation, helping you identify performance bottlenecks.

# Profile a few training batches with torch.profiler to identify
# which operations (linear, cross_entropy, relu) consume the most CPU time.
from torch.profiler import profile, ProfilerActivity

# Profile execution to find performance bottlenecks
with profile(activities=[ProfilerActivity.CPU], record_shapes=True) as prof:
    for i, (images, labels) in enumerate(train_loader):
        images = images.view(images.size(0), -1)
        outputs = model(images)
        loss = criterion(outputs, labels)
        # Compute gradients via backpropagation
        loss.backward()
        if i >= 4:
            break

            print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=5))

0.4.2 Common Mistakes and How to Fix Them

#!/usr/bin/env python3
"""Lab 0.3: FashionMNIST Classifier in PyTorch (from scratch)."""

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# -- Hyperparameters ------------------------------------------
BATCH_SIZE = 64
LEARNING_RATE = 1e-3
NUM_EPOCHS = 10
HIDDEN_DIM = 256

# -- Device ---------------------------------------------------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Training on: {device}")

# -- Data -----------------------------------------------------
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.2860,), (0.3530,)),
    ])

train_data = datasets.FashionMNIST("./data", train=True, download=True, transform=transform)
test_data = datasets.FashionMNIST("./data", train=False, download=True, transform=transform)

train_loader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(test_data, batch_size=BATCH_SIZE, shuffle=False)

# -- Model ----------------------------------------------------
class FashionClassifier(nn.Module):
    def __init__(self, hidden_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Flatten(), # (B,1,28,28) -> (B,784)
            nn.Linear(784, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, 10),
            )

    def forward(self, x):
        return self.net(x)

        model = FashionClassifier(HIDDEN_DIM).to(device)
        print(model)
        print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")

        # -- Loss and Optimizer ---------------------------------------
        criterion = nn.CrossEntropyLoss()
        optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)

        # -- Training -------------------------------------------------
    def train_one_epoch(model, loader, criterion, optimizer, device):
        model.train()
        total_loss, correct, total = 0.0, 0, 0

        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)

            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            total_loss += loss.item() * labels.size(0)
            correct += (outputs.argmax(1) == labels).sum().item()
            total += labels.size(0)

            return total_loss / total, correct / total

            # -- Evaluation -----------------------------------------------
def evaluate(model, loader, criterion, device):
    model.eval()
    total_loss, correct, total = 0.0, 0, 0

    with torch.no_grad():
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)

            total_loss += loss.item() * labels.size(0)
            correct += (outputs.argmax(1) == labels).sum().item()
            total += labels.size(0)

            return total_loss / total, correct / total

            # -- Run ------------------------------------------------------
            for epoch in range(NUM_EPOCHS):
                train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer, device)
                test_loss, test_acc = evaluate(model, test_loader, criterion, device)

                print(f"Epoch {epoch+1:2d}/{NUM_EPOCHS} "
                    f"Train Loss: {train_loss:.4f} Acc: {train_acc:.4f} "
                    f"Test Loss: {test_loss:.4f} Acc: {test_acc:.4f}")

                # -- Save -----------------------------------------------------
                torch.save({
                    "model_state_dict": model.state_dict(),
                    "optimizer_state_dict": optimizer.state_dict(),
                    "test_acc": test_acc,
                    }, "fashion_classifier_checkpoint.pth")
                print(f"\nModel saved. Final test accuracy: {test_acc:.4f}")

0.4.3 Modern PyTorch: Performance and Scale

The training loop and model patterns covered so far are the foundation of every PyTorch project. However, modern deep learning, particularly LLM training and inference, demands tools that go beyond the basics. PyTorch 2.x introduced a compiler, and the ecosystem provides built-in support for mixed precision and distributed training. This section covers the three most important performance tools you will encounter in practice.

0.4.3.1 torch.compile and PyTorch 2.x

Starting with PyTorch 2.0, torch.compile transforms your eager-mode model into an optimized graph that runs significantly faster. Under the hood, it uses TorchDynamo to capture the computation graph from Python bytecode, then passes that graph to the TorchInductor compiler backend, which generates optimized Triton (GPU) or C++/OpenMP (CPU) kernels. The key insight is that you do not need to change your model code at all; you simply wrap it with torch.compile() and let the compiler handle fusion, memory planning, and kernel selection.

torch.compile offers three compilation modes, each trading compile time for runtime speed:

A few common pitfalls to watch for: (1) the first call triggers compilation, so you will see a one-time latency spike; (2) data-dependent control flow (e.g., if x.sum() > 0) causes "graph breaks" that reduce optimization opportunities; and (3) not all custom CUDA extensions are supported yet. When in doubt, start with default mode and profile.

# torch.compile: wrap a model for optimized GPU kernel generation.
# The first call triggers compilation; subsequent calls run faster.
import torch

# Define a simple model
model = MyTransformerBlock(d_model=512, n_heads=8).cuda()

# Without torch.compile: standard eager execution
output_eager = model(input_tensor)

# With torch.compile: optimized execution
compiled_model = torch.compile(model, mode="reduce-overhead")

# First call triggers compilation (slow), subsequent calls are fast
output_compiled = compiled_model(input_tensor)

# In benchmarks, expect 1.3x to 2x speedup on Transformer blocks

Advanced torch.compile: Dynamic Shapes, Fullgraph Mode, and Debugging

Getting the most out of torch.compile in production requires understanding three additional concepts beyond the basic wrapper. First, dynamic shapes: by default, the compiler assumes fixed input shapes and triggers a full recompilation whenever the shape changes. For NLP workloads where sequence lengths vary across batches, this causes repeated compilations that negate any speedup. Setting dynamic=True tells the compiler to generate shape-generic kernels that work across a range of input sizes, at the cost of slightly less aggressive optimization for any single shape. In Transformer training with variable-length sequences, dynamic=True is almost always the right choice.

Second, fullgraph mode: the fullgraph=True option tells the compiler to capture the entire model as a single graph, which enables global optimizations but will raise an error if any graph break occurs. This is useful for validating that your model is fully compilable before deploying to production. If graph breaks are present, the compiler silently falls back to partial compilation, which may deliver only modest speedups. Running with fullgraph=True during development ensures you catch and eliminate graph breaks early.

Third, debugging and profiling: the torch._dynamo module exposes configuration flags that help you understand what the compiler is doing. Setting torch._dynamo.config.verbose = True logs every graph break with a traceback, making it straightforward to identify problematic code patterns. The torch.utils.benchmark module provides a clean way to compare eager and compiled execution times with statistically meaningful measurements.

# Strict mode: fails if any graph break is detected
compiled_strict = torch.compile(model, fullgraph=True)

# Dynamic shapes: avoid recompilation when input sizes change
compiled_dynamic = torch.compile(model, dynamic=True)

# Combine max-autotune with fullgraph for production
compiled_prod = torch.compile(
 model,
 mode="max-autotune",
 fullgraph=True,
 dynamic=True,
)

# Debugging: see what the compiler is doing
import torch._dynamo
torch._dynamo.config.verbose = True # Log graph breaks with tracebacks
torch._dynamo.config.suppress_errors = False # Fail loudly on issues

# Profile compiled vs. eager to measure actual speedup
import torch.utils.benchmark as bench

timer_eager = bench.Timer(
 stmt="model(x)",
 globals={"model": model, "x": input_tensor},
)
timer_compiled = bench.Timer(
 stmt="compiled_model(x)",
 globals={"compiled_model": compiled_prod, "x": input_tensor},
)

print(f"Eager: {timer_eager.timeit(100).mean * 1000:.2f} ms")
print(f"Compiled: {timer_compiled.timeit(100).mean * 1000:.2f} ms")

torch.export: Deployment Beyond Python

PyTorch 2.x also introduced torch.export, which captures a model as a clean, self-contained graph representation suitable for deployment outside of Python. While torch.compile accelerates training and eager-mode inference, torch.export targets production deployment scenarios: shipping a model to a mobile device, embedding it in a C++ application, or converting it to a format consumed by a purpose-built serving stack. The exported graph can be lowered to backends like ExecuTorch (for edge and mobile devices) or AOTInductor (for server deployment without the Python runtime overhead).

import torch
# torch.export: capture a deployment-ready graph
from torch.export import export
# Define example inputs for tracing
example_input = torch.randn(1, 128, 512).cuda()
# Export the model (captures the full graph)
exported = export(model, (example_input,))
# The exported program can be serialized and loaded without Python
torch.export.save(exported, "model_exported.pt2")
# For server deployment with AOTInductor (generates a .so library)
# torch._inductor.aot_compile(model, (example_input,))

Combining torch.compile with Mixed Precision

In practice, torch.compile and mixed precision are used together rather than in isolation. The compiler is aware of autocast regions and can fuse operations across precision boundaries, generating kernels that perform the cast and the computation in a single step. This combination typically yields the best results: mixed precision reduces memory traffic and enables Tensor Core utilization, while the compiler eliminates kernel launch overhead and fuses adjacent operations. The following example shows the recommended production pattern that combines both techniques.

# Combine torch.compile (max-autotune) with BF16 autocast.
# The compiler fuses cast and compute into single GPU kernels.
import torch
from torch.amp import autocast

# Compile the model first
model = MyTransformerBlock(d_model=512, n_heads=8).cuda()
compiled_model = torch.compile(model, mode="max-autotune", dynamic=True)

optimizer = torch.optim.AdamW(compiled_model.parameters(), lr=3e-4)

for batch_x, batch_y in train_loader:
    batch_x, batch_y = batch_x.cuda(), batch_y.cuda()
    optimizer.zero_grad()

    # BF16 autocast inside the compiled model: the compiler fuses casts
    with autocast(device_type="cuda", dtype=torch.bfloat16):
        output = compiled_model(batch_x)
        loss = criterion(output, batch_y)

        loss.backward()
        optimizer.step()

        # On Ampere+ GPUs, this pattern typically yields 2x to 3x throughput
        # improvement over eager FP32 execution.

0.4.3.2 Mixed Precision Training with torch.amp

Modern GPUs have specialized hardware (Tensor Cores) that operate much faster on 16-bit floating-point numbers than on 32-bit. Mixed precision training uses 16-bit for most operations (forward pass, backward pass) while keeping a 32-bit master copy of the weights for the optimizer update. This roughly halves memory usage and can double training throughput.

PyTorch provides torch.amp (Automatic Mixed Precision) with two components: torch.amp.autocast automatically selects the right precision for each operation (matmuls in FP16/BF16, reductions in FP32), and torch.amp.GradScaler prevents underflow by scaling the loss before the backward pass and unscaling gradients before the optimizer step. On Ampere GPUs (A100, RTX 3090) and newer, BF16 (bfloat16) is preferred over FP16 because it has the same exponent range as FP32, which eliminates most overflow/underflow issues and makes GradScaler unnecessary.

# Mixed-precision training with GradScaler (FP16) and autocast.
# GradScaler prevents gradient underflow; skip it when using BF16.
import torch
from torch.amp import autocast, GradScaler

model = MyModel().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
scaler = GradScaler() # Only needed for FP16; skip for BF16

for epoch in range(num_epochs):
    for batch_x, batch_y in train_loader:
        batch_x, batch_y = batch_x.cuda(), batch_y.cuda()
        optimizer.zero_grad()

        # Forward pass in mixed precision
        with autocast(device_type="cuda", dtype=torch.float16):
            output = model(batch_x)
            loss = criterion(output, batch_y)

            # Backward pass with gradient scaling
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()

            # For BF16 (preferred on Ampere+ GPUs), simply use:
                # with autocast(device_type="cuda", dtype=torch.bfloat16):
                    # output = model(batch_x)
                    # loss = criterion(output, batch_y)
                    # loss.backward() # No scaler needed
                    # optimizer.step()

0.4.3.3 Distributed Data Parallel (DDP)

When a single GPU is not enough, torch.nn.parallel.DistributedDataParallel (DDP) is the standard way to scale training across multiple GPUs (or multiple machines). DDP replicates the model on each GPU, splits each batch across the replicas, and synchronizes gradients with an all-reduce operation after each backward pass. Because each GPU processes a different slice of the data, the effective batch size scales linearly with the number of GPUs.

Concretely, DDP estimates the gradient of the per-example loss as the average over the global batch of size $B = K \cdot b$, where $K$ is the number of replicas and $b$ the per-GPU micro-batch:

$$ g \;=\; \frac{1}{B} \sum_{i=1}^{B} \nabla_{\theta} \ell(x_i; \theta) \;=\; \frac{1}{K} \sum_{k=1}^{K} \underbrace{\frac{1}{b} \sum_{i \in \text{shard}_k} \nabla_{\theta} \ell(x_i; \theta)}_{g_k\, \text{computed on GPU}\, k} $$

That second equality is exactly what the all-reduce computes: each GPU first averages over its own shard locally, then the ring all-reduce averages the $g_k$ across replicas. The result is mathematically identical to running a single GPU at the full batch size $B$, which is why the linear scaling rule (multiply the learning rate by $K$) usually holds.

DDP is preferred over the older DataParallel because it avoids the GIL bottleneck and overlaps communication with computation. Setting it up requires initializing a process group and wrapping your model, but the training loop itself remains almost identical to the single-GPU version. For LLM training at larger scales, you will encounter FSDP (Fully Sharded Data Parallel), which shards both parameters and gradients across GPUs. We will revisit distributed training in Chapter 06 when we discuss pretraining.

import os
# Distributed Data Parallel: initialize a process group, wrap the model,
# and train with automatic gradient synchronization across GPUs.
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
# Initialize the process group (one process per GPU)
dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
# Create model and wrap with DDP
model = MyModel().cuda(local_rank)
model = DDP(model, device_ids=[local_rank])
# Training loop is the same as single-GPU
for batch_x, batch_y in train_loader:
    optimizer.zero_grad()
    output = model(batch_x.cuda(local_rank))
    loss = criterion(output, batch_y.cuda(local_rank))
    loss.backward() # DDP handles gradient sync automatically
    optimizer.step()
    # Launch with: torchrun --nproc_per_node=4 train.py

0.4.3.4 DDP in Practice: What Happens Under the Hood

Understanding DDP's mechanics helps you debug distributed training issues and make informed choices about scaling. When you wrap a model with DistributedDataParallel, three things happen at initialization: (1) the model parameters are broadcast from rank 0 to all other processes, ensuring every GPU starts with identical weights; (2) DDP registers backward hooks on every parameter, which trigger gradient synchronization automatically; and (3) parameters are grouped into "buckets" for communication efficiency, so that all-reduce operations overlap with backward computation.

The bucket-based overlap is critical for performance. Rather than waiting until all gradients are computed and then performing a single all-reduce, DDP starts synchronizing the gradients of later layers (which finish their backward pass first) while earlier layers are still computing. This overlap means that for well-balanced models, communication is almost entirely hidden behind computation. You can control bucket size with the bucket_cap_mb parameter (default: 25 MB).

A few practical details matter when using DDP:

• DistributedSampler: Each GPU must see a different subset of the data. Use torch.utils.data.distributed.DistributedSampler with your DataLoader to ensure non-overlapping splits. Remember to call sampler.set_epoch(epoch) at the start of each epoch so that shuffling differs across epochs.
• Batch size scaling: The effective batch size is per_gpu_batch_size * num_gpus. If you increase from 1 GPU to 8 GPUs, the effective batch size grows 8x. You may need to adjust the learning rate accordingly (linear scaling rule: multiply learning rate by the same factor).
• Saving checkpoints: Only save from rank 0 to avoid duplicate writes. Guard your save logic with if dist.get_rank() == 0.
• Launching: Use torchrun (or torch.distributed.launch) to spawn one process per GPU. For multi-node training, you also need to set --nnodes, --node_rank, and --master_addr.

from torch.utils.data import DataLoader
import torch
# Complete DDP training setup with DistributedSampler
from torch.utils.data.distributed import DistributedSampler
sampler = DistributedSampler(train_dataset, shuffle=True)
train_loader = DataLoader(train_dataset, batch_size=32, sampler=sampler)
for epoch in range(num_epochs):
    sampler.set_epoch(epoch) # Ensure different shuffling each epoch
    for batch_x, batch_y in train_loader:
        optimizer.zero_grad()
        output = model(batch_x.cuda(local_rank))
        loss = criterion(output, batch_y.cuda(local_rank))
        loss.backward()
        optimizer.step()
        # Save only from rank 0
        if dist.get_rank() == 0:
            torch.save(model.module.state_dict(), f"checkpoint_epoch_{epoch}.pt")

Note the use of model.module.state_dict() rather than model.state_dict() when saving. The DDP wrapper adds a .module attribute that references the original model. Saving through .module produces a state dict compatible with non-DDP loading, which is almost always what you want.

Exercises

import torch; from torch import nn
X = torch.randn(100, 5); y = X @ torch.randn(5, 1) + 0.1 * torch.randn(100, 1)
model = nn.Linear(5, 1)
opt = torch.optim.SGD(model.parameters(), lr=0.01)
loss_fn = nn.MSELoss()
for epoch in range(100):
  opt.zero_grad()
  pred = model(X)
  loss = loss_fn(pred, y)
  loss.backward()
  opt.step()
  if epoch % 10 == 0: print(f"epoch {epoch} loss {loss.item():.4f}")