Phase 00 - Lesson 12

Depuración y Profiling

Los peores bugs de IA no provocan fallos. Entrenan en silencio sobre basura y reportan una hermosa curva de loss.

Tipo: Construir Lenguaje: Python Requisitos previos: Lección 1 (Entorno de Dev), familiaridad básica con PyTorch Tiempo: ~60 minutos

Objetivos de Aprendizaje

  • Usar breakpoint() condicional y debug_print para inspeccionar las formas de los tensores, los dtypes y los valores NaN durante el entrenamiento
  • Hacer profiling de los bucles de entrenamiento con cProfile, line_profiler y tracemalloc para encontrar cuellos de botella
  • Detectar bugs comunes de IA: incompatibilidad de formas, loss NaN, fuga de datos y tensores en el dispositivo equivocado
  • Configurar TensorBoard para visualizar curvas de loss, histogramas de pesos y distribuciones de gradientes

El Problema

El código de IA falla de forma diferente al código normal. Una app web se cae con un stack trace. Un bucle de entrenamiento mal configurado se ejecuta durante 8 horas, quema USD 200 en tiempo de GPU y produce un modelo que predice la media de cada entrada. El código nunca dio error. El bug era un tensor en el dispositivo equivocado, un .detach() olvidado o etiquetas filtrándose hacia las features.

Necesitas herramientas de depuración que atrapen estas fallas silenciosas antes de que desperdicien tu tiempo y tu cómputo.

El Concepto

La depuración de IA opera en tres niveles:

graph TD
    L3["3. Training Dynamics<br/>Loss curves, gradient norms, activations"] --> L2
    L2["2. Tensor Operations<br/>Shapes, dtypes, devices, NaN/Inf values"] --> L1
    L1["1. Standard Python<br/>Breakpoints, logging, profiling, memory"]

La mayoría de la gente salta directo al nivel 3 (mirar fijamente TensorBoard). Pero el 80% de los bugs de IA viven en los niveles 1 y 2.

Constrúyelo

Parte 1: Depuración con Print (Si, Funciona)

La depuración con print se desprecia. No debería. Para código de tensores, un print dirigido supera recorrer un depurador, porque necesitas ver formas, dtypes y rangos de valores todo a la vez.

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()}")

Llama a esto después de cada operación sospechosa. Cuando encuentres el bug, quita los prints. Sencillo.

Parte 2: Depurador de Python (pdb y breakpoint)

El depurador integrado esta subestimado para el trabajo con IA. Coloca breakpoint() en tu bucle de entrenamiento e inspecciona tensores de forma interactiva.

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()

    loss.backward()
    optimizer.step()

Cuando el depurador te deja dentro, comandos útiles:

  • p outputs.shape para verificar formas
  • p loss.item() para ver el valor del loss
  • p torch.isnan(outputs).sum() para contar NaNs
  • p model.fc1.weight.grad para verificar gradientes
  • c para continuar, q para salir

Esto es depuración condicional. Solo te detienes cuando algo se ve mal. Para una ejecución de entrenamiento de 10.000 pasos, eso importa.

Parte 3: Logging en Python

Reemplaza las instrucciones de print por logging cuando tu depuración vaya más alla de una verificación rápida.

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_size=%d", lr, batch_size)
logger.warning("Loss spike detected: %.4f at step %d", loss.item(), step)
logger.error("NaN loss at step %d, stopping", step)

El logging te da marcas de tiempo, niveles de severidad y salida a archivo. Cuando una ejecución de entrenamiento falla a las 3 de la madrugada, quieres un archivo de registro, no una salida de terminal que se desplazó fuera de la pantalla.

Parte 4: Medir el Tiempo de Secciones de Código

Saber a dónde se va el tiempo es el primer paso hacia la optimización.

import time

class Timer:
    def __init__(self, name=""):
        self.name = name

    def __enter__(self):
        self.start = time.perf_counter()
        return self

    def __exit__(self, *args):
        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)

with Timer("backward pass"):
    loss.backward()

Hallazgo común: la carga de datos toma el 60% del tiempo de entrenamiento. La solución es num_workers > 0 en tu DataLoader, no una GPU más rápida.

Parte 5: cProfile y line_profiler

Cuando necesitas más que temporizadores manuales:

python -m cProfile -s cumtime train.py

Esto muestra cada llamada a función ordenada por tiempo acumulado. Para profiling línea por línea:

pip install line_profiler
@profile
def train_step(model, data, target):
    output = model(data)
    loss = F.cross_entropy(output, target)
    loss.backward()
    return loss

# Run with: kernprof -l -v train.py

Parte 6: Profiling de Memoria

Memoria de CPU con tracemalloc

import tracemalloc

tracemalloc.start()

# your code here
model = build_model()
data = load_dataset()

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

Memoria de CPU con memory_profiler

pip install memory_profiler
from memory_profiler import profile

@profile
def load_data():
    raw = read_csv("data.csv")       # watch memory jump here
    processed = preprocess(raw)       # and here
    return processed

Ejecuta con python -m memory_profiler your_script.py para ver el uso de memoria línea por línea.

Memoria de GPU con PyTorch

import torch

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")

Cuando llegas a OOM (Out of Memory):

  1. Reduce el tamaño del batch (lo primero que hay que probar, siempre)
  2. Usa torch.cuda.empty_cache() para liberar la memoria en cache
  3. Usa del tensor seguido de torch.cuda.empty_cache() para grandes intermedios
  4. Usa precision mixta (torch.cuda.amp) para reducir el uso de memoria a la mitad
  5. Usa gradient checkpointing para modelos muy profundos

Parte 7: Bugs Comunes de IA y Como Atraparlos

Incompatibilidad de Forma (Shape Mismatch)

El bug más frecuente. Un tensor tiene forma [batch, features] cuando el modelo espera [batch, channels, height, width].

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()

Ejecuta esto una vez con un batch de ejemplo. Mapea cada transformación de forma en tu modelo.

Loss NaN

Un loss NaN significa que algo explotó. Causas comunes:

  • Tasa de aprendizaje demasiado alta
  • División por cero en un loss personalizado
  • Logaritmo de cero o de un número negativo
  • Gradientes que explotan en RNNs
def detect_nan(model, loss, step):
    if torch.isnan(loss):
        print(f"NaN loss at step {step}")
        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}")
                if torch.isinf(param.grad).any():
                    print(f"  Inf gradient in {name}")
        return True
    return False

Fuga de Datos (Data Leakage)

Tu modelo alcanza el 99% de exactitud en el conjunto de prueba. Suena estupendo. Es un bug.

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)} samples in both train and test")
        return True
    return False

Verifica también la fuga temporal: usar datos futuros para predecir el pasado. Ordena por marca de tiempo antes de dividir.

Dispositivo Equivocado

Los tensores en dispositivos diferentes (CPU vs GPU) provocan errores en tiempo de ejecución. Pero a veces un tensor se queda en silencio en la CPU mientras todo lo demás esta en la GPU, y el entrenamiento simplemente se ejecuta lento.

def check_devices(model, *tensors):
    model_device = next(model.parameters()).device
    print(f"Model device: {model_device}")
    for i, t in enumerate(tensors):
        if t.device != model_device:
            print(f"  WARNING: tensor {i} on {t.device}, model on {model_device}")

Parte 8: Fundamentos de TensorBoard

TensorBoard te muestra lo que está ocurriendo dentro del entrenamiento a lo largo del tiempo.

pip install tensorboard
from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter("runs/experiment_1")

for step in range(num_steps):
    loss = train_step(model, batch)

    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()

Lánzalo:

tensorboard --logdir=runs

Que buscar:

  • El loss no disminuye: Tasa de aprendizaje demasiado baja, o problema en la arquitectura del modelo
  • El loss oscila descontroladamente: Tasa de aprendizaje demasiado alta
  • El loss llega a NaN: Inestabilidad numérica (ver la sección de NaN más arriba)
  • El loss de entrenamiento disminuye, el de validación aumenta: Overfitting
  • Los histogramas de pesos colapsan a cero: Gradientes que se desvanecen (vanishing gradients)
  • Los histogramas de gradientes explotan: Se necesita recorte (clipping) de gradientes

Parte 9: Depurador de VS Code

Para depuración interactiva, configura VS Code con un launch.json:

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

Coloca breakpoints haciendo clic en el margen (gutter). Usa el panel Variables para inspeccionar las propiedades de los tensores. La Debug Console permite ejecutar expresiones Python arbitrarias durante la ejecución.

Útil para recorrer pipelines de preprocesamiento de datos donde quieres ver cada transformación.

Usalo

Aquí esta el flujo de depuración que atrapa la mayoría de los bugs de IA:

  1. Antes del entrenamiento: Ejecuta check_shapes con un batch de ejemplo. Verifica que las dimensiones de entrada y salida coincidan con lo esperado.
  2. Primeros 10 pasos: Usa debug_print en el loss, las salidas y los gradientes. Confirma que nada sea NaN y que los valores esten en rangos razonables.
  3. Durante el entrenamiento: Registra el loss, la tasa de aprendizaje y las normas de los gradientes. Usa TensorBoard para la visualizacion.
  4. Cuando algo se rompe: Coloca breakpoint() en el punto de la falla. Inspecciona los tensores de forma interactiva.
  5. Para el rendimiento: Mide el tiempo de tu carga de datos vs el forward vs el backward pass. Haz profiling de memoria si estas cerca de OOM.

Entrégalo

Ejecuta el script del toolkit de depuración:

python phases/00-setup-and-tooling/12-debugging-and-profiling/code/debug_tools.py

Consulta outputs/prompt-debug-ai-code.md para un prompt que ayuda a diagnosticar bugs específicos de IA.

Ejercicios

  1. Ejecuta debug_tools.py y lee la salida de cada sección. Modifica el modelo dummy para introducir un NaN (pista: divide por cero en el forward pass) y observa cómo el detector lo atrapa.
  2. Haz profiling de un bucle de entrenamiento con cProfile e identifica la función más lenta.
  3. Usa tracemalloc para encontrar que línea de tu pipeline de carga de datos asigna más memoria.
  4. Configura TensorBoard para una ejecución de entrenamiento simple e identifica si el modelo está en overfitting.
  5. Usa breakpoint() dentro de un bucle de entrenamiento. Práctica inspeccionar formas de tensores, dispositivos y valores de gradientes desde el prompt del depurador.
0 lifetime access. Curriculum based on AI Engineering from Scratch by Rohit Ghumare (MIT, used under attribution).