torch.profiler explicado: cómo leer trazas en PyTorch desde cero

Lo que no puedes medir, no puedes optimizar.

Ya sea que intentes exprimir más tokens por segundo a un modelo de lenguaje (LLM), recortar milisegundos en inferencia o simplemente entender por qué tu loop de entrenamiento corre más lento de lo que prometía la hoja técnica, el camino tarde o temprano pasa por profiling.

El problema es que la curva de entrada es empinada. Las trazas son paredes densas de rectángulos de colores. Los eventos tienen nombres intimidantes. La mayoría de los tutoriales asume que ya sabes leerlos. Por eso, aun sabiendo que deberíamos estar profileando, abrir una traza se siente como una tarea que se deja para después. Este post, y la serie que abre, intenta bajar esa barrera.

Es la primera entrega de Profiling en PyTorch, una serie en la que se construye poco a poco la habilidad de leer trazas y usarlas para guiar la optimización. El plan:

• Parte 1 (este post): partir con la operación más simple posible, una multiplicación de matrices seguida de una suma con bias, y aprender a leer lo que devuelve el profiler.
• Parte 2: escalar a nn.Linear y a un MLP pequeño, usar las trazas para motivar optimizaciones y mirar los kernels por debajo.
• Parte 3: integrar todo en modelos de lenguaje con transformers.

El recorrido es desde el punto de vista de alguien que recién parte. Sin prerrequisitos más allá de PyTorch básico. La estructura es deliberadamente guiada por preguntas: se abre una traza, se pregunta "un momento, ¿por qué pasa eso?" y se persigue la respuesta hasta que algo encaja. Al final deberías saber:

• cómo configurar torch.profiler y qué entrega exactamente,
• cómo leer la tabla del profiler y la traza (canal CPU, canal GPU y los huecos sospechosos entre medio),
• la cadena de eventos desde una llamada en Python hasta un kernel CUDA,
• qué cambia (y, más interesante, qué no cambia) cuando le pones torch.compile encima.

Antes de empezar, dos definiciones que ayudan:

• Un kernel de GPU es un programa que corre en paralelo en muchos hilos de la GPU.
• La CPU programa y lanza esos kernels.

Normalmente no escribes kernels GPU a mano; cuando usas una operación de PyTorch, se traduce automáticamente a uno o más kernels que hacen el trabajo en la GPU.

¿Qué operación se usa de ejemplo?

El script completo del post está en 01_matmul_add.py. Conviene abrirlo en una pestaña aparte. Los autores corren los scripts en una NVIDIA A100-SXM4-80GB.

YouTube

Como bien apuntó la Dra. Sara Hooker, así como nosotros somos en su mayoría agua, las redes neuronales profundas son en su mayoría multiplicaciones de matrices. Por fundamentales que sean, sería una pena partir esta travesía con otra cosa.

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

La suma matricial junto a la multiplicación imita cómo interactúan pesos y bias dentro de una neurona. Esa suma además abre el camino para hablar de compilación más adelante.

Para profilear se usa el módulo torch.profiler. Los pasos son:

• Tener el código a profilear listo (aquí def fn, que envuelve el matmul y la suma).
• Anotar el algoritmo. Aunque es opcional, se recomienda. record_function anota la función como matmul_add, fácil de ubicar en las trazas.

def step():
    with torch.profiler.record_function("matmul_add"):
        return fn(x, w, b)

• Envolver el código con el context manager torch.profiler.profile:

with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],
) as prof:
    for _ in range(5):
        step()
        prof.step()

• Exportar el perfil:

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

El profiler exporta dos artefactos distintos:

• La tabla del profiler: el resumen estadístico del algoritmo. Responde qué está tomando más tiempo. Útil para encontrar hotspots: eventos que tardan más, que son cuellos de botella o que se disparan muchas veces.
• La traza del profiler: la vista temporal de la ejecución. Responde cuándo y por qué ocurrió una operación, mostrando las actividades de la CPU y de la GPU. Útil para investigar qué kernels se lanzaron, demoras en lanzarlos y solapamientos entre CPU y GPU.

¿Qué dice la tabla en un matmul de 64x64?

Ejecutado el script en una máquina con GPU:

uv run 01_matmul_add.py --size 64

Aparece una carpeta traces/01_matmul_add con los dos artefactos 64_bf16_cold_eager.json y 64_bf16_cold_eager.txt. El .txt es la tabla del profiler, con una primera columna que lista los eventos disparados dentro del scope. Las demás columnas indican el tiempo del evento en CPU, GPU u otros dispositivos declarados en activities.

Vale la pena mirar la columna # of Calls para saber cuántas veces se disparó cada evento, y distinguir entre Self CPU/CUDA (tiempo dentro del evento, sin hijos) y CPU/CUDA total (evento más hijos). En matmul_add, el CPU total es la suma del tiempo propio más los hijos que se gatillan.

Las dos últimas líneas dicen:

Self CPU time total: 2.314ms
Self CUDA time total: 23.104us

El tiempo en CPU está en milisegundos y el de GPU en microsegundos. Para poner perspectiva: el tiempo en GPU (el kernel ampere_bf16_s16816gemm...) es menos del 1% del tiempo gastado en CPU (la operación matmul_add). La GPU está ociosa casi todo el tiempo, una bandera roja inmediata. Pasa porque la GPU resuelve un matmul chico muy rápido, así que el código se va casi entero en preparar los kernels, lanzarlos, mover datos y juntar resultados. A esto se le llama régimen overhead-bound.

La forma más simple de salir de ese régimen es usar matmuls más grandes:

uv run 01_matmul_add.py --size 4096

Las dos últimas líneas pasan a:

Self CPU time total: 4.908ms
Self CUDA time total: 4.495ms

Ahora ambos están en milisegundos: subiendo el tamaño se materializa más tiempo de GPU. El kernel ampere_bf16_s16816gemm_.. pasa a dominar el CUDA time en lugar de la operación matmul_add en CPU. Cambiamos de régimen overhead bound a régimen compute bound.

¿Cómo se ve la traza temporal y por qué hay huecos?

Para ver la cadena de dispatch hay que abrir los artefactos .json. Se pueden subir a Perfetto UI o generar links directos con uvx trace-util traces -b traces.

En la traza, el ancho de cada barra indica la duración del evento, el anidamiento vertical es la jerarquía de llamadas, el canal CPU muestra los eventos en CPU y el canal GPU muestra las ejecuciones reales de kernels. Los espacios vacíos son tiempos de espera o idle.

El script corre con configuración por defecto: tamaño 64, dtype bf16, sin compilación y sin warmup. Sobre la traza:

schedule = torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1)

wait se salta inicializaciones ruidosas (ProfilerStep#0), warmup corre sin grabar (ProfilerStep#1) y active es lo que aparece efectivamente en la traza.

¿Por qué ProfilerStep#2 se demora tanto?

ProfilerStep#2 toma más tiempo que los demás, y lo mismo ocurre con la anotación matmul_add. La pista no está en la anotación sino dentro: hay una "ventana muerta" de unos 228 µs entre entrar a record_function("matmul_add") y el dispatch real de aten::matmul. Las causas son varias: workspace allocations, heurísticas de cuBLAS (la librería propietaria de NVIDIA para álgebra lineal acelerada en GPU) o lazy loading de módulos. La salida estándar es correr más pasos de warmup antes de profilear.

En profiling, warmup es correr los eventos un par de veces antes de medirlos: el trabajo previo de la GPU (los punteros antes mencionados, allocations, etc.) son esfuerzos one-shot que no queremos profilear. El script de Hugging Face usa dos etapas de warmup: una manual antes de entrar al profiler y otra dentro, controlada por el argumento warmup del schedule.

uv run 01_matmul_add.py --warmup

La serie continuará con la Parte 2, que escala a nn.Linear, un MLP pequeño y conexión con torch.compile para mirar los kernels que efectivamente terminan corriendo en la GPU.