Multiplicación de matrices en la GPU

Multiplicación de matrices en la GPU Potenciando el rendimiento

Cómo lograr un rendimiento de multiplicación de matrices de última generación en CUDA.

Este blog surgió de la repentina realización de lo poco que sabía sobre cómo funciona la multiplicación de matrices en la GPU. Después de haber realizado tantos proyectos de ML, siento que debo entender cómo funciona la operación más importante en ML: ¿Qué es esta cosa llamada “Tensor Core”? ¿Por qué todos dicen que “el movimiento de datos es el cuello de botella”? ¿Qué tan rápido pueden llegar a ser las GPUs?

Para responder a estas preguntas, decidí que debo salir de mi burbuja de PyTorch y adentrarme en el abismo de CUDA. Escribí este blog para documentar todo lo que he aprendido y espero que cualquier persona que lo lea no tenga que pasar por el dolor de buscar en la documentación o el código de CUDA como yo lo hice.

Si hay algo que he aprendido en este viaje, es que la multiplicación de matrices concurrente es DIFÍCIL. La eficiente multiplicación de matrices depende en gran medida del hardware específico que estés utilizando y del tamaño del problema que estés tratando de resolver. No hay una solución única para todo.

Suficiente quejarse, ¡vamos a profundizar!

• Cómo ampliar los DataFrames de Pandas con métodos personalizados para potenciar la funcionalidad y legibilidad del código
• No olvides los intervalos de confianza para tu producto de ML
• Aumentando LLMs con RAG

Recapitulación de la arquitectura de la GPU

Recordemos cómo funcionan las GPUs (NVIDIA). Una GPU logra la paralelización ejecutando muchos hilos (threads). Cada hilo se ejecuta en un único núcleo de CUDA, aunque en un momento dado, solo un subconjunto de los hilos está activo, por lo que puede haber muchos más hilos que núcleos de CUDA disponibles. Cada hilo, activo o no, tiene su propio conjunto de registros.

Un grupo de 32 hilos se conoce como un grupo de hilos (warp). Todos los hilos en un warp deben ejecutarse juntos (o estar inactivos juntos). En la mayoría de los casos, hay muchos warps inactivos en comparación con los warps activos, y el planificador de warps es el encargado de elegir qué warps ejecutar en un momento dado. Esto permite que la GPU oculte la latencia de los accesos a la memoria programando la ejecución de otros warps mientras un warp está esperando datos.

Un grupo de warps se conoce como un bloque de hilos (threadblock). Todos los warps en un bloque de hilos se ejecutan en el mismo Multiprocesador de Transmisión (Streaming Multiprocessor) (SM). Cada bloque de hilos tiene su propia memoria compartida (shared memory) a la que todos los hilos en el bloque de hilos pueden acceder.

Nota: Arquitecturas más recientes

A partir de la arquitectura Volta en adelante, cada hilo también tiene su propio contador de programa y pila de llamadas, etc. Esto significa que cada hilo en un warp puede ejecutar instrucciones diferentes al mismo tiempo.

La arquitectura Volta también introdujo Tensor Cores especializados en resolver multiplicaciones de matrices de tamaños específicos. Cada warp activo tiene acceso a un Tensor Core.

En la última arquitectura Hopper, hay un concepto de clusters de bloques de hilos (threadblock clusters) que representa un grupo de bloques de hilos. Esto le da al usuario un control más preciso sobre la programación de bloques de hilos y permite que la memoria compartida de un bloque de hilos sea accedida por otros bloques de hilos en el mismo cluster.

Paralelizar la multiplicación de matrices

Supongamos que queremos calcular:

Diremos que el tamaño del problema es (M, N, K) en este caso. Para paralelizar esta operación, podemos dividir A y B en matrices más pequeñas, multiplicar las matrices individualmente y concatenar los resultados para formar C.

Específicamente, podemos dividir A por filas (es decir, dividir M en trozos de tamaño M’) y B por columnas (es decir, dividir N en trozos de tamaño N’) para obtener:

Podemos ver que cada sub-matriz en C es independiente entre sí, por lo que podemos paralelizar fácilmente el cálculo de cada sub-matriz.

En la práctica, K podría ser demasiado grande para cargar directamente en la memoria y realizar cálculos. En su lugar, una implementación típica también dividirá K en fragmentos de tamaño K’, iterará sobre cada fragmento y acumulará (sumando) los resultados parciales. Esto se conoce como reducción en serie de K. (A diferencia de la reducción en paralelo de K, ver sección a continuación). Matemáticamente, esto se ve así:

Nota: Añadir relleno

En cualquier momento en que el tamaño del problema no sea divisible por el tamaño de la partición, necesitamos añadir relleno. Esto se hace típicamente de forma implícita cuando cargamos los datos particionados (𝐴ᵢ,ₖ y 𝐵ₖ,ⱼ) en la memoria de nivel inferior, donde nos aseguramos de que la partición cargada (de tamaño M’×K’ para 𝐴ᵢ,ₖ y K’×N’ for 𝐵ₖ,ⱼ) siempre esté “completa” añadiendo ceros. Se debe tener cuidado especial al escribir los resultados de nuevo en la memoria global para evitar errores de desbordamiento.

En un alto nivel, tres particiones anidadas se utilizan para paralelizar la multiplicación de matrices en la GPU:1. La primera partición ocurre a nivel de bloque de hilos. Cada bloque de hilos es responsable de calcular Cᵢ,ⱼ = Aᵢ Bⱼ.2. La segunda partición ocurre a nivel de gavilla. El problema a nivel de bloque de hilos Cᵢ,ⱼ se particiona aún más para que cada gavilla sea responsable de calcular Cᵢ,ⱼ⁽ᵐⁿ⁾ = Aᵢ⁽ᵐ⁾ Bⱼ⁽ⁿ⁾.3. La tercera partición ocurre a nivel de instrucción. Algunas instrucciones esperan entradas de tamaños particulares. Por ejemplo, las Tensor Cores de segunda generación operan en problemas de tamaño (16, 8, 8) para fp16, mientras que una implementación directa en los núcleos CUDA por multiplicación escalar simplemente opera en tamaño (1, 1, 1). El problema a nivel de gavilla se particiona aún más para que cada fragmento tenga un tamaño adecuado para la instrucción: Cᵢ,ⱼ⁽ᵐⁿ⁾⁽ᵃᵇ⁾ = Aᵢ⁽ᵐ⁾⁽ᵃ⁾ Bⱼ⁽ⁿ⁾⁽ᵇ⁾.

Hay una buena razón por la que necesitamos tres niveles de partición, como veremos en la siguiente sección.

Redundancia de datos

La multiplicación de matrices puede volverse fácilmente limitada por la memoria si recuperamos datos de memoria global en registros cada vez que realizamos un cálculo. La observación clave es que muchos de los sub-entradas Aᵢ y Bⱼ se reutilizan en diferentes multiplicaciones de submatrices. Por ejemplo, se requiere Aᵢ para Cᵢ,₁, Cᵢ,₂, … y se requiere Bⱼ para C₁,ⱼ, C₂,ⱼ, … . Podemos obtener el mejor rendimiento si podemos minimizar el movimiento de datos redundantes y reutilizar los datos cargados tanto como sea posible.

En CUDA, existen tres tipos de memoria accesible para el usuario:

Esta es una visión general de cómo se utiliza cada tipo de memoria:

1. Cada bloque de hilos primero cargará sus entradas requeridas desde la memoria global a la memoria compartida. Los accesos posteriores a esos datos serán atendidos por la memoria compartida en lugar de por la más lenta memoria global.
2. Dentro de cada bloque de hilos, cada gavilla cargará primero sus entradas requeridas desde la memoria compartida a los registros. Los accesos posteriores a esos datos serán atendidos directamente por los rápidos registros.

Sumergiéndonos en los detalles

Nivel de bloque de hilos

En el nivel de bloque de hilo, el problema se divide en subproblemas de tamaño (M’, N’, K’). Así, cada bloque de hilo es responsable de calcular un fragmento de C, denotado como:

Se minimiza el movimiento redundante de datos cargando las subentradas Aᵢ,ₖ y Bₖ,ⱼ en memoria compartida. Cuando terminamos de calcular Aᵢ,ₖ Bₖ,ⱼ, se cargará el siguiente fragmento (Aᵢ,ₖ₊₁ y Bₖ₊₁,ⱼ).

Nivel de hilo

En el nivel de hilo, el subproblema se divide aún más en sub-subproblemas de tamaño (M”’, N”’, K”’). Así, cada warp es responsable de calcular un fragmento de Cᵢ,ⱼ, denotado como Cᵢ,ⱼ⁽ᵐ ⁿ⁾:

Se minimiza el movimiento redundante de datos cargando las subsubentradas Aᵢ,ₖ⁽ᵐ ˡ⁾ y Bₖ,ⱼ⁽ˡ ⁿ⁾ en registros. Cualquier acceso a Aᵢ,ₖ⁽ᵐ ˡ⁾ y Bₖ,ⱼ⁽ˡ ⁿ⁾ dentro de un warp será atendido por los registros rápidos.

Nota: Distribuir datos en registros

Vale la pena destacar que los registros son solo de nivel de hilo. Esto significa que las entradas en un registro no pueden ser accedidas por otros hilos en un warp. La forma exacta en que se dividen Aᵢ,ₖ⁽ᵐ ˡ⁾ y Bₖ,ⱼ⁽ˡ ⁿ⁾ en los registros de cada hilo depende de la instrucción específica utilizada. La documentación de NVIDIA sobre Instrucciones de multiplicación-acumulación de matrices a nivel de warp proporciona una descripción detallada para cada instrucción.

Nivel de la tensor core

Para realizar la multiplicación de matrices, utilizamos las Tensor Cores en la GPU. Mi GPU (RTX 2060) tiene las Tensor Cores de segunda generación, que están especializadas en resolver problemas fp16 de tamaño (M’’‘, N’’‘, K’’’) = (16, 8, 8). Por lo tanto, dividimos aún más Cᵢ,ⱼ⁽ᵐ ⁿ⁾ para que se ajuste al tamaño esperado por la instrucción:

Aquí, todas las entradas ya están en los registros y, por lo tanto, el costo de movimiento de datos es mínimo.

Nota: Tensor Cores

Las operaciones de Tensor Core son instrucciones a nivel de warp, lo que significa que todos los hilos en un warp deben ejecutar la instrucción de Tensor Core al mismo tiempo, preparando colaborativamente los datos para ser consumidos por una Tensor Core.

Elección de los tamaños de partición

Entonces, dado que queremos minimizar el movimiento de datos, ¿deberíamos simplemente elegir un tamaño de partición lo más grande posible para utilizar toda la memoria compartida y los registros, ¿verdad? Bueno, no exactamente.

Tamaño de partición del bloque de hilo

Asintóticamente, a medida que el tamaño del problema aumenta, sí, queremos usar tanta memoria compartida y registros como sea posible. Sin embargo, para tamaños de problemas pequeños, podríamos enfrentar dos problemas:

1. Tener un tamaño de partición grande significa que tendremos menos bloques de hilo. Dado que cada bloque de hilo solo puede ejecutarse en un SM, esto podría significar que no podemos utilizar todos los SM.
2. Para tamaños de problemas que no son divisibles por el tamaño de la partición, tendremos que agregar más relleno a las entradas. Esto reduce la eficiencia, ya que se realizará menos cálculo en entradas significativas.

Una implementación típica podría usar un tamaño de partición de (M’, N’, K’) = (128, 256, 32).

Tamaño de partición de warps

En general, tener un tamaño de partición de warps grande significa que habrá menos movimiento de datos redundante, pero a costa de tener menos warps. Tener muy pocos warps significa que no podremos ocultar la latencia de los accesos a la memoria (porque podríamos quedarnos sin otros warps para programar mientras el warp actual espera los datos).

Una implementación típica podría usar un tamaño de partición de (M’‘, N’‘, K’’) = (64, 64, 32).

Tamaño de partición de instrucciones

Esto está completamente determinado por las instrucciones que admita tu GPU. Para mi RTX 2060, la instrucción ptx para la multiplicación de matrices Tensor Core fp16 (con acumulación fp16) es mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16, que espera entradas de tamaño (16, 8, 8).

Aún más optimizaciones

Las técnicas anteriores pueden acercarnos al rendimiento máximo teórico de la GPU cuando el tamaño del problema es grande. Sin embargo, para tamaños de problemas más pequeños, no son tan eficientes. Hay dos técnicas comunes para mejorar aún más el rendimiento de la multiplicación de matrices: reducción paralela-K y pipeline de software.

Reducción paralela-K

En casos en los que M y N son pequeños, es posible que tengamos solo algunos bloques de hilos. Por ejemplo, si (M’, N’) = (128, 256) y el tamaño del problema original tiene M ≤ 128 y N ≤ 256, solo tendremos un bloque de hilos, ¡y solo utilizamos una fracción de la potencia de cálculo de la GPU! (Por ejemplo, mi RTX 2060 tiene 30 SM, por lo que para maximizar la utilización queremos al menos 30 bloques de hilos).

En casos en los que K es grande (aunque M y N son pequeños), podemos aprovechar más paralelismo utilizando reducción paralela-K. Recordemos que en la reducción serial-K, cada bloque de hilos itera sobre la siguiente suma:

y acumula los resultados intermedios en Cᵢ,ⱼ. En la reducción paralela-K, en cambio, asignamos a cada bloque de hilos la tarea de calcular solo un elemento de la suma (es decir, Aᵢ,ₖ Bₖ,ⱼ). Esto nos permite aumentar el número de bloques de hilos en un factor de K/K’, aprovechando así más SM.

La única salvedad es que ahora necesitamos asignar más memoria para almacenar los resultados de cada bloque de hilos e invocar un segundo kernel para realizar una reducción final sobre los resultados parciales y obtener Cᵢ,ⱼ.

Pipeline de software

Normalmente, CUDA oculta la latencia de los accesos a memoria programando la ejecución de otros warps mientras un warp espera datos. Esto requiere que tengamos suficientes warps para ocultar la latencia.

Sin embargo, el número de warps suele ser relativamente pequeño al hacer GEMM. Esto se debe a que el número de warps está limitado por “El número de registros disponibles por bloque de hilos dividido por el número de registros requeridos por warp”. Para la multiplicación de matrices, utilizamos muchos registros por warp para maximizar el reuso de datos. Como resultado, es posible que no tengamos suficientes warps para ocultar la latencia.

“Los elementos del acumulador suelen ocupar al menos la mitad del presupuesto total de registros de un hilo”. — Documentación de CUTLASS

Para mitigar este efecto, podemos utilizar pipeline de software. En esencia, podemos (manualmente) cargar previamente de forma asíncrona las entradas para la próxima iteración del bucle utilizando instrucciones especiales. Mientras se cargan las entradas, podemos seguir calculando en la iteración actual. Se resume con el siguiente diagrama:

Esto es posible gracias a que la GPU es como cualquier CPU moderna: puede realizar accesos a memoria y operaciones aritméticas en pipeline siempre que no haya dependencia de datos entre ellos. Esto se conoce como paralelismo a nivel de instrucción.

Multiplicación de matrices en acción

Si quieres ver cómo todos estos conceptos se unen en una implementación real, echa un vistazo a mi implementación de entrenamiento de MNIST desde cero con CUDA. Allí, entrené un perceptrón multicapa en MNIST utilizando CUDA, logrando una aceleración 6 veces más rápida que PyTorch optimizado para un tamaño oculto de 128:

GitHub – andylolu2/cuda-mnist

Contribuye al desarrollo de andylolu2/cuda-mnist creando una cuenta en GitHub.

github.com

Referencias

1. Documentación de CUTLASS2. Documentación de CUDA3. Ejemplos de CUTLASS

We will continue to update Zepes; if you have any questions or suggestions, please contact us!