NVIDIA (II): la arquitectura Volta (V100)

CATEGORÍA

Rubén Rodríguez Abril

La arquitectura Volta de NVIDIA, presentada en 2017, revolucionó el diseño de las GPUs con la introducción de los Tensor Cores, áreas del chip especializadas en la multiplicación matricial. Si bien las GPUs existen desde hace décadas, no fue hasta la década de 2010 que el aprendizaje profundo comenzó a influir en el diseño de los chips de NVIDIA, desde las arquitecturas Kepler (2012) y Volta (2017) en adelante. Este artículo ofrece una visión personal sobre los aspectos más destacados de Volta.

Aunque en la actualidad, el público general (y los inversores) asocian las GPUs con el desarrollo de la IA, en realidad el contacto entre ambos es relativamente reciente. De hecho, se remonta al trabajo seminal “Large-scale Deep Unsupervised Learning using Graphics Processors” del año 2009, que por primera vez puso de manifiesto los beneficios de entrenar modelos de Deep Learning en GPUs, en términos de escalabilidad y tiempo de aprendizaje.

Por este motivo, aunque las tarjetas gráficas tienen décadas de existencia, no es hasta la década de los 2010s cuando el aprendizaje profundo comenzó a influenciar el diseño de los chips, particularmente en las arquitecturas Volta y Ampère, que son el objeto del presente y el siguiente artículo, y cuyas tarjetas V100 y A100 han sido el soporte silíceo de los revolucionarios modelos de lenguaje surgidos en la actualidad (el autor escribe en Agosto de 2024).

Los siguientes párrafos no pretenden ser un white paper exhaustivo de la arquitectura Volta, sino una visión personal de la misma ofrecida por el autor, destacando los aspectos que le parecen más llamativos.

La Arquitectura Volta: la llegada de los Tensor Cores

Cuando me piden una definición breve de Deep Learning, a mí me gusta ofrecer la siguiente: “Variedad del Aprendizaje de Máquina, ligeramente inspirada en el sistema nervioso de los animales, y consistente en la aplicación sucesiva de la multiplicación matricial sobre los datos de entrada”.

Influenciados por el trabajo de McCulloch y de Pitts, los primeros especialistas en Aprendizaje Profundo solían concebir sus modelos en términos de neuronas artificiales (siendo cada una de ellas una concatenación de dos funciones, excitación y activación). Sin embargo, la llegada de los redes convolucionales y de los transformers sugirió considerar a las capas densamente conectadas, a las capas convolucionales o a las unidades de atención como enormes matrices que se multiplicaban por los datos de activación de la capa anterior conforme la información avanzaba hacia la salida.

La multiplicación matricial se reveló, pues, como la principal operación de los modelos, la que más recursos computacionales consumía y la que les proporcionaba buena parte de su potencia. Para su perfección de un modo lo más eficiente posible, NVIDIA introdujo un nuevo tipo de sección en sus chips: Los núcleos tensoriales o Tensor Cores (un tensor es una generalización del concepto de matriz a varias dimensiones).

La multiplicación matricial

En una multiplicación matricial, dos matrices factores, una de dimensión m x k y otra de dimensión k x n se combinan para dar lugar a una matriz-producto de tamaño m x n. La matriz-producto tiene, pues, la misma cantidad de filas (m) que la primera y la misma cantidad de columnas (n) que la segunda. La dimensión interna de la multiplicación matricial, k, es denominada profundidad, y es de gran importancia en el ámbito de las GPUs, ya que determina el número de pasos aritméticos que hay que realizar para calcular cada elemento de la nueva matriz (véase en la imagen).

Figura 1: La multiplicación matricial puede ser considerada concatenación de productos escalares de vectores-fila con vectores-columnas. En la ilustración, las matrices A y B se multiplican para dar lugar a la matriz-producto C. Para calcular un elemento de esta última se debe realizar el producto escalar del correspondiente vector-fila de A con el correspondiente vector-columna de B. Ambos vectores deben de tener la misma longitud, equivalente a la profundidad de la multiplicación matricial (en este caso, equivalente a 7).

Para facilitar el cálculo de productos escalares (y por lo tanto, matriciales), la arquitectura Kepler (año 2012) introdujo la instrucción FMA (Fused Multiplication and Addition), que multiplica dos factores y sumaba el resultado a un acumulador. Es ejecutada en los núcleos CUDA.

Sin embargo, pronto los equipos de diseño de NVIDIA consideraron conveniente la creación de secciones de chip especializadas en la multiplicación matricial. Con este objetivo, la arquitectura Volta introdujo por primera vez los Tensor Cores, que operan con arreglo a la siguiente ecuación lineal:

D = AB + C

donde D es la matriz resultado, C es la matriz acumulada (puede ser igual a cero), y A y B las matrices que se están multiplicando. El formato numérico es el de semiprecisión: A y B almacenan números de 16 bits; C y D, de 32 bits.

En cada ciclo del reloj, un Tensor Core de Volta puede calcular una matriz producto de tamaño 4×4, y con una profundidad de 4. Si la matriz es de tamaño superior, pueden utilizarse paralelamente varios Tensor Cores, uno de ellos para cada sección de 4×4. Si la profundidad de la multiplicación excede de 4, es necesario utilizar varios ciclos del reloj, tal y como se muestra en la imagen.

Figura 2: Dos matrices, una de 4×8 (en turquesa) y otra de 8×8 son multiplicados en un SM de Volta. La dimensión de profundidad viene representada verticalmente. Dado que la matriz-producto tiene un tamaño de 8×4, es necesario utilizar dos Tensor Cores. Y como la profundidad es de 8, necesitamos emplear dos ciclos de reloj. Los cuadrados en gris representan multiplicaciones, que son sumadas verticalmente, agregándose la suma a los acumuladores en verde (representan los elementos de la matriz D). En el siguiente ciclo del reloj, la matriz D anterior se convierte en la nueva matriz C.

Tensor Cores vs núcleos CUDA

Tanto si se realiza en los Tensor Cores como los núcleos CUDA, la multiplicación de dos matrices 4×4 abarca el mismo número de operaciones aritméticas (64 sumas-multiplicaciones). La razón de la mayor eficiencia de los primeros radica en que en un Tensor Core es necesario invocar una única instrucción FMA, frente a las 64 que serían necesarias en un núcleo CUDA, con la consiguiente sobrecarga de los circuitos de control, que deben extraer las instrucciones del caché L0, descodificarlas y enviarlas a las unidades aritméticas para su ejecución.

En la arquitectura Volta, cada vez que se va realizar una multiplicación matricial involucra los pasos siguientes:

-Los coeficientes de las matrices son transportados desde L1 a los registros del hilo (hay 255 por cada uno de ellos).

-Una instrucción de código máquina FMA ordena al Tensor Core a realizar una multiplicación matricial. Los argumentos de la misma son los registros donde se guardan las matrices A, B, C y D.

En pseudocódigo, la sintaxis de esta última instrucción sería:

HMMA Rdest, Rsrc1, Rsrc2, Racc

donde HHMA significa Half-Precision Matrix Multiply-Accumulate, mientras que los argumentos señalan respectivamente los registros de destino, los de las dos matrices-fuente y los de la matriz acumuladora.

Los SMs de la arquitectura Volta: operaciones aritméticas

En Volta, cada SM se divide en cuatro secciones (quads, cuadrantes). Cada cuadrante tiene una estructura interna muy parecida a la de la arquitectura Fermi (con la salvedad delos Tensor Cores), como se puede comprobar en esta imagen:

Figura 3. Un SM se divide en cuatro sectores idénticos. Cada uno de ellos tiene la misma estructura que en el caso de la arquitectura Fermi (analizada en el artículo anterior), pero con un par de salvedades: Aparecen los Tensor Cores y hay unidades (CUDA cores) dedicadas a la aritmética FP64 (coma flotante en 64 bits). En el exterior hay cuatro unidades de textura, que no tienen demasiada relevancia en el ámbito del Deep Learning. Fuente: NVIDIA.

Existe una cierta correspondencia entre las unidades aritméticas que integran un SM y las diferentes operaciones que se ejecutan en los modelos de Deep Learning:

-Los Tensor Cores están especializados en la multiplicación matricial, y por ello se utilizan para calcular el paso de la información a través capas densamente conectadas (como las mini-redes orientadas hacia adelante de los transformers), capas convolucionales (en visión artificial) o unidades de atención.

-Los núcleos CUDA son utilizados fundamentalmente para llevar a cabo operaciones elemento a elemento, como por ejemplo, los productos de Hadamard y de Frobenius (usados en la transferencia neuronal de estilo), las procesos de incremento o decremento de resolución en Visión Artificial o la propia función ReLU.

-Las unidades de funciones especiales (SFU) calculan funciones trascendentes, como las logarítmicas, las exponenciales o las trigonométricas. Son utilizadas para calcular, entre otras, la función logística (muy utilizada como función de activación en los años 90), la función GeLU, o la función softmax, que tanta importancia tiene en la clasificación de imágenes o los modelos de procesamiento de lenguaje natural.

Usos de la arquitectura Volta en Deep Learning

Figura 4: Un acelerador (“tarjeta gráfica”) V100. La GPU (GV100) se encuentra en el centro. A su alrededor se sitúan chips de memoria HBM2, los sistemas de refrigeración y alimentación, así como el interfaz que conecta con la placa base (PCIe) o con otras GPUs (NVLink 2.0).

Los aceleradores V100 han tenido una extraordinaria importancia en el desarrollo de los modelos de lenguaje de OpenAI, pues fueron utilizados en el entrenamiento de GPT-2 y GPT-3. Además, los mismos sirvieron de base para el aprendizaje de DALL-E, un modelo derivado de GPT-3 de que se hizo sumamente popular en el gran público, porque era capaz de crear gráficos con cierta calidad artística a partir de una entrada en texto proporcionada por el usuario.