¡Hola GPU, ¿qué hay de mi matriz?

Hi GPU, how about my matrix?

Una guía suave para entender cómo las GPUs realizan la multiplicación de matrices

Foto de Thomas Foster en Unsplash

La multiplicación de matrices; el santo grial de las redes neuronales profundas y los gigantes modernos de comprensión del lenguaje. Como MLEs o científicos de datos, nuestros dedos son demasiado rápidos para escribir tf.matmul o torch.matmul y nunca miramos atrás. Pero no me digas que nunca has tenido la infatuación de milisegundos de saber qué podría estar sucediendo con esa matriz cuando entra en la GPU. Si lo hiciste, estás en el lugar correcto. Acompáñame en un viaje a través de las fascinantes complejidades dentro de una GPU.

Te explicaré cómo estas potencias de cálculo machacan los números. Aprenderás tres cosas impresionantes que las GPUs hacen poco conocidas, cuando se enfrentan a matrices. Al final de esta publicación de blog, tendrás una buena comprensión de cómo funciona la multiplicación de matrices dentro de las GPUs.

GEMM: Un verdadero tesoro 💎 para una GPU

GEMM o multiplicación generalizada de matrices es el kernel que se ejecuta cuando las GPUs realizan la multiplicación de matrices.

C = a (A.B) + b C

Aquí, a y b son escalares, A es una matriz MxK, B es una matriz KxN, y por lo tanto, C es una matriz MxN. ¡Es fácil como eso! Puede que te preguntes por qué existe esa adición final. Resulta que este es un patrón bastante común para redes neuronales (por ejemplo, agregar sesgo, aplicar ReLU, agregar conexiones residuales).

Truco #1: El producto externo está fuera de este mundo 👽

Si se te pide que escribas un algoritmo de multiplicación de matrices desde los principios básicos, esto es lo que harás (a menos que te hayan regalado una GPU en lugar de un cerebro, ¡eso ahorraría dinero para un MLE!).

for (int i = 0; i < M; ++i)    for (int j = 0; j < N; ++j)        for (int k = 0; k < K; ++k)             C[i][j] += A[i][k] * B[k][j];

Aquí tienes una visualización animada que muestra lo que hace esto.

