Google (II): TPUv2 y TPUv3
Rubén Rodríguez Abril
El éxito de Google en la creación del chip de su inferencia TPUv1 (cuyo rendimiento en modelos neurales se demostró muy superior al de las CPUs y GPUs contemporáneas), motivó el inicio de nuevos trabajos para el diseño de un supercomputador capaz de entrenar grandes modelos de lenguaje, así como un ecosistema completo de herramientas.
El entrenamiento es un proceso computacionalmente mucho más intensivo que la inferencia. Centenares de ejemplos, integrados en un lote, se hacen pasar paralelamente a través de todas las capas del modelo hasta de llegar hasta la salida. Allí, se calculan la función de pérdida y su gradiente. Éste último es retropropagado a través de toda la red y su valor local para una determinada unidad/neurona es utilizado para adaptar los pesos de la función de pérdida. Es necesario una mayor capacidad de paralelismo, de cómputo, mayor programabilidad y una nueva aritmética en coma flotante que supere las limitaciones de los 8 bits que afectan a TPUv1.
Para atender a estas nuevas necesidades, los ingenieros de Google presentaron sus nuevos chips TPUv2 y TPUv3, diseñados específicamente para trabajar en paralelo en centros de datos integrados por centenares de nodos. Los reto a los que se enfrentaron los ingenieros durante el diseño fueron los siguientes:
-Creación de una infraestructura de red que permita la creación de un supercomputador a partir de múltiples chips/nodos conectados entre sí.
-Modificaciones en la arquitectura TPU de los chips.
-Un nuevo compilador (XLA).
-Un nuevo tipo de aritmética en coma flotante (Bfloat16).
Topología de la red
El supercomputador está compuesto de 256 (16×16) chips TPUv2 integrados en una estructura toroidal denominada pod, en el modo que se muestra en la imagen. Los chips se comunican entre sí a través de interconexiones ICI de gran ancho de banda. Cada chip tiene cuatro interconexiones ICI que transmiten información a un ratio de 496Gbits/s. En el caso de TPUv3, el número de chips de la red asciende a 1024.
Dado que los lotes de entrenamiento se componen de centenares de ejemplos, los más práctico es, durante el entrenamiento distribuir cada uno de estos últimos entre los diferentes chips (sharding), que se encargarán de realizar los cálculos de activación de las neuronas (a medida que la información progresa hacia adelante) y de retropropagación del gradiente de la función de pérdida.
Figura 1. Topología toroidal de un pod. Cada nodo de la red se ocupa del cálculo de una capa o de una parte de la msima. La información relativa a los nuevos pesos sinápticos es transmitida por cada nodo sus vecinos a través de conexiones y routers ICI. Fuente: Google.
Los diferentes nodos/chips trabajan individualmente y sólo necesitan comunicarse para realizar operaciones colectivas (allreduce), que involucren a todo el sistema, como el cálculo del valor medio de una capa (necesario para una normalización por lotes) o el cómputo de la actualización de los pesos sinápticos. Por lo general, cada vez que hay que realizar una operación global, los nodos se intercambian información a través de las interconexiones, y luego cada uno de los nodos procede a realizar la operación con la información obtenida del resto.
El entrenamiento se realiza, pues, de un modo síncrono, y todos los nodos actualizan sus pesos sinápticos a la vez. Un mecanismo de refrigeración (por aire, en el caso de TPUv2, y líquida, en el caso de TPUv3) asegura que la temperatura de la red se mantenga a un nivel constante.
Arquitectura del chip
Figura 2. Cada chip se compone de dos núcleos como los señalados en esta imagen (no confundir con los Tensor Cores de NVIDIA). Los núcleos se componen de un secuenciador, que distribuye las instrucciones, una unidad vectorial y otra unidad que realiza operaciones de trasposición matricial y permutación de componentes de vectores. A la izquierda, una memoria HBM de 8/16 GiB de capacidad y una conexión PCIe con la CPU y la placa base. A la derecha, un router de interconexión, que lo comunica con otros nodos de la red. Fuente: AMD.
Cada chip se compone de dos núcleos. Y cada núcleo incorpora las siguientes secciones principales (y otras secundarias):
1- Routers ICI, que se utilizan para enrutar información entre chips, a lo largo de la red toroidal, en los términos reseñados por los párrafos anteriores.
2- Memoria de banda ancha (HBM), situada fuera de la TPU pero unida físicamente a ella a través de un circuito integrado, denominado interposer, y que contiene 32 pistas de 128 bits. Los chips de memoria DRAM se acumulan en cuatro pilas diferentes. El acceso a ellos es regulado por un controlador en el modo señalado en la imagen.
Figura 3. A la derecha, varios chips de memoria DRAM se apilan par formar una memoria HBM, que se comunica con el procesador a través de un “interposer”. Fuente: AMD.
3- Secuenciador del núcleo (Core sequencer). Es el circuito de control más importante, encargado de extraer las instrucciones almacenadas previamente por el software y el compilador en una memoria específica del chip (Imem). Las instrucciones son del tipo VLIW (Very Long Instruction Word). Consisten en una cadena de nada menos que 322 bits que puede lanzar 8 operaciones aritméticas diferentes: dos escalares, dos vectoriales, carga y descarga de vector, carga y descarga de datos en la cola de la unidad matricial.
4- Unidad de procesamiento vectorial (Vectorial processing unit). Realiza operaciones sobre vectores. Dispone de una memoria (Vmem) capaz de almacenar 32 x 128 vectores de 32 bits, así como 32 registros de matrices de dimensión 128 x 8, con coeficientes de 32 bits. Esta unidad se encarga del cálculo de funciones de activación, así como de las operaciones de normalización tanto por lotes como por capas
5- Unidad matricial (MXU). Se trata de un array sistólico de 128×128, que tiene el mismo modo de operación que en el caso de TPUv1. Su aritmética es de semiprecisión: toma como factores a números de 16 bits y acumula resultados de 32 bits.
Los circuitos de control son más bien escasos y el chip no es capaz de gestionar la ejecución de varios hilos a la vez, al contrario que los chips de NVIDIA. El paralelismo en las TPUs es representado por la unidad vectorial, que procesa vectores de dimensión 128×8 (paralelismo a nivel de datos) y es capaz de lanzar 8 instrucciones aritméticas a la vez (paralelismo a nivel de instrucciones).
Una nueva aritmética
Figura 4. Comparación entre los formatos FP32, FP16 y BF16. Las operaciones que toman como entrada números FP16 y producen números FP32 se denominan “de semiprecisión”.
Aunque los creadores de TPUv2 tenían la intención originaria de que las operaciones aritméticas del chip (particularmente las multiplicaciones matriciales) se realizasen en formato de semiprecisión, pronto se dieron cuenta que un exponente de tan sólo 5 bits generaba a menudo errores de desbordamiento. Por este motivo, se decidió crear un nuevo tipo aritmético, BP16, en el que el tamaño del exponente se incrementaba en 3 bits, a expensas de la mantisa. La precisión de los resultados (que se asocia con la mantisa) apenas se vio afectada.
Rendimiento de TPUv2 y TPUv3
Las TPUs v2 y v3 de Google ofrecen un rendimiento notable en tareas de aprendizaje profundo. La TPUv2 tiene una capacidad de procesamiento de hasta 45 TeraFLOPs por núcleo y 11.5 PetaFLOPs en un pod completo. En comparación, la TPUv3 duplica el rendimiento por núcleo a 90 TeraFLOPs, alcanzando más de 100 PetaFLOPs en un pod completo. TPUv2 puede entrenar al modelo ResNet-50 en 2-4 horas y a GPT-2 en unos pocos días. TPUv3 puede completar estas tareas en menos de una hora y en menos de un día, respectivamente.
En comparación con la arquitectura Volta de NVIDIA, específicamente la V100, las TPUs de Google están diseñadas con un enfoque específico para el entrenamiento de redes neuronales profundas a gran escala. Aunque la V100 es más versátil, siendo adecuada para una variedad más amplia de cargas de trabajo, las TPUs destacan en tareas optimizadas para TensorFlow, lo que les permite un rendimiento superior en escenarios específicos de Machine Learning.
Lecturas Recomendadas
