Profiling en PyTorch: guía inicial de torch.profiler | Keryc
La regla es sencilla: lo que no puedes perfilar no puedes optimizar. ¿Quieres exprimir más tokens por segundo de un LLM, quitar milisegundos a la inferencia o entender por qué tu loop de entrenamiento va más lento de lo que promete la hoja de especificaciones? En algún punto tienes que abrir un rastro del profiler y leerlo.
Pero abrir ese trazo suele sentirse como entrar a una pared de rectángulos de colores con nombres intimidantes. Esta guía técnica, primera de una serie, baja esa rampa de entrada: partimos de lo más simple, una multiplicación de matrices seguida de una suma de bias, y aprendemos a leer lo que torch.profiler nos devuelve.
Preparando el experimento y qué exporta el profiler
La función que usamos es mínima y pedagógica:
def fn(x, w, b): return torch.add(torch.matmul(x, w), b)
Anotar la región con torch.profiler.record_function('matmul_add') es opcional pero útil: facilita saltar a esa etiqueta en las trazas. El flujo típico de profiling en el script es:
definir la región a medir (aquí fn) y envolverla en record_function.
ejecutar bajo torch.profiler.profile(...) con actividades CPU y CUDA.
llamar prof.key_averages().table(...) para la tabla estadística.
exportar la traza con prof.export_chrome_trace(trace_path) para verla en Perfetto.
El profiler produce dos artefactos complementarios:
La tabla del profiler: resumen estadístico que responde "qué se lleva la mayor parte del tiempo". Ideal para encontrar hotspots.
La traza temporal (.json): vista tiempo-espacio de CPU vs GPU, con los kernels, huecos, y dependencias.
Ejemplo de comandos usados en la guía:
uv run 01_matmul_add.py --size 64
y para runs más grandes o con warmup/compile:
uv run 01_matmul_add.py --size 4096 --warmup --compile
Leer la tabla: Self vs Total y el síntoma de overhead
Al abrir la tabla verás columnas como Self CPU time total y Self CUDA time total. Atención a esta señal frecuente:
Si Self CPU time total está en milisegundos y Self CUDA time total en microsegundos, el patrón típico es overhead-bound: la CPU hace más trabajo preparatorio que lo que el GPU tarda en calcular.
En el ejemplo con tamaño 64 observamos algo así como:
Self CPU time total: 2.314 ms
Self CUDA time total: 23.104 us
Eso significa que el GPU está casi siempre inactivo: el matmul pequeño se ejecuta rápido en GPU, pero la mayor parte del tiempo se gasta en lanzar kernels, asignar workspace o ejecutar heurísticas de cuBLAS. Solución simple: aumentar la carga por llamada (matrices más grandes o batching) o fusionar operaciones.
Con --size 4096 el balance cambia:
Self CPU time total: 4.908 ms
Self CUDA time total: 4.495 ms
Ambos ahora en ms, estamos en una región compute-bound donde el GPU domina: el objetivo usual cuando escalar a modelos reales.
Leer la traza: CPU lane, GPU lane y las ventanas muertas
La traza muestra dos lanes principales: CPU y GPU. Los elementos clave que debes buscar y qué significan:
Un espacio entre la entrada de record_function('...') y el primer aten::* es una "ventana muerta". Puede ser allocation de workspace, heurísticas de cuBLAS o lazy loading. ¿Qué hacer? Warmup.
El scheduler de profiling controla qué pasos se registran. Ejemplo de schedule común:
wait salta inicializaciones; warmup ejecuta sin registrar; active es lo que ves en la traza.
Si ves un offset de varios ms entre el momento en que la CPU envía trabajo y cuando el GPU empieza, puede aparecer una entrada llamada "Activity Buffer Request" en la lane de GPU: el propio profiler refiriéndose a su buffer de eventos o a alguna solicitud de memoria.
Si un paso es mucho más ancho que los demás (primer ProfileStep, por ejemplo) es cold-start: agrega warmup antes de medir.
Señales en la traza y llamadas CUDA que importan
Busca estas llamadas y ten en mente su significado:
cudaOccupancyMaxActiveBlocksPerMultiprocessor antes de un cudaLaunchKernel: es una consulta de planificación. Suele aparecer para kernels pesados como GEMM o convoluciones. Indica que cuBLAS está eligiendo heurísticamente una variante.
cudaLaunchKernel sin la ocupancy query: típicamente kernels elementwise o reducciones, de huella fija y liviana.
cudaDeviceSynchronize al final de la ventana activa: el profiler fuerza un sincronizado para volcar eventos. Si esa sincronización cubre microsegundos de trabajo GPU pero milisegundos de sincronización, es otro signo de overhead-bound.
cudaMemcpyAsync que no escribiste: muchas veces es un Device-to-Device copy que prepara la memoria antes de un GEMM con epílogo (por ejemplo bias add). Inductor o cuBLAS pueden hacer ese memcpy implícito.
Variabilidad en tiempos de kernel
Incluso con el mismo kernel y los mismos datos, verás variabilidad entre invocaciones. Causas comunes:
Boost de reloj en GPUs después de periodos de inactividad.
Thermal throttling o gestión de energía.
Tareas del driver o housekeeping del sistema.
La moraleja: no te fíes solo de la media. La traza te muestra la distribución real.
Qué cambia con torch.compile y qué no
torch.compile (por ejemplo con el backend predeterminado TorchInductor) captura fragmentos de ejecución y genera código: pipeline general resumido:
TorchDynamo captura ejecución Python en un grafo FX.
AOTAutograd prepara forward/backward cuando hay gradientes.
Inductor baja el grafo a código GPU/CPU optimizado.
En la traza verás filas como "Torch-Compiled Region" y llamadas a TorchDynamo Cache Lookup en cada paso: Dynamo revisa si los shapes y metadatos casan con la cache. Ese chequeo se paga en cada llamada.
Un detalle importante observado en el ejemplo:
Inductor puede fusionar al nivel del dispatcher: tu torch.add(torch.matmul(x, w), b) puede convertirse en aten::addmm(b, x, w).
Sin embargo, eso no garantiza un único kernel GPU nuevo. En el caso 4096x4096 bf16, Inductor escogió usar cuBLAS y arregló la bias add como un epílogo de GEMM, acompañada de un Memcpy DtoD para preparar el buffer de salida.
Resultado práctico: compile puede aumentar el costo por llamada en la CPU (Dynamo + AOTAutograd + Inductor) para operaciones muy pequeñas. Para modelos con docenas o cientos de ops, ese overhead se amortiza. Para una operación aislada, es una tasa que te come tiempo.
Diagnóstico rápido: señales y cómo interpretarlas
Usa esto como chequeo cuando mires una traza:
Self CPU time total mucho mayor que Self CUDA time total (CPU ms vs GPU us): overhead-bound. Haz el trabajo más grande o fusiona.
Self CPU ≈ Self CUDA en ms: compute-bound. GPU es el cuello de botella.
Un evento domina cuda_time_total: ahí empieza la optimización.
Un evento con muchas llamadas (# of Calls grande): potencial para batching o fusión.
cudaMemcpyAsync inesperado antes de GEMM: copia DtoD para epílogo de addmm.
Actividad Buffer Request en GPU lane: el profiler o la runtime está pidiendo/llenando buffers y puede explicar offsets entre lanes.
Variación en tiempos del mismo kernel entre steps: lee la traza completa, no solo la media.
Recomendaciones prácticas rápidas
Empieza perfilando con casos simples y viste cómo cambian las métricas al aumentar tamaño y batching.
Usa warmup (tanto fuera como dentro del profiler) para eliminar costos de cold-start.
Si el CPU domina, intenta fusionar operadores o mover trabajo al GPU en bloques más grandes.
Revisa la traza completa en Perfetto para encontrar gaps y memcpy implícitos; la tabla por sí sola no lo contará todo.
Cuando uses torch.compile, inspecciona el árbol CPU: Dynamo cache hits, AOT wrapper y el Call CompiledFxGraph te dirán cuánto se está pagando en CPU.
Al terminar este primer ejercicio con un matmul + bias tienes ya varios lentes para leer trazas más complejas: distinguir overhead de cómputo, identificar planificación cuBLAS, entender copias implícitas y ver qué hace realmente torch.compile. En los siguientes posts subiremos la complejidad: primero nn.Linear y un MLP, y luego modelos grandes con transformers. La idea es construir la intuición paso a paso para que cuando abras una traza no sea terror sino curiosidad de detective.