Multiplicación basada en el producto interno de dos matrices (Recreado por el autor - fuente de inspiración: https://www.adityaagrawal.net/blog/architecture/matrix_multiplication )

¿Pero sabías que a las GPUs les desagrada esta implementación 🤔? Para entender por qué es así, necesitas entender la arquitectura de memoria de la GPU,

Para todas las comparaciones y especificaciones, usaré las especificaciones de la GPU Nvidia A100.

Una GPU tiene tres niveles principales de memoria,

  • Memoria global o HBM (lo que normalmente se refiere como memoria de GPU y lo que se ve cuando se ejecuta nvidia-smi)
  • Memoria compartida (una memoria local que está dedicada a un solo multiprocesador de transmisión [o SM] y compartida entre hilos que se ejecutan en ese SM)
  • Registros (asignados individualmente a hilos para llevar a cabo su carga de trabajo)

Esto es lo que parece,

La jerarquía de memoria típica de una GPU (cachés L0/L1/L2 ignorados por simplicidad)

Lo primero que hay que tener en cuenta es que la memoria compartida (denominada SRAM a partir de ahora) es mucho más pequeña que la HBM, y mucho menos que los registros. Por lo tanto, tu matriz no cabrá allí (en la mayoría de las ocasiones). Si volvemos a nuestra animación, para una sola fila de A todas las columnas de B deben ser recuperadas, y repetir el proceso para todas las filas en A. Esto significa que la GPU necesita hacer muchas lecturas para calcular la salida. La HBM (~1,5TB/s) es varias magnitudes más lenta que la SRAM (~19TB/s).

Para ponerlo en números, digamos que quieres multiplicar una matriz de 10x20 y una matriz de 20x30, necesitas leer las columnas de B 10x30=300 veces. ¿Hay una mejor manera de hacer esto?

¡Resulta que un simple truco puede hacer mucho aquí! Simplemente cambia el orden de los bucles, de modo que k se convierta en el bucle externo. ¡Y listo! 😮

for (int k = 0; k < K; ++k)     for (int i = 0; i < M; ++i)        for (int j = 0; j < N; ++j)            C[i][j] += A[i][k] * B[k][j];

No tocamos el cálculo real, solo el orden de los bucles, así que deberíamos obtener el mismo resultado que antes. ¡Así es como se ve ahora la multiplicación de matrices!

Multiplicación de dos matrices basada en el producto externo (Recreado por el autor - fuente de inspiración: https://www.adityaagrawal.net/blog/architecture/matrix_multiplication )

Como puedes ver, solo traemos una columna de A y una fila de B a la vez y nunca miramos hacia atrás. Esto requiere muchas menos lecturas que la implementación original. La única diferencia es que antes estábamos calculando el producto interno entre dos vectores, ahora estamos calculando el producto externo.

La diferencia entre el producto interno y el producto externo se muestra en verde para dos vectores (azul y amarillo).

Pero aún así, necesitamos toda la matriz C en SRAM, lo que podría ser demasiado grande para caber en SRAM. ¿Qué hace CUDA entonces? Eso nos lleva al segundo truco.

Truco #2: Dividir y conquistar (y acumular)

¡No te preocupes! No voy a bombardearte con matemáticas complejas ni con algoritmos de Leetcode. Lo principal a tener en cuenta es que una matriz es una disposición en 2D de bloques individuales. La siguiente animación hace justicia a lo que estoy tratando de explicar.

Puedes iterar cada bloque en A y B y aún así calcular la respuesta exacta para el bloque correspondiente de C

El resultado del bloque verde 💚 es la franja de color azul claro de A 💙 y la franja de color amarillo claro de B 💛. Llevando esto un paso más allá, para calcular la salida, puedes traer un bloque de esa franja de A y un bloque de la franja de B a la vez, calcular la salida y acumular el resultado en el bloque verde.

Esto nos da un marco flexible donde podemos cargar un bloque (o mosaico) de cualquier tamaño de A y B y aún así calcular la respuesta final. No tenemos que detenernos allí, podemos seguir dividiendo recursivamente el problema en problemas aún más pequeños. es decir, la matriz se divide en mosaicos, los mosaicos se dividen en fragmentos y los fragmentos en valores individuales.

Usando el enfoque de teselado, el problema puede descomponerse de manera recursiva

Y esto se presta muy bien a la arquitectura de ejecución de procesos en una GPU. Hay tres capas para la ejecución del kernel en una GPU. Para simplificar, diremos que un SM ejecuta un solo bloque de hilos (aunque en la práctica los ejecutan concurrentemente, para reducir algo conocido como el efecto de la cola).

  • Hilos
  • Conjuntos de hilos (una colección de 32 hilos)
  • Bloques de hilos (una colección de varios conjuntos de hilos)

El número exacto de hilos en un bloque de hilos depende de una arquitectura específica. Por ejemplo, un A100 tiene las siguientes especificaciones.

  • Máximo de 2048 hilos por SM
  • Máximo de 1024 hilos por bloque
  • Máximo de 32 bloques de hilos por SM

Volviendo al teselado, se ha encontrado que (heurísticamente) un teselado de matriz de tamaño 256x128 por bloque de hilos da una eficiencia razonable para la mayoría de los problemas. Por lo tanto, es un tamaño de tesela común utilizado por CUDA.

Tal vez haya escuchado acerca de una buena práctica de mantener el tamaño del lote, el tamaño de la dimensión oculta como potencias de dos. ¡De esto es de donde viene esto! Cuando las dimensiones de su matriz son potencias de dos, será completamente divisible en un conjunto de teselas sin resto. Si no lo es, su código será menos eficiente.

Los cálculos de GPU son más eficientes cuando las dimensiones de su matriz son potencias de 2

¿Qué sucede cuando no es una potencia de 2?

Lo que sucede es un efecto conocido como cuantificación de teselas. En otras palabras, si tiene una dimensión de fila de tesela de 128 pero su matriz tiene 257 elementos en una fila, necesitará no dos, sino tres teselas en una fila (es decir, 256+1). Esto se ilustra a continuación.

Solo porque teníamos un elemento extra en las filas, debemos dedicar dos bloques de hilos completos

El problema con esto es que el bloque de hilos hace la misma cantidad de cálculos independientemente de los datos útiles que residan en él. Entonces, está perdiendo la oportunidad de hacer cálculos útiles en su GPU, lo que lleva a ineficiencias.

Un efecto similar se conoce como cuantificación de onda, donde la matriz es demasiado grande y los SM colectivamente no pueden ajustarla de una sola vez. Entonces, la GPU necesita hacer el cálculo en 2 “ondas”. Sin embargo, esto es menos preocupante para las GPUs modernas, ya que aprovechan la concurrencia para reducir la cuantificación de ondas.

La cuantificación de teselas ocurre cuando un bloque de hilos tiene que derramar datos parcialmente, la cuantificación de ondas ocurre cuando los SM tienen que derramar datos.

Truco #3: Uno es mejor que dos

El truco final es la fusión del kernel. Con mayor frecuencia, es más rápido hacer todas las operaciones en un solo kernel que tener dos kernels llamados uno después del otro. ¿Por qué? Porque un kernel necesita escribir los datos en HBM y el otro necesita leerlos de vuelta. Ya hablamos sobre lo lento que es esto. Un enfoque mejor es simplemente combinar las dos operaciones en una sola. Algunos ejemplos son,

  • matmul + bias + relu
  • conv + bias + relu
  • batchnorm + relu

Entonces, como se ve aquí (estoy seguro de que Pytorch tiene un glosario similar), hay muchos kernels fusionados ofrecidos a través de TensorFlow que combinan operaciones cómodas en un solo kernel. En el código, significa algo como esto,

for (int m = 0; m < M; m += Mtile)     for (int n = 0; n < N; n += Ntile)        for (int k = 0; k < K; ++k)            float tmp = 0            for (int i = 0; i < Mtile; ++i)                for (int j = 0; j < Ntile; ++j)                     int row = m + i;                    int col = n + j;                    tmp += A[row][k] * B[k][col];                    // Do other things                    C[row][col] = tmp

En otras palabras, valoramos nuestra variable tmp hasta que terminamos todos nuestros cálculos. Solo entonces escribiremos el resultado en C.

Conclusión

Eso es todo amigos. Espero que haya sido una excursión agradable a través de los detalles técnicos de una GPU. Si estás interesado en la versión audiovisual, aquí está el enlace a mi video de YouTube.

Para resumir, discutimos tres cosas que hacen que las GPUs sean realmente rápidas en la multiplicación de matrices:

  • Las GPUs abandonan la implementación de producto interno más amigable de matmul y adoptan la implementación de producto externo más eficiente en la lectura de matmul.
  • Las GPUs dividen las matrices en bloques más pequeños (y los bloques en fragmentos) y distribuyen la carga de cálculo entre bloques de hilos, warps y hilos.
  • Las GPUs emplean la fusión de kernel para unir funcionalidades comúnmente co-ocurrentes, mejorando la eficiencia de la GPU.

Si te ha gustado esta historia, no dudes en suscribirte a Zepes, y recibirás notificaciones de contenido nuevo mío, así como acceso completo a miles de historias de calidad de otros autores.

Como miembro de Zepes, una parte de tu tarifa de membresía va a los escritores que lees, y tienes acceso completo a cada historia…

thushv89.medium.com

A menos que se indique lo contrario, todas las imágenes son del autor.

Referencias:

  • https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html
  • https://developer.nvidia.com/blog/cutlass-linear-algebra-cuda/
  • https://developer.nvidia.com/blog/nvidia-ampere-architecture-in-depth/

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

Share:

Was this article helpful?

93 out of 132 found this helpful

Discover more

Inteligencia Artificial

LLM (Modelos de Lenguaje Grandes) para un Mejor Aprendizaje del Desarrollador de tu Producto

Explora cómo se pueden aprovechar los LLMs (Modelos de Lenguaje Grandes) y las aplicaciones de LLM para una educación...

Inteligencia Artificial

Visión del PM Modi sobre la regulación de la IA en India Cumbre B20 2023

A medida que el B20 Summit India 2023 llegaba a su fin en Delhi, las palabras del primer ministro Narendra Modi conti...

Inteligencia Artificial

Este artículo de IA propone un método de generación de memoria recursivo para mejorar la consistencia conversacional a largo plazo en modelos de lenguaje grandes

Los chatbots y otras formas de sistemas de comunicación de dominio abierto han experimentado un aumento de interés e ...

Inteligencia Artificial

Últimos avances en el campo de la IA multimodal (ChatGPT + DALLE 3) + (Google BARD + extensiones) y muchos más…

La IA multimodal es un campo de la Inteligencia Artificial (IA) que combina diferentes tipos de datos (modalidades), ...

Ciencias de la Computación

Cómo la inteligencia artificial protege (y ataca) tu bandeja de entrada.

Las empresas, como Google, están buscando formas en que la inteligencia artificial y el aprendizaje automático puedan...