Debugging and Profiling Techniques for Developing and Optimizing AI Systems: A Three-Layer Framework

The AI Engineering from Scratch curriculum teaches a systematic three-layer framework that combines standard Python debugging tools, tensor validation utilities, and training dynamics visualizers to diagnose issues ranging from import errors to multi-agent production pipeline failures.

The AI Engineering from Scratch repository provides a battle-tested curriculum for building production-ready machine learning systems. According to the source code in phases/00-setup-and-tooling/12-debugging-and-profiling/docs/en.md, effective debugging requires a layered approach that catches bugs at the Python level, validates tensor operations, and monitors training dynamics. This methodology scales from simple scripts to complex multi-agent architectures.

Layer 1: Standard Python Debugging and Profiling

Standard Python tooling forms the foundation for catching syntax errors, import failures, and performance bottlenecks before they contaminate tensor operations.

Interactive Debugging with breakpoint()

Instead of sprinkling print statements, use Python's built-in breakpoint() to inspect variables interactively. In phases/00-setup-and-tooling/12-debugging-and-profiling/code/debug_tools.py, the curriculum demonstrates conditional breakpoints that trigger only when anomalies occur:

def training_step(model, batch, criterion, optimizer):
    inputs, labels = batch
    outputs = model(inputs)
    loss = criterion(outputs, labels)

    if loss.item() > 100 or torch.isnan(loss):
        breakpoint()  # Drops into pdb or VS Code debugger

    loss.backward()
    optimizer.step()

When execution pauses, inspect tensor properties with commands like p outputs.shape or p torch.isnan(outputs).sum().

Structured Logging for Reproducibility

Replace print statements with Python's logging module to preserve timestamps and severity levels across distributed runs. The repository recommends this configuration:

import logging

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[logging.FileHandler("training.log"), logging.StreamHandler()],
)
logger = logging.getLogger(__name__)

logger.info("Starting training: lr=%.4f, batch=%d", lr, batch_size)
logger.warning("Loss spike: %.4f at step %d", loss.item(), step)
logger.error("NaN loss at step %d – aborting", step)

Performance Timing and Profiling

Identify bottlenecks using the Timer context manager and system profilers. The Timer class from the lesson wraps code sections to measure wall-clock time:

import time

class Timer:
    def __init__(self, name=""):
        self.name = name
    def __enter__(self):
        self.start = time.perf_counter()
        return self
    def __exit__(self, *exc):
        elapsed = time.perf_counter() - self.start
        print(f"[{self.name}] {elapsed:.4f}s")

with Timer("data loading"):
    batch = next(dataloader_iter)
with Timer("forward pass"):
    outputs = model(batch)

For function-level analysis, use cProfile from the shell:

python -m cProfile -s cumtime phases/00-setup-and-tooling/12-debugging-and-profiling/code/debug_tools.py

And for line-by-line granularity, apply the line_profiler decorator:

from line_profiler import profile

@profile
def train_step(model, data, target):
    out = model(data)
    loss = F.cross_entropy(out, target)
    loss.backward()
    return loss

# Run: kernprof -l -v debug_tools.py

Memory Profiling for CPU and GPU

Track allocations to prevent out-of-memory (OOM) errors. For CPU debugging, use tracemalloc:

import tracemalloc
tracemalloc.start()

# ... training loop ...

snapshot = tracemalloc.take_snapshot()
for stat in snapshot.statistics("lineno")[:10]:
    print(stat)

Or use the memory_profiler decorator:

from memory_profiler import profile

@profile
def load_data():
    raw = read_csv("data.csv")
    processed = preprocess(raw)
    return processed

# Run: python -m memory_profiler debug_tools.py

For GPU memory in PyTorch, print detailed summaries:

if torch.cuda.is_available():
    print(torch.cuda.memory_summary())
    print(f"Allocated: {torch.cuda.memory_allocated()/1e9:.2f} GB")
    print(f"Cached: {torch.cuda.memory_reserved()/1e9:.2f} GB")

When OOM occurs, the curriculum recommends this mitigation hierarchy:

1. Reduce batch size.
2. Call torch.cuda.empty_cache().
3. Delete large intermediates with del tensor.
4. Enable mixed precision with torch.cuda.amp.
5. Apply gradient checkpointing.

Layer 2: Tensor Operation Validation

After confirming Python-level correctness, validate tensor shapes, devices, and numerical stability before errors propagate downstream.

The debug_print Utility

Insert the debug_print helper from phases/00-setup-and-tooling/12-debugging-and-profiling/code/debug_tools.py after suspicious operations to verify tensor metadata:

def debug_print(name, tensor):
    print(
        f"{name}: shape={tensor.shape}, dtype={tensor.dtype}, "
        f"device={tensor.device}, "
        f"min={tensor.min().item():.4f}, max={tensor.max().item():.4f}, "
        f"mean={tensor.mean().item():.4f}, "
        f"has_nan={tensor.isnan().any().item()}"
    )

# Usage: debug_print("logits", logits)

Shape-Hook Inspection

Register forward hooks to trace tensor dimensions through every layer. This technique, documented in the debug-checklist at phases/03-deep-learning-core/13-debugging-neural-networks/outputs/skill-debug-checklist.md, prints a complete shape graph:

def check_shapes(model, sample_input):
    print(f"Input: {sample_input.shape}")
    hooks = []
    def make_hook(name):
        def hook(module, inp, out):
            in_shape = inp[0].shape if isinstance(inp, tuple) else inp.shape
            out_shape = out.shape if hasattr(out, "shape") else type(out)
            print(f"  {name}: {in_shape} → {out_shape}")
        return hook
    for name, module in model.named_modules():
        hooks.append(module.register_forward_hook(make_hook(name)))
    with torch.no_grad():
        model(sample_input)
    for h in hooks:
        h.remove()

Running this once reveals mismatches where a layer expects [B, C, H, W] but receives [B, F].

NaN and Inf Detection

Numerical instability kills training silently. Implement guards that halt execution immediately after backward passes:

def detect_nan(model, loss, step):
    if torch.isnan(loss):
        print(f"NaN loss at step {step}")
        for n, p in model.named_parameters():
            if p.grad is not None:
                if torch.isnan(p.grad).any():
                    print(f"  NaN grad in {n}")
                if torch.isinf(p.grad).any():
                    print(f"  Inf grad in {n}")
        return True
    return False

Data Leakage Validation

Before training, verify that train and test sets contain no overlapping identifiers:

def check_data_leakage(train_set, test_set, id_column="id"):
    train_ids = set(train_set[id_column].tolist())
    test_ids  = set(test_set[id_column].tolist())
    overlap = train_ids & test_ids
    if overlap:
        print(f"DATA LEAKAGE: {len(overlap)} IDs appear in both splits")
        return True
    return False

Layer 3: Training Dynamics and Production Observability

Monitor optimization behavior and implement enterprise-grade debugging for deployed systems.

TensorBoard Visualization

Track loss curves, learning rates, and weight distributions using the implementation shown in phases/00-setup-and-tooling/12-debugging-and-profiling/docs/en.md:

from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter("runs/exp1")

for step in range(num_steps):
    loss = train_step(...)
    writer.add_scalar("loss/train", loss.item(), step)
    writer.add_scalar("lr", optimizer.param_groups[0]["lr"], step)

    if step % 100 == 0:
        for name, param in model.named_parameters():
            writer.add_histogram(f"weights/{name}", param, step)
            if param.grad is not None:
                writer.add_histogram(f"grads/{name}", param.grad, step)

writer.close()

# Launch: tensorboard --logdir=runs

Visual cues indicate specific problems: flat loss suggests dead gradients or low learning rate, oscillating loss indicates excessive learning rate, and diverging validation loss signals overfitting.

VS Code Debugger Integration

Configure graphical debugging by creating .vscode/launch.json as referenced in the lesson:

{
    "version": "0.2.0",
    "configurations": [
        {
            "name": "Debug Training",
            "type": "debugpy",
            "request": "launch",
            "program": "${file}",
            "console": "integratedTerminal",
            "justMyCode": false
        }
    ]
}

The Variables pane displays tensor.shape, tensor.device, and tensor.dtype instantly without manual printing.

Time-Travel Debugging for Multi-Agent Systems

For production AI agents built with LangGraph, utilize time-travel debugging covered in phases/11-llm-engineering/16-langgraph-state-machines/docs/en.md. This technique uses graph.get_state_history(thread_id) to retrieve checkpoint logs and graph.invoke(..., checkpoint_id=...) to branch execution from any prior state, enabling precise post-mortem analysis of multi-step failures without re-running the entire pipeline.

Summary

The AI Engineering from Scratch curriculum provides a systematic debugging hierarchy:

• Layer 1 (Python): Use breakpoint(), structured logging, cProfile, and memory profilers to catch environment and logic errors.
• Layer 2 (Tensors): Validate shapes with forward hooks, detect NaN/Inf values, and check for data leakage before training begins.
• Layer 3 (Dynamics): Monitor training with TensorBoard histograms, leverage VS Code's graphical debugger, and apply time-travel debugging for complex agent pipelines.

These techniques, implemented in phases/00-setup-and-tooling/12-debugging-and-profiling/code/debug_tools.py and related modules, scale from research notebooks to production deployments.

Frequently Asked Questions

What is the difference between using print statements and logging for AI debugging?

Print statements buffer output and lack timestamps, making them unsuitable for long-running distributed training jobs. The logging module writes to both console and persistent files with severity levels, ensuring you can analyze failure sequences after crashes or node failures.

How do you detect shape mismatches in PyTorch neural networks?

Register forward hooks using module.register_forward_hook() to print input and output shapes for every layer. Run a single dummy batch through the model to generate a complete shape graph, immediately revealing dimensions that conflict between consecutive layers.

What causes NaN losses in neural networks and how do you fix them?

NaN losses typically originate from exploding gradients, invalid data inputs, or numerical instability in loss functions. Implement a detect_nan function that checks torch.isnan() on loss and gradients immediately after backward(), then apply gradient clipping, reduce learning rates, or add epsilon values to denominators as appropriate.

How does time-travel debugging work for AI agents?

Time-travel debugging, as implemented in LangGraph, persists state checkpoints throughout agent execution. Developers call graph.get_state_history() to list previous states, then invoke graph.invoke() with a specific checkpoint_id to restart execution from that exact point, enabling precise inspection of decision branches without re-running the entire pipeline.

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