Google (I): TPUv1

CATEGORÍA

Rubén Rodríguez Abril

TPUs (Tensor Processing Units, Unidades de Procesamiento Tensorial), es el nombre con el que Google ha designado a una serie de chips especializados en realizar tareas de aprendizaje profundo. En este artículo presentamos su primera versión, TPUv1, cuyo rendimiento era del orden a 15-20 veces superior al de las CPUs y GPUs de su generación

Google ha realizado importantes contribuciones al desarrollo reciente de la Inteligencia Artificial. En el ámbito de los grandes modelos de lenguaje, introdujo la arquitectura de transformers. En el plano del hardware, sus ingenieros diseñaron el primer ASIC (chip específico) especializado en cálculos de redes neuronales artificiales, y dotado de una unidad especifica para operaciones matriciales. Nos estamos refiriendo a los célebres TPUs (Tensor Processing Units) de Google, accesibles al programador medio y amateur gracias a los servicios de Google Cloud.

La primera versión de las serie de chips TPU, TPUv1, fue presentado al público en la edición anual de ISCA que tuvo lugar en Toronto en Junio de 2017. La segunda versión, TPUv2, está especializada en tareas de entrenamiento. La primera será analizada en el presente artículo.

Arquitectura de TPUv1

TPUv1 tiene analogías con GPU y con FPU. Al igual que los procesadores gráficos, puede realizar una enorme cantidad de cálculos en paralelo y su conexión con la CPU se realiza a través de un bus PCIe, y de una ranura SATA de la placa base, lo que facilita su incorporación a cualquier sistema (basta su mero enchufe a la ranura).

Sin embargo, las GPUs (al menos, las de NVIDIA) reciben las instrucciones en bloque desde la CPU, las guardan en su propio memoria DRAM, y desde allí son extraídas (fetching) por sus propios circuitos de control, y colocadas en la cola para su ejecución. Las TPUs no son tan autónomas. No disponen de una memoria RAM propia. La operación de fetching es realizada por la propia CPU, que extrae las instrucciones desde la memoria de la placa base y las envía al búfer/cola de la TPU. Lo cual les permite ahorrar espacio en circuitos de control. Y en ello, son más parecidas a una FPU.

En TPU, no hay unidades de control que gestionen procesos computacionales en paralelo (los hilos de las GPUs de NVIDIA). No hay unidades de carga y descarga de datos, ni tampoco circuitos especializados en la gestión de instrucciones de ramificación (que son más bien escasas en los modelos de Aprendizaje Profundo, que son deterministas). Todo ello es lo que posibilita que los circuitos de control se reduzcan al mínimo, y su área pueda ser asignada a circuitos aritméticos o de memoria.

Por lo demás, el chip TPUv1 fue fabricado a 28 nms, corre a 700 mHZ y tiene un consumo de 40 vatios.

Figura 1: A la izquierda, diagrama de una TPU. A la derecha descripción del área física que ocupa cada unidad. Las unidades aritméticas se dibujan en amarillo, las de entrada y salida en verde, y las de memoria en azul. Fuente: Google.

En la TPU, los circuitos de control se reducen al mínimo. Así que en el análisis del chip distinguiremos entre tipos de áreas: aritméticas (o de computación), de memoria y de entrada y salida.

Unidades aritméticas

La Unidad de Multiplicación Matricial (MXU) realiza, como su propio nombre indica, cálculos de matrices. Es capaz de computar 65.536 (256×256) operaciones de multiplicación-acumulación (esenciales en el cálculo de productos escalares entre dos vectores y de productos vectoriales) por ciclo del reloj, en aritmética de 8 bits.

Los diferentes módulos de esta matriz forman un array sistólico: un retículo cuatrado de 256×256 unidades de multiplicación-acumulación, cada una de las cuales, en cada ciclo del reloj, realiza una multiplicación de dos elementos, suma el resultado con el valor de la casilla de arriba, y pasa los resultados a la casilla de abajo. Los factores de cada multiplicación son dos números de 8 bits, pero el resultado es de 16 bits.

Figura 2. A la izquierda, se describe la multiplicación ordinaria de un vector (en rojo oscuro) por una matriz cuadrada. El vector avanza hacia la derecha, y en cada momento, se realiza un producto escalar con el correspondiente vector columna. Los coeficientes se multiplican y luego se suman. Los resultados de cada suma son almacenados en el vector-resultado colocado verticalmente. A la derecha se describe cómo funciona el procedimiento en un array sistólico. Los coeficientes de la matriz han sido precargados. Los vectores avanzan en diagonal. Cada vez que atraviesan una celdilla, ésta realiza la multiplicación del coeficiente vectorial con el matricial, y pasa el resultado a la celda de abajo. El resultado, es que en el momento 8, un producto escalar del vector amarillo y dos del vector rojo han sido almacenados en los vectores-resultado.

La unidad de multiplicación matricial está concebida para implementar las capas lineales de las redes neuronales artificiales. En este esquema, los pesos sinápticos son precargados en la matriz procedentes de un búfer FIFO (una cola). Tras ello, los sucesivos vectores de entrada entran en la matriz procedentes del búfer unificado Los resultados se guardan en los acumuladores.

La unidad de activación aplica funciones de esta categoría (ReLU, sigmoide, etc.) a los datos del acumulador. También puede realizar operaciones de normalización (por capas), incremento y reducción (pooling) de resolución. Su salida es almacenada en el búfer unificado.

Áreas de memoria

Los acumuladores son las pequeñas unidades que guardan los resultados obtenidos por cada una de las celdillas. Están situados físicamente, dentro del chip, debajo de la unidad matricial. El búfer unificado (de memoria SRAM) guarda los resultados de salida de cada capa de la red neuronal, que sirven a su vez de entrada para la capa siguiente.

Por encima del chip, pero fuera del mismo, se sitúa una memoria DRAM DDR3, en la que se guardan los pesos sinápticos del modelo. Dichos pesos son transferidos progresivamente al búfer de pesos, a medida que la información va atravesando las diferentes capas de la arquitectura.

No hay cachés de memoria SRAM (a diferencia de lo que sucede en chips de otras marcas como NVIDIA) lo cual supone un gran ahorro en potencia.

Áreas de entrada y salida

Son el bus PCIe, que comunica con la CPU y la placa base, y los puertos DD3, que lo hacen con la memoria externa.

Ejecución de instrucciones

El juego de instrucciones ISA de una TPU se compone de doce instrucciones, de las cuales podemos destacar las siguientes:

1. Read_Host_Memory, transfiere datos desde la memoria de la placa base al búfer unificado.

2. Read_Weights copia los pesos sinápticos de una capa desde la memoria DDR3 al búfer de pesos.

3. MatrixMultiply/Convolve realiza una operación de multiplicacion matricial. Los pesos sinápticos se precargan en la matriz procedentes de su búfer y tras ello. El resultado es guardado en el acumulador.

4. Activate aplica una función de activación a los datos del acumulador y guarda el resultado en el búfer unificado.

5. Write_Host_Memory transmite datos desde la TPU (búfer unificado) a la memoria de la CPU.

Al inicio de cada fase de inferencia, se ejecuta la instrucción 1, para transferir datos de entrada y pesos sinápticos a la TPU. Después, se van repitiendo sucesivamente, en bucle, las instrucciones 2 a 4, en cada una de las capas del modelo. Cuando se llega a la salida, los resultados se transfieren a la CPU mediante la instrucción 5.

Conclusiones

Gracias a la reducción de sus circuitos y la supresión de los cachés el rendimiento de la TPU, durante la ejecución de tareas de inferencia de un modelo de aprendizaje profundo, era del orden de 15 a 30 veces superior que el de las CPUs y las GPUs. Tratándose del rendimiento por vatio los resultados eran aun mejores.

El siguiente modelo de TPU, TPUv2, está enfocado no tanto en la inferencia como en el entrenamiento de las arquitecturas de Aprendizaje Profundo. A él dedicaremos nuestro próximo artículo en esta serie.