Arquitectura de TPU
Las unidades de procesamiento de tensor (TPU) son circuitos integrados para aplicaciones específicas (ASIC) diseñados por Google para acelerar las cargas de trabajo de aprendizaje automático. TPU de Cloud es un Google Cloud servicio que pone las TPUs a disposición de los usuarios como un recurso escalable.
Las TPUs están diseñadas para realizar operaciones de matriz rápidamente, por lo que son ideales para cargas de trabajo de aprendizaje automático. Puedes ejecutar cargas de trabajo de aprendizaje automático en TPUs mediante frameworks como PyTorch y JAX.
¿Cómo funcionan las TPUs?
Para entender cómo funcionan las TPUs, es útil saber cómo abordan otros aceleradores los retos computacionales del entrenamiento de modelos de aprendizaje automático.
Cómo funciona una CPU
Una CPU es un procesador de uso general basado en la arquitectura de Von Neumann. Esto significa que una CPU funciona con software y memoria de la siguiente manera:
La mayor ventaja de las CPUs es su flexibilidad. Puedes cargar cualquier tipo de software en una CPU para muchos tipos de aplicaciones. Por ejemplo, puedes usar una CPU para procesar textos en un ordenador, controlar motores de cohetes, ejecutar transacciones bancarias o clasificar imágenes con una red neuronal.
Una CPU carga valores de la memoria, realiza un cálculo con los valores y almacena el resultado en la memoria para cada cálculo. El acceso a la memoria es lento en comparación con la velocidad de cálculo y puede limitar el rendimiento total de las CPUs. A menudo se denomina cuello de botella de von Neumann.
Cómo funciona una GPU
Para conseguir un mayor rendimiento, las GPUs contienen miles de unidades aritméticas y lógicas (ALUs) en un solo procesador. Una GPU moderna suele contener entre 2500 y 5000 ALUs. El gran número de procesadores significa que puedes ejecutar miles de multiplicaciones y sumas simultáneamente.
Esta arquitectura de GPU funciona bien en aplicaciones con paralelismo masivo, como las operaciones de matrices en una red neuronal. De hecho, en una carga de trabajo de entrenamiento típica para el aprendizaje profundo, una GPU puede proporcionar un rendimiento superior en un orden de magnitud al de una CPU.
Sin embargo, la GPU sigue siendo un procesador de uso general que tiene que admitir muchas aplicaciones y software diferentes. Por lo tanto, las GPUs tienen el mismo problema que las CPUs. Para cada cálculo en los miles de ALUs, una GPU debe acceder a los registros o a la memoria compartida para leer los operandos y almacenar los resultados de los cálculos intermedios.
Cómo funciona una TPU
Google ha diseñado las TPUs de Cloud como un procesador de matrices especializado en cargas de trabajo de redes neuronales. Las TPUs no pueden ejecutar procesadores de texto, controlar motores de cohetes ni realizar transacciones bancarias, pero sí pueden gestionar operaciones de matrices masivas que se usan en redes neuronales a gran velocidad.
La tarea principal de las TPU es el procesamiento de matrices, que es una combinación de operaciones de multiplicación y acumulación. Las TPUs contienen miles de multiplicadores de acumulación que están conectados directamente entre sí para formar una matriz física de gran tamaño. Esto se denomina arquitectura de matriz sistólica. Cloud TPU v3 contiene dos matrices sistólicas de 128 x 128 ALUs en un solo procesador.
El host de TPU transmite datos a una cola de entrada. La TPU carga datos de la cola de entrada y los almacena en la memoria HBM. Cuando se completa el cálculo, la TPU carga los resultados en la cola de outfeed. A continuación, el host de la TPU lee los resultados de la cola de salida y los almacena en la memoria del host.
Para realizar las operaciones de matriz, la TPU carga los parámetros de la memoria HBM en la unidad de multiplicación de matriz (MXU).
A continuación, la TPU carga los datos de la memoria HBM. A medida que se ejecuta cada multiplicación, el resultado se pasa al siguiente multiplicador-acumulador. El resultado es la suma de todos los resultados de la multiplicación entre los datos y los parámetros. No se requiere acceso a la memoria durante el proceso de multiplicación de matrices.
Por lo tanto, las TPUs pueden alcanzar un alto rendimiento computacional en los cálculos de redes neuronales.
Arquitectura del sistema de TPU
En las siguientes secciones se describen los conceptos clave de un sistema de TPU. Para obtener más información sobre los términos habituales del aprendizaje automático, consulta el glosario de aprendizaje automático.
Si no has usado nunca las TPU de Cloud, consulta la página principal de la documentación de las TPU.
Chip TPU
Un chip de TPU contiene uno o varios TensorCores. El número de Tensor Cores depende de la versión del chip de TPU. Cada Tensor Core consta de una o varias unidades de multiplicación de matrices (MXUs), una unidad vectorial y una unidad escalar. Para obtener más información sobre los Tensor Cores, consulta el artículo A Domain-Specific Supercomputer for Training Deep Neural Networks (Superordenador específico de un dominio para entrenar redes neuronales profundas).
Una MXU se compone de 256 x 256 (TPU v6e) o 128 x 128 (versiones de TPU anteriores a la v6e) multiplicadores de acumulación en una matriz sistólica. Las MXUs proporcionan la mayor parte de la potencia de cálculo de un Tensor Core. Cada MXU puede realizar 16.000 operaciones de multiplicación y acumulación por ciclo. Todas las multiplicaciones usan entradas bfloat16, pero todas las acumulaciones se realizan en formato de número FP32.
La unidad vectorial se usa para cálculos generales, como activaciones y softmax. La unidad escalar se usa para el flujo de control, el cálculo de direcciones de memoria y otras operaciones de mantenimiento.
pod de TPU
Un pod de TPUs es un conjunto contiguo de TPUs agrupadas en una red especializada. El número de chips de TPU de un pod de TPU depende de la versión de la TPU.
Segmento
Un slice es un conjunto de chips que se encuentran en el mismo pod de TPU y están conectados mediante interconexiones de alta velocidad entre chips (ICI). Las porciones se describen en términos de chips o Tensor Cores, según la versión de la TPU.
Forma del chip y topología del chip también hacen referencia a las formas de las porciones.
Multicorte frente a corte único
Multislice es un grupo de slices que amplía la conectividad de las TPU más allá de las conexiones de interconexión entre chips (ICI) y aprovecha la red de centros de datos (DCN) para transmitir datos más allá de un slice. ICI sigue transmitiendo los datos de cada segmento. Gracias a esta conectividad híbrida, Multislice permite el paralelismo entre slices y te permite usar un mayor número de núcleos de TPU para un solo trabajo que el que puede admitir un solo slice.
Las TPUs se pueden usar para ejecutar un trabajo en una sola porción o en varias. Consulta la introducción a Multislice para obtener más información.
Cubo de TPU
Una topología de 4x4x4 de chips de TPU interconectados. Esto solo se aplica a las topologías 3D (a partir de la TPU v4).
SparseCore
Los SparseCores son procesadores de flujo de datos que aceleran los modelos mediante operaciones dispersas. Uno de los principales casos prácticos es la aceleración de modelos de recomendación, que dependen en gran medida de las incrustaciones. La versión 5p incluye cuatro SparseCores por chip, mientras que la versión 6e incluye dos SparseCores por chip. Para obtener una explicación detallada sobre cómo se pueden usar los SparseCores, consulta Análisis detallado de SparseCore para modelos de inserciones grandes (LEM). Puedes controlar cómo usa el compilador de XLA SparseCores mediante marcas de XLA. Para obtener más información, consulta Flags de XLA de TPU.
Resistencia de ICI de TPU de Cloud
La resiliencia de ICI ayuda a mejorar la tolerancia a fallos de los enlaces ópticos y los conmutadores de circuitos ópticos (OCS) que conectan las TPUs entre cubos. Las conexiones ICI de un cubo usan enlaces de cobre que no se ven afectados. La resiliencia de ICI permite que las conexiones ICI se enruten por OCS y por fallos ópticos de ICI. Como resultado, mejora la disponibilidad de programación de las porciones de TPU, pero se produce una degradación temporal del rendimiento de ICI.
En las versiones 4 y 5p de TPU de Cloud, la resiliencia de ICI está habilitada de forma predeterminada en los slices que tienen un cubo o más. Por ejemplo:
- v5p-128 al especificar el tipo de acelerador
- 4x4x4 al especificar la configuración del acelerador
Versiones de TPU
La arquitectura exacta de un chip de TPU depende de la versión de TPU que utilices. Cada versión de TPU también admite diferentes tamaños y configuraciones de slice. Para obtener más información sobre la arquitectura del sistema y las configuraciones admitidas, consulta las siguientes páginas:
Arquitectura de nube de TPU
Google Cloud ofrece TPUs como recursos de computación a través de máquinas virtuales de TPU. Puedes usar directamente las VMs de TPU para tus cargas de trabajo o usarlas a través de Google Kubernetes Engine o Vertex AI. En las siguientes secciones se describen los componentes clave de la arquitectura de la nube de TPU.
Arquitectura de las máquinas virtuales de TPU
La arquitectura de máquina virtual de TPU te permite conectarte directamente a la máquina virtual conectada físicamente al dispositivo de TPU mediante SSH. Una VM de TPU, también denominada "trabajador", es una máquina virtual que ejecuta Linux y tiene acceso a las TPUs subyacentes. Tienes acceso root a la VM, por lo que puedes ejecutar código arbitrario. Puedes acceder a los registros de depuración y a los mensajes de error del compilador y del tiempo de ejecución.
Un único host, varios hosts y subhosts
Un host de TPU es una máquina virtual que se ejecuta en un ordenador físico conectado al hardware de TPU. Las cargas de trabajo de TPU pueden usar uno o varios hosts.
Una carga de trabajo de un solo host está limitada a una VM de TPU. Una carga de trabajo de varios hosts distribuye el entrenamiento en varias VMs de TPU. Una carga de trabajo de subhost no usa todos los chips de una VM de TPU.
Arquitectura de nodos de TPU (obsoleta)
La arquitectura de nodos de TPU consta de una VM de usuario que se comunica con el host de la TPU a través de gRPC. Si usas esta arquitectura, no podrás acceder directamente al host de la TPU, lo que dificulta la depuración de errores de entrenamiento y de TPU.
Cambiar de la arquitectura de nodos de TPU a la de máquinas virtuales de TPU
Si tienes TPUs que usan la arquitectura de nodos de TPU, sigue estos pasos para identificarlas, eliminarlas y volver a aprovisionarlas como VMs de TPU.
Ve a la página TPUs:
Busca tu TPU y su arquitectura en el encabezado Arquitectura. Si la arquitectura es "Máquina virtual de TPU", no tienes que hacer nada. Si la arquitectura es "Nodo de TPU", debes eliminar y volver a aprovisionar la TPU.
Elimina la TPU y vuelve a aprovisionarla.
Consulta Gestionar TPUs para obtener instrucciones sobre cómo eliminar y volver a aprovisionar TPUs.