Perfilado en PyTorch: de nn.Linear a un MLP fusionado

Introducción

En la segunda entrega de la serie sobre perfilado en PyTorch avanzamos un peldaño: pasamos de la operación manual matmul-add a la implementación estándar que usan la mayoría de los modelos, nn.Linear. Usamos trazas del perfilador para entender qué pasa en el CPU y en la GPU cuando ejecutamos una sola capa lineal y cuándo tiene sentido usar torch.compile.

Los scripts que acompañan el ejemplo son 02_linear.py, 03_simple_mlp.py y 03_kernels_mlp.py; la ejecución se realizó en una NVIDIA A100-SXM4-80GB. Si trabajan desde América Latina, recordar que plataformas como Hugging Face permiten probar estos scripts fácilmente con Dev Mode en Spaces o con Jobs, lo que facilita replicar las trazas localmente o en la nube.

Rápido recordatorio: kernels GPU y scheduling

Una GPU ejecuta kernels que corren en muchos hilos en paralelo. Es la CPU la que programa y lanza estos kernels; en las trazas del perfilador gran parte del “overhead” que se observa corresponde a ese trabajo de scheduling y a operaciones pequeñas en CPU que preparan la llamada al kernel. Diferenciar entre un régimen limitado por overhead (scheduler/CPU) y uno limitado por cómputo (GPU) es clave para optimizar.

De matmul-add a nn.Linear

nn.Linear envuelve exactamente la multiplicación de matrices y la suma de bias que vimos antes, pero con la diferencia de que almacena peso y bias como parámetros y expone un forward de uso habitual:

y = x @ w.T + b

Al ejecutar 02_linear.py con parámetros como —batch 1024 —in_dim 32 —out_dim 64 y analizar la traza del perfilador vemos el comportamiento real de PyTorch en CPU y GPU.

¿Qué hace la transposición (aten::t)?

En la traza aparece un aten::t (transpose) justo antes de la llamada a aten::addmm. Es importante entender que esa transpose no necesariamente copia datos ni lanza un kernel en GPU: aten::t reescribe metadatos del tensor (forma y stride) en CPU para representar la vista transpuesta. No implica un lanzamiento de kernel; se puede comprobar observando que no hay actividad CUDA asociada a esa fila en la traza.

Esto es relevante para el rendimiento: cambiar metadatos es barato comparado con mover datos en memoria. Sin embargo, según el layout que el kernel espere, puede que haya pasos previos para acomodar esa vista.

¿Por qué no aparecen kernels separados de multiplicación y suma?

En la traza de la capa lineal no se ve una aten::add separada (la suma del bias). Eso se debe a que el bias se integra dentro del kernel de multiplicación como un epílogo. En el contexto de GEMM (General Matrix Multiply), una epílogo es una pequeña operación que el kernel realiza justo antes de escribir el resultado en la memoria principal (HBM). Operaciones típicas de epílogo son añadir bias, aplicar una activación simple o escalar por una constante.

La ventaja es evidente: evitar una segunda lectura/escritura a memoria principal reduce tráfico de memoria, que suele ser el cuello de botella más caro.

En PyTorch, nn.Linear llama a torch.nn.functional.linear que luego despacha a aten::linear; si detecta bias, aten::linear invoca aten::addmm(bias, x, weight) en lugar de hacer matmul y add por separado. El kernel cuBLAS usado por addmm incluye la variante con bias en su epílogo, por eso en GPU solo se ve el kernel de GEMM y no un kernel extra para la suma.

¿Aporta algo torch.compile a una sola capa Linear?

Al compilar el forward de una sola nn.Linear y comparar las trazas eager y compiladas se observa:

• En GPU: se sigue viendo el mismo kernel cuBLAS GEMM con bias.
• En CPU: aparece la misma aten::addmm en las trazas.
• En CPU compilado: algunas filas adicionales asociadas a compile, pero sin cambios en el kernel principal.

La conclusión práctica es que, para una sola operación GEMM-con-bias, torch.compile no tiene mucho por fusionar: el kernel ya viene optimizado (incluye el epílogo) y no queda trabajo de fusión significativo. Esto no es un bug: compile necesita una cadena de operaciones mayor para poder transformar y fusionar en kernels más grandes.

Observaciones sobre el dispatch y los layouts en traces compilados

Si comparan con cuidado las trazas eager y compiladas notarán que la cadena de despacho en CPU suele ser más larga en modo eager: por ejemplo, aten::linear puede aparecer paseándose por aten::t antes de llegar a aten::addmm. En la traza compilada algunas de estas entradas en CPU desaparecen o se ven reducidas; la compilación reorganiza el flujo y puede eliminar pasos explícitos de metadata en CPU, o mover responsabilidades hacia la unidad que ejecuta el kernel.

No es necesario entrar en detalles internos para retener la idea: compile puede reducir overhead de CPU en escenarios donde hay varias operaciones encadenadas, pero para un único GEMM con epílogo ya optimizado no hay mucho que ganar.

Subiendo un escalón: el MLP

Para comprobar cómo se comporta el compilador con operaciones encadenadas, los siguientes scripts apilan varias nn.Linear con activaciones intermedias formando un bloque MLP (en el ejemplo, tres capas con activación entre ellas). Es precisamente en ese tipo de cadenas donde torch.compile puede tener oportunidades reales de fusión, eliminar pasos intermedios de scheduling y reducir overhead de dispatch en CPU.

Si su intención es optimizar modelos reales —por ejemplo modelos de clasificación o embeddings usados por equipos en LATAM— conviene perfilar tanto el caso de una sola capa como el del bloque completo; las ganancias de compile suelen verse cuando varias GEMM y operaciones menores se pueden combinar.

Recomendaciones prácticas para equipos y tomadores de decisión

• No asuman que compile siempre acelerará: primero perfilar. Para una sola nn.Linear no verán mejora significativa.
• Usen el perfilador de PyTorch y analicen CPU vs GPU lanes: si el tiempo está dominado por dispatch y pequeñas operaciones CPU, compile o fusión pueden ayudar; si está dominado por un gran kernel GEMM, optimicen HIP/BLAS y el batch size.
• Entiendan las epílogues: operaciones como la suma de bias o activaciones simples pueden estar ya integradas en el kernel de GEMM y no aparecerán como kernels separados.
• Experimenten en infraestructuras compartidas: Hugging Face facilita ejecutar estos scripts en GPUs (Spaces, Dev Mode o Jobs), lo que permite a equipos en LATAM reproducir y discutir las trazas sin invertir en hardware propio.

Conclusión

nn.Linear no introduce kernels adicionales para bias porque aprovecha la epílogo del GEMM; aten::t suele ser una operación de metadatos y no un kernel GPU. torch.compile no hace magia en una única capa lineal ya que el kernel principal ya está optimizado. Las ventajas de compile aparecen cuando hay varias operaciones encadenadas que se pueden fusionar para reducir overhead de CPU y tráfico de memoria. Perfilar con cuidado es la forma más eficiente de decidir dónde invertir tiempo de optimización.

Si desean, el próximo artículo de esta serie analiza en detalle el MLP de tres capas y cómo cambia la traza cuando hay más oportunidades de fusión.