# Debugging and Profiling Techniques for Developing and Optimizing AI Systems: A Three-Layer Framework > Learn AI debugging and profiling with a three-layer framework. Master Python tools, tensor validation, and training visualization for robust AI systems. Optimize your AI development now. - Repository: [Rohit Ghumare/ai-engineering-from-scratch](https://github.com/rohitg00/ai-engineering-from-scratch) - Tags: how-to-guide - Published: 2026-06-18 --- **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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/00-setup-and-tooling/12-debugging-and-profiling/code/debug_tools.py), the curriculum demonstrates conditional breakpoints that trigger only when anomalies occur: ```python 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: ```python 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: ```python 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: ```bash 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: ```python 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`: ```python 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: ```python 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: ```python 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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/00-setup-and-tooling/12-debugging-and-profiling/code/debug_tools.py) after suspicious operations to verify tensor metadata: ```python 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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/03-deep-learning-core/13-debugging-neural-networks/outputs/skill-debug-checklist.md), prints a complete shape graph: ```python 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: ```python 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: ```python 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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/00-setup-and-tooling/12-debugging-and-profiling/docs/en.md): ```python 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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/.vscode/launch.json) as referenced in the lesson: ```json { "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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/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.