Profiling en PyTorch: guía inicial con torch.profiler

Por qué perfilar importa

Lo que no se puede perfilar no se puede optimizar. Ya sea que estén tratando de obtener más tokens por segundo de un modelo de lenguaje, reducir milisegundos de latencia en inferencia o entender por qué el ciclo de entrenamiento va más lento de lo esperado, el camino pasa por el profiling. Sin embargo, leer trazas puede intimidar: tablas densas, eventos con nombres enrevesados y vistas temporales con muchas capas. Esta guía busca bajar esa barrera de entrada.

Este es el primer artículo de una serie sobre profiling en PyTorch. Aquí nos concentramos en lo más simple: una multiplicación de matrices seguida por una suma de bias, y aprendemos a interpretar lo que nos devuelve torch.profiler. En artículos posteriores escalaremos a nn.Linear, una pequeña MLP y, finalmente, modelos Transformer.

Conceptos clave antes de empezar

Un par de definiciones que ayudan a entender las trazas:

• GPU kernel: es un programa que se ejecuta en paralelo en muchas hebras del GPU.
• Lanzamiento de kernels: la CPU es la que agenda y lanza esos kernels en la GPU.

Cuando ejecutan operaciones de PyTorch en GPU, internamente se traduce en uno o más kernels que hacen el trabajo pesado. Entender la cadena completa —llamada Python → operaciones de alto nivel → kernels CUDA— es central para interpretar las trazas.

El ejemplo mínimo: matmul + add

La operación que usaremos es intencionalmente simple:

def fn(x, w, b): return torch.add(torch.matmul(x, w), b)

Esta composición reproduce la interacción básica entre pesos y bias en una neurona y es suficiente para observar cómo PyTorch lanza kernels y cómo se comportan CPU y GPU en el tiempo.

Cómo ejecutar el profiler (paso a paso)

Los pasos básicos para obtener una traza útil con torch.profiler son:

1. Tener el código listo: en este caso fn y una función step() que envuelva la llamada.
2. (Opcional pero recomendable) Anotar regiones con torch.profiler.record_function("matmul_add") para que la traza muestre etiquetas legibles.
3. Ejecutar el perfilador con el context manager torch.profiler.profile, solicitando actividades CPU y CUDA:

with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],
) as prof:
    for _ in range(5):  # varias iteraciones para warm-up
        step()
        prof.step()

4. Exportar los artefactos: la tabla de promedios y la traza en formato Chrome (JSON):

• prof.key_averages().table(sort_by="cuda_time_total", row_limit=15)
• prof.export_chrome_trace(trace_path)

En el ejemplo usado para esta guía, el autor ejecutó el script en una máquina con una NVIDIA A100-SXM4-80GB y guardó los resultados en traces/01_matmul_add con archivos como 64_bf16_cold_eager.json y 64_bf16_cold_eager.txt.

Nota para equipos en Latinoamérica: si no disponen de acceso a GPUs de última generación, pueden ejecutar los mismos pasos en GPUs más modestas o en instancias en la nube para aprender a leer las trazas; el patrón de eventos y el modo de interpretación es el mismo.

¿Qué entregan el profiler table y la traza?

El profiler produce dos artefactos complementarios:

• La tabla del profiler: un resumen estadístico que responde “¿qué toma más tiempo?”. Es útil para identificar hotspots —eventos que consumen la mayor porción de tiempo o que se repiten mucho.

• La traza (Chrome trace): una vista temporal que responde “¿cuándo y por qué ocurrió una operación?”. Muestra actividades en CPU y GPU, solapamientos, latencias entre el lanzamiento de un kernel y su ejecución, y ayuda a localizar demoras en la cadena de ejecución.

Al abrir el .txt con la tabla, verán filas por evento y columnas con tiempos en CPU y CUDA. Presten atención a:

• # of Calls: cuántas veces se llamó ese evento.
• Self CPU/CUDA vs CPU/CUDA Total: “Self” mide el tiempo dentro del evento excluyendo sus hijos; “Total” incluye el tiempo de los hijos. Esto es clave para distinguir si el tiempo se gasta en la operación en sí o en llamadas que ésta provoca.

Cómo leer la traza temporal

En la traza visual se distinguen típicamente:

• La “lane” de CPU, donde se ven llamadas de alto nivel y el lanzamiento de kernels.
• La “lane” de GPU, donde aparecen los kernels ejecutándose.
• Los huecos entre ellas: pueden indicar latencia de host, sincronizaciones implícitas o colas de ejecución.

Etiquetar la región con record_function("matmul_add") facilita localizar el bloque de interés en la traza y entender qué llamadas Python dieron lugar a qué kernels CUDA.

Buenas prácticas básicas

• Ejecuten varias iteraciones antes de recoger datos definitivos para evitar medir el estado “cold” del GPU.
• Ordenen la tabla por cuda_time_total o por la métrica que más les importe (CPU o memoria).
• Usen record_function para marcar bloques lógicos del código; esto acelera el diagnóstico.

Qué sigue en la serie

En las próximas entregas escalaremos el ejemplo hacia nn.Linear y una pequeña MLP, extraeremos optimizaciones basadas en las trazas y exploraremos los kernels subyacentes. Finalmente, cerraremos el ciclo aplicando lo aprendido a modelos de lenguaje.

Cierre y recomendación práctica

Comiencen probando el script de ejemplo y abran tanto la tabla como la traza. La primera lectura suele responder “qué está pasando” y la segunda “por qué y cuándo”. Para equipos en Latinoamérica que buscan optimizar costos en nube o exprimir instancias locales, dominar el profiling es una habilidad de alto retorno: permite decisiones informadas sobre batching, precisión numérica, uso de memoria y cambios en la arquitectura.

Si quieren seguir avanzando: identifiquen un hotspot en su modelo real, perfílenlo con la misma metodología y usen la traza para decidir la siguiente optimización (por ejemplo, fusionar operaciones, cambiar tipos numéricos o ajustar el data pipeline). En la siguiente parte de esta serie veremos cómo esas decisiones nacen directamente de lo que muestran las trazas.