NVIDIA (III): La arquitectura Ampère (A100)
Rubén Rodríguez Abril
La arquitectura Ampère de NVIDIA, representada por la A100, perfeccionó las anteriores (Volta y Turing) sin cambios revolucionarios, enfocándose en mejoras en aritmética de coma flotante, el principio de escasez (sparsity), optimización de la transferencia de datos y la creación de múltiples GPUs virtuales. La A100 ha sido crucial en el entrenamiento de modelos avanzados de IA, como GPT-3, DALL-E 2/3 y Stable Diffusion.
La arquitectura Ampère de NVIDIA (representada por su emblemática tarjeta A100) continuó desarrollando las innovaciones de su predecesora, Volta. A mi juicio, ninguno de los cambios introducidos fue verdaderamente revolucionario, sino que continuaron tendencias de desarrollo ya existentes en arquitecturas anteriores. En este artículo, incidiremos en cuatro aspectos: las mejoras en la aritmética de coma flotante, el principio de escasez (sparsity), mejoras en la transferencia de datos así como la posibilidad de crear varias GPUs virtuales.
Dentro de A100, sus 108 multiprocesadores SM, analizados en los artículos anteriores, se agrupan en clústeres. Hay dos tipos de clústeres:
-TPCs (Texture Processor Clusters), encargados del procesamiento de texturas e integrados por 2 SMs.
-GPCs (General Processor Clusters), encargados de tareas de procesamiento general, e integrados por 8 TPCs y otros 16 SMs sueltos. Hay 7 GPCs en todo el chip.
Figura 1. Detalle de la estructura de una GPU A100. PCIExpress se encarga de las comunicaciones con la CPU y la placa base, mientras que el GigaThread Engine (en naranja) distribuye el trabajo entre los clústers. Los cachés L1 se denominan “memoria local”, el caché L2 “memoria compartida” y HBM2 “memoria global”. Fuente: NVIDIA.
Innovaciones aritméticas
Se sustancian, fundamentalmente, en la introducción de dos nuevos formatos de aritmética en coma flotante, así como la expansión de los Tensor Cores a operaciones con todo tipo de números (en Volta sólo realizaban tareas de semiprecisión).
Aritmética en coma flotante
El formato de números en coma flotante fue introducido por primera vez por el español Leonardo Torres Quevedo. Es una adaptación a la aritmética binaria de la representación en decimales (que como su propio nombre indica usa la base 10). En la notación de coma flotante cada número es descrito por tres variables, la mantisa, el exponente y el signo, mediante la siguiente ecuación:
N = (-1)sm2x
donde N es el número representado, s es el signo (0, si es positivo, 1, si es negativo), m es la mantisa y x el exponente. La mantisa es, por así decirlo, el núcleo del número. Y el exponente representa el número de puestos que tiene que, en aritmética binaria, tiene que desplazarse la coma hacia la izquierda (si el exponente es negativo) o hacia la derecha (si el exponente es positivo). El desplazamiento hacia la derecha equivale a la agregación tantos ceros a la mantisa como señale el exponente. Así, el número binario 1,1 (1,5, en decimal) tiene como mantisa 11 y como exponente -1. El número 1100 (12, en decimal) tiene la misma mantisa, pero un exponente de 2.
En los sistemas de computación modernos, las tres variables se integran como secciones de una única cadena de bits. El estándar IEEE754 describe tres tipos de formatos:
-Precisión doble (double, 64 bits): 11 bits de exponente y 52 bits de mantisa.
-Precisión simple (float, 32 bits): 8 bits de exponente y 23 bits de mantisa.
-Semiprecisión (half, 16 bits): 5 bits de exponente y 10 bits de mantisa.
Figura 2: En la imagen se describen tres formatos de números en coma flotante, cada uno de un tamaño diferente (16, 32 y 64 bits). El primer bit (naranja) es de signo. Le siguen los bits de exponente (en amarillo) y mantisa (en azul). Fuente: NVIDIA.
Los sistemas de multiprecisión son aquellos en los que el procesador es capaz de manipular números de varios formatos, mientras que en los de precisión mixta el sistema produce números de 32 bits a partir de datos de entrada en 16 bits.
Los formatos numéricos TF32 y BFLOAT16
La arquitectura Ampere de NVIDIA (a la que pertenecen sus tarjetas A100) introduce un nuevo formato matemático denominado TensorFloat 32 (TF32), en el que se usa un exponente de 8 bits (como en un número de precisión simple), pero la mantisa queda reducida a tan sólo 10 bits, sumando un total de 19 bits, tal y como se muestra en la imagen.
Figura 3. Comparación del formato TF32 con los formatos FP16 y FP32. La mantisa (turquesa oscuro) de TF32 es idéntica a la del formato de semiprecisión (FP16), mientras que el exponente (verde claro) es similar a la del formato de precisión simple (FP32). Fuente: NVIDIA.
El formato BFLOAT16 tiene una mantisa de 7 bits y un exponente de 8.
La sustitución del formato de FP32 por el de TF32 (de tan sólo 19 bits) en los modelos A100 supuso un incremento de la velocidad de entrenamiento de algunos modelos de Deep Learning (en el caso de BERT, del orden de 3 a 6 veces), sin que la eficiencia de dichos modelos se viesen apenas afectada.
Figura 4. En el diagrama se realiza una comparación entre la velocidad de entrenamiento del modelo BERT. Fuente: NVIDIA.
Tensor CORES: generalización a todo tipo de formatos numéricos
En Volta, los Tensor Cores sólo podían realizar operaciones sobre números FP16. En Turing (2018) y en Ampere (2020), en cambio, el circuitaje de los Tensor Cores es mejorado de tal manera que también pueden realizar operaciones también sobre números enteros en INT8 (entero de 8 bits), INT4 y INT1 (booleano).
Figura 5. Esquema con las diferentes velocidades de computación de un Tensor Core de Turing (arquitectura inmediatamente predecesora de Ampere). En enteros de 4 y 8 bits, los cálculos son 4 y dos veces más rápidos que en el caso de los números en coma flotante de 16 bits. Fuente: NVIDIA.
Como se ve, cuanto menor es la longitud en bits de los operandos, más operaciones (en este caso, multiplicaciones-adiciones) por segundo se pueden realizar. Recordemos que en una multiplicación de dos matrices, con dimensión m x k y k x n, la cantidad de multiplicaciones-adiciones que hay que realizar es de m x n x k (un producto escalar de tamaño k por cada elemento de la matriz de destino). En A100, los Tensor Cores pueden realizar una multiplicación matricial 16x8x16 (frente a 4x4x4 ede Volta) por cada ciclo del reloj.
Principio de escasez (sparsity)
Los autores del trabajo de GoogLeNet, introdujeron por primera vez el principio de escasez (sparsity principle), procedente de la neurología, en el ámbito del Aprendizaje Profundo. Los investigadores se habían percatado que sólo un pequeño porcentaje de neuronas (en torno al 5%) está activa en cada momento y se preguntaron si algo parecido podría implementarse en las redes neuronales artificiales, con el objeto de impedir el overfitting, y facilitar las tareas de generalización.
En la actualidad, el principio de escasez implica que buena parte (o la mayoría) de los coeficientes de las matrices de los modelos de IA son nulos, y que la mayor parte de la información se almacena en áreas reducidas de la matriz. Así sucede en el caso de las capas convolucionales y de las matrices de atención (en el que las que los coeficientes, a medida que nos alejamos de la diagonal tienden rápidamente a 0). De este modo, se evita que el modelo memorice datos y se favorece que las capas del mismo almacenen información relativa a patrones y no bits aislados de información.
Figura 6. Tras el preentrenamiento, la matriz es “podada”, anulando dos de cada cuatro coeficientes de la matriz. Los coeficientes no nulos se siguen actualizando durante la fase de afinamiento. En la inferencia, cuando llega el momento de multiplicar pesos por activaciones de entrada, los mismos elementos que están a cero en el vector-fila (verde), también se anulan en el vector-columna (en rojo). De este modo, la profundidad k de la multiplicación matricial se reduce a la mitad. Fuente: NVIDIA.
La arquitectura Ampère supuso la introducción de este princpio de escasez (sparsity) también en el hardware. Los ingenieros de NVIDIA realizaron experimentos en los que durante la fase de inferencia (evaluación) pusieron a cero la mitad de los coeficientes de las matrices. Para su sorpresa, el rendimiento de los sistemas de IA apenas de vio afectado. Consiguientemente, introdujeron un modelo de “podado” de matrices, en cuya virtud cada matriz de pesos sinápticos es dividida en ternas de cuatro coeficientes, de los cuales se anula a los dos con valores más bajos. La velocidad de los Tensor Cores se multiplica por dos.
Transferencia de datos
Los SMs de los chips de NVIDIA no son ajenos al problema de cuello de botella de Von Neumann: Por lo general, la capacidad del procesador para realizar cálculos es mayor que la tasa de transferencia de datos desde o hacia la memoria, lo que provoca que se infrautilice el procesador. En una arquitectura de Von Neumann, que es aquella que usa una memoria y un bus comunes para código y datos, este programa se agrava.
Las arquitecturas de NVIDIA palían este problema adoptando una arquitectura híbrida, en las que en cada SM las instrucciones y los datos son almacenados en cachés diferentes (L0 y L1). Las instrucciones son enviadas a los núcleos por la unidad de envío (dispatch unit) y los datos a los registros por las unidades de carga y descarga. El procedimiento de prefetching envía los datos (p.e. los pesos sináptico de un modelo de Aprendizaje Profundo) a los cachés y a los registros mucho antes de que sean utilizados.
GPU multi-instancia (MiG)
Figura 7: En la imagen, el chip se distribuye en tres GPUs virtuales. En la primera de ellas corren cuatro procesos (“compute instances”) programados en PyTorch y TensorFlow. Fuente: NVIDIA.
Los procesadores GA100 permiten la creación de hasta 7 GPUs virtuales completamente aisladas entre sí, de tal manera que los procesos que corren en una de ellas no afecten a los que se ejecutan en las demás. El principio de la virtualización es simple: Los diferentes SMs del chip así como el espacio de direcciones de memoria (recordemos que cada dirección es un número de 64 bits) se asignan a cada una de las instancias de GPU. A su vez, dentro de estas últimas pueden correr procesos en paralelo, tal y como se muestra en la imagen.
Los enormes frutos de A100
Figura 8: A Stable Diffusion (entrenada con NVIDIA A100) se le pide que pinte la escena de la Navegación de San Brandán en la que pájaros comienzan a entonar salmos a la llegada de la embarcación. Éste es el resultado.
Las tarjetas A100 han tenido una relevancia extraordinaria en el desarrollo del Deep Learning, ya que fueron utilizadas para entrenar poderosos modelos de lenguaje, como GPT-2, GPT-3 (modelo subyacente a la primera versión del célebre interfaz ChatGPT), Claude y Mistral. Además, los revolucionarios sistemas de texto a imagen, como Midjourney, DALL-E 2, DALL-E 3 y Stable Diffusion. también utilizaron el chip para fijar el valor de sus pesos sinápticos.
