Dos números en una GPU
Ya sabes cómo hacer c = a + b en Python. Lo has hecho mil veces. Aquí vamos a hablar de ello en tensores de PyTorch. Un tensor es solo un arreglo de números. Ponerlo en una GPU significa que ese arreglo vive en la memoria de la GPU en lugar de la memoria de la CPU normal. Cuando a y b son dos tensores en una GPU, esa línea se ejecuta lo suficientemente rápido como para que nunca pienses en ello.
Ahora redúcelo. Digamos que a y b son dos flotantes individuales, ambos en la GPU. La misma línea. c = a + b. ¿Qué es lo que realmente se ejecuta en el chip?
La respuesta es un kernel. En este mundo, un kernel es un pequeño programa que la GPU ejecuta sobre algún dato. No es el kernel del sistema operativo con el que arranca tu laptop. No es el núcleo matemático de un libro de álgebra lineal. La palabra se reutiliza mucho y eso no es tu culpa. En el mundo de las GPU, un kernel simplemente significa: una pequeña función que se le dice a la GPU que ejecute, ahora mismo, en paralelo, sobre los datos que le entregaste.
Al final de esto, serás capaz de mirar un fragmento de código de PyTorch y contar cuántos kernels ejecutará la GPU. Eso suena como un truco pequeño, y lo es, pero también es el primer paso para sacudir la sensación de la GPU como una caja negra, esa sensación cuando tu modelo es lento y no sabes por qué. Cada pregunta que puedas hacer sobre el rendimiento de la GPU eventualmente vuelve a "qué kernels se ejecutaron y qué estaban haciendo". Así que aquí es donde empezamos.
Tu primer kernel
Hagamos a y b un poco más grandes. Tensores de longitud 8 esta vez. Sigue siendo una línea de PyTorch: c = a + b.
Cuando ejecutas esto, tu CPU (la máquina que hace el Python real) le dice a la GPU: oye, ejecuta este programa sobre estos datos. Esa instrucción se llama un lanzamiento. Lo que se lanza es un kernel: un programa, listo para ejecutarse. Los lanzamientos en sí son baratos, microsegundos cada uno. Lo que está alrededor del lanzamiento (los datos que van a la GPU, los resultados que vuelven) es donde reside el costo real, y eso es lo que estaremos contando a medida que avanzamos.
Dentro del kernel, el trabajo real lo hacen pequeños trabajadores llamados hilos. Una GPU tiene miles de ellos disponibles. Para nuestra suma de longitud 8, 8 hilos toman el trabajo: el hilo 0 maneja el elemento 0, el hilo 1 maneja el elemento 1, y así sucesivamente hasta el hilo 7. Cada hilo ejecuta el mismo pequeño programa: leer un elemento de a, leer el elemento correspondiente de b, sumarlos, escribir el resultado en c.
(En la práctica, la GPU lanza hilos en grupos de tamaño fijo llamados warps, siempre 32 hilos en tarjetas NVIDIA, y enmascara los extras cuando tu arreglo no se divide uniformemente. Puedes ignorarlo por ahora.)
Así que tenemos una línea de PyTorch, un lanzamiento, un kernel, 8 hilos haciendo 8 sumas. Ahora contemos lo que realmente cruzó el chip. Para hacer la suma, cada hilo necesita su elemento de a y su elemento de b. Eso son 8 lecturas de a y 8 lecturas de b. Luego cada hilo escribe su resultado en c. Eso son 8 escrituras.
Esas lecturas y escrituras van a la gran memoria que está justo al lado del chip de la GPU. En tarjetas de centro de datos (A100, H100), esa memoria se llama HBM (memoria de alto ancho de banda). En tarjetas de consumo (RTX 4090) y las T4 estilo Colab que la gente probablemente usará para probar esto, se llama GDDR. De cualquier manera, es memoria rápida al lado del chip, y la llamaremos simplemente memoria de la GPU. Es rápida, pero no es gratis, y cada viaje a ella tiene un costo.
Un kernel = un lanzamiento = un pase sobre los datos. Haga lo que haga el kernel dentro de su cuerpo, las lecturas y escrituras en sus bordes (los viajes a la memoria de la GPU para obtener entradas, el viaje de regreso para escribir salidas) son la parte que cuesta. Esa es la forma completa.
Nada de esto cambia cuando los tensores se hacen más grandes. La misma línea de PyTorch, el mismo kernel, solo más hilos. Si a y b tienen un millón de elementos cada uno, la GPU lanza el mismo kernel con un enjambre más grande de hilos. Las matemáticas escalan, los bytes movidos escalan, el modelo mental no. Una línea, un kernel.
Qué sucede entre dos operaciones
1c = (a + b).relu()
Sabes Python lo suficiente como para saber que esto son dos operaciones, una suma y luego una relu, encadenadas. En un intérprete, eso son dos llamadas a funciones. En una GPU, en PyTorch modo eager, eso son dos lanzamientos de kernel: uno para la suma, uno para la relu. Hasta ahora, nada sorprendente.
Lo que realmente es interesante es lo que sucede entre los dos kernels.
Cuando la suma termina, tiene que poner su resultado en algún lugar. Ese lugar es la memoria de la GPU. La suma escribe un arreglo intermedio completo (llámalo tmp) a la memoria. Luego, un momento después, la relu se lanza, y su primer trabajo es leer ese mismo arreglo tmp de vuelta desde la memoria. Lee todo, aplica relu a cada elemento, escribe el resultado en c.
Cuenta el tráfico de memoria para esos dos kernels:
- La suma: lee
a, leeb, escribetmp. Tres transferencias del tamaño del arreglo. - La relu: lee
tmp, escribec. Dos más.
Cinco transferencias del tamaño del arreglo en total. Compara eso con la suma de longitud 8 por sí sola de la sección anterior, que fueron tres. Agregar .relu() a la cadena no solo te costó el cómputo de la relu. Te costó un viaje de ida y vuelta completo del arreglo a través de la memoria de la GPU, porque tmp tuvo que ser escrito solo para que el siguiente kernel pudiera leerlo de vuelta.
Nada se almacenó en caché. tmp no pudo quedarse en un registro o en una caché local rápida. Salió a la memoria de la GPU (del tipo lento y lejano) y volvió inmediatamente. Los dos kernels son extraños entre sí. Tienen que pasarse datos a través del único medio que ambos saben usar: la memoria de la GPU.
¿Por qué PyTorch lo hace así? Porque en modo eager, cuando escribes a + b, PyTorch lo ejecuta \ahora**. No sabe que estabas a punto de llamar a .relu() a continuación. Cada operación se despacha en el momento en que se ejecuta su línea de Python. No hay plan, no hay anticipación. Cada operación es independiente, produce un arreglo real y lo pasa a lo que venga después a través de la memoria.
Este es el patrón que debes recordar. Cada valor intermedio en tu código de PyTorch se escribe físicamente en la memoria de la GPU y es leído de vuelta por la siguiente operación. Todos y cada uno. Eso es lo que realmente mide el "conteo de kernels". Cada kernel extra es otro viaje de ida y vuelta que tus datos tienen que hacer a través de la memoria de la GPU.
Fusión: dos operaciones, un kernel
Imagina un kernel que hace todo esto de una sola vez: lee su elemento de a, lee su elemento de b, los suma, aplica relu al resultado (todo dentro del kernel, en un pequeño espacio de trabajo por hilo que nunca sale del chip), y solo entonces escribe el valor final en c. El intermedio (a + b) sigue existiendo, pero solo dentro del kernel, en el espacio de trabajo privado de cada hilo. Nunca se escribe en la memoria de la GPU. tmp, como arreglo real, no existe en absoluto.
Cuenta las transferencias ahora. Lecturas de a: 1 por elemento. Lecturas de b: 1 por elemento. Escrituras de c: 1 por elemento. Tres transferencias del tamaño del arreglo. Mismas matemáticas que la versión de dos kernels, pero dos viajes de ida y vuelta menos.
En longitud 8, esto es un error de redondeo. A nadie le importa. En longitud 1 millón, o 100 millones, esos viajes de ida y vuelta extra se convierten en una gran parte del tiempo de ejecución, y el reloj de pared lo refleja. \Por qué** el tráfico de memoria termina dominando así es el tema completo del Artículo 2 de esta serie, así que dejaré el "por qué" aquí. El punto ahora es solo: mismas matemáticas, menos viajes, más rápido en la práctica.
Ese truco, combinar operaciones que habrían sido kernels separados en un solo kernel para que el intermedio nunca tenga que visitar la memoria de la GPU, tiene un nombre. Se llama fusión. Esa es toda la palabra. Esa es toda la idea.
Ahora la parte incómoda. Escribir ese kernel combinado a mano parece fácil para suma + relu. Dos operaciones. Una línea de "cómputo" en el medio. Pero el código real de PyTorch tiene docenas de operaciones encadenadas, cada una con sus propias formas, tipos de datos y reglas de broadcasting. Escribir un kernel fusionado que maneje todo eso correctamente es un trabajo de ingeniería real. Generalmente no escribirías estos kernels elementales de rutina a mano.
Buenas noticias: PyTorch incluye una herramienta que hace la reescritura por ti, exactamente para este tipo de caso, automáticamente. Se llama torch.compile.
Casi seguro has visto torch.compile usado así:
1model = torch.compile(model)
Una línea. Alguien en internet te dijo que hace las cosas más rápidas. Esto es lo que realmente hace, en español claro: en lugar de ejecutar tus operaciones una por una como lo hace el modo eager, torch.compile captura las operaciones de tensores que realiza tu función, busca oportunidades para combinarlas y genera código optimizado. Llamadas posteriores que coincidan con las mismas suposiciones pueden reutilizar ese código.
La fusión que hicimos en papel arriba (suma y relu, compartiendo un kernel, tmp nunca tocando memoria) es exactamente el tipo de cosa que torch.compile hará a tu código automáticamente, siempre que las operaciones sean lo suficientemente simples. Cuando la gente dice que torch.compile "hace PyTorch más rápido", esta es una gran parte de lo que quieren decir.
Para los casos en que torch.compile no puede fusionar por sí solo (operaciones personalizadas que no reconoce, reducciones inusuales, diseños de memoria extraños), alguien todavía tiene que escribir un kernel a mano. Para eso sirven herramientas como Triton y CUDA. Un artículo aparte.
Compruébalo tú mismo
Todo lo anterior ha sido contar kernels en papel. Es hora de contarlos en una GPU real. Si tienes alguna máquina con una GPU CUDA a mano (una estación de trabajo, un notebook de Colab, una instancia en la nube), puedes ejecutar esto tú mismo en unos minutos.
La herramienta es torch.profiler. Está integrada en PyTorch. Todo lo que hace es registrar lo que la GPU realmente hizo mientras tu código se ejecutaba, y te devuelve una tabla que puedes leer.
Paso 1: la versión eager.
Envuelve la línea de dos operaciones en una función para tener algo que llamar:
1import torch2from torch.profiler import profile, ProfilerActivity34def add_relu(a, b):5 return (a + b).relu()67a = torch.randn(1_000_000, device="cuda")8b = torch.randn(1_000_000, device="cuda")910with profile(activities=[ProfilerActivity.CUDA]) as prof:11 add_relu(a, b)12 torch.cuda.synchronize()1314print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
torch.cuda.synchronize() está ahí solo para asegurarse de que la GPU haya terminado antes de leer los tiempos. El trabajo de la GPU se ejecuta de forma asíncrona, y sin la sincronización a veces medirías la sobrecarga del lanzamiento en lugar del trabajo real del kernel.
Paso 2: lee la salida.
Tu salida real del profiler tendrá más filas de las que podrías esperar. Un montón de líneas de asignación de memoria y contabilidad de PyTorch estarán mezcladas. Las filas que nos importan son los kernels de CUDA, las funciones reales que la GPU ejecutó. Busca filas con kernel en el nombre. Las dos se verán más o menos así:
1vectorized_elementwise_kernel<...BinaryFunctor_add...> 1 12us2vectorized_elementwise_kernel<...threshold_kernel...> 1 9us
Los nombres exactos de las plantillas cambian entre versiones de PyTorch (relu a menudo aparece como threshold porque esa es la operación subyacente, y add a veces como CUDAFunctor_add). No intentes analizar todo. Solo cuenta las filas. Dos filas. Dos kernels. Uno para la suma, uno para la relu. Exactamente lo que dijimos que sucedería hace una sección.
Paso 3: la versión compilada.
Un cambio de línea. Envuelve la función en torch.compile:
1compiled = torch.compile(add_relu)
Ahora aquí hay un detalle que vale la pena saber antes de ejecutarlo. La \primera llamada a una función compilada con `torch.compile` es lenta, a veces extrañamente lenta, porque es cuando `torch.compile` hace su trabajo: analizar tu código, descubrir qué fusionar, generar el kernel fusionado. Si perfilas la primera llamada, estás midiendo el paso de compilación, no el kernel. Así que el patrón es: llámala una vez para calentarla, descarta el resultado, luego** perfilas.
1compiled(a, b) # calentamiento, descarta el resultado2torch.cuda.synchronize()34with profile(activities=[ProfilerActivity.CUDA]) as prof:5 compiled(a, b)6 torch.cuda.synchronize()78print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
Paso 4: lee la salida de nuevo.
Esta vez:
1triton_poi_fused_add_relu_0 1 14us
Una fila. Un kernel. El nombre incluso te dice lo que hizo: suma y relu fusionadas. Mismas matemáticas que antes, un lanzamiento en lugar de dos.
Acabas de hacer lo que este artículo ha estado hablando, en una oración: le pediste a PyTorch que combinara dos operaciones en un kernel, observaste el profiler y confirmaste que el conteo pasó de dos a uno. Fusión, en la naturaleza, en tu máquina.
Si quieres verlo de manera más dramática, pruébalo con algunos tamaños de tensor diferentes. En longitud 100, ambas versiones se ejecutan tan rápido que la diferencia se pierde en el ruido. En longitud 10 millones o 100 millones, la versión compilada comienza a adelantarse claramente, porque el viaje de ida y vuelta que eliminamos es una parte real del trabajo a esa escala.
Contar kernels ya no es un consejo abstracto. Tienes una manera de verificarlo.
Resumamos
Aquí está todo de una vez.
Tu código de PyTorch, cuando se ejecuta en una GPU, se convierte en una secuencia de kernels. Cada kernel es un lanzamiento, un pase sobre tus datos, un viaje de ida y vuelta a través de la memoria de la GPU para obtener entradas y escribir salidas. Las operaciones simples se convierten en un kernel. Las cadenas de operaciones se convierten en un kernel \por operación** por defecto, con intermedios viajando de ida y vuelta a través de la memoria entre ellos. torch.compile puede fusionar cadenas simples por ti para que esos intermedios nunca toquen la memoria. Menos kernels generalmente significa menos tráfico de memoria. Y menos tráfico de memoria generalmente significa más velocidad.
