Arquitectura de las TPU
Las unidades de procesamiento tensorial (TPU) son circuitos integrados específicos de aplicaciones (ASIC) diseñados por Google para acelerar las cargas de trabajo de aprendizaje automático. Cloud TPU es un Google Cloud servicio que pone las TPU a disposición como un recurso escalable.
Las TPU están diseñadas para realizar operaciones de matrices con rapidez, lo que las hace ideales para las cargas de trabajo de aprendizaje automático. Puedes ejecutar cargas de trabajo de aprendizaje automático en TPU con frameworks como Pytorch y JAX.
¿Cómo funcionan las TPU?
Para comprender cómo funcionan las TPU, es útil entender cómo otros aceleradores abordan los desafíos computacionales de la capacitación de modelos de AA.
Cómo funciona una CPU
La CPU es un procesador de propósito general basado en la arquitectura de von Neumann. Eso significa que una CPU funciona con software y memoria, de la siguiente manera:
El mayor beneficio de las CPUs es su flexibilidad. Puedes cargar cualquier tipo de software en una CPU para muchos tipos de aplicaciones diferentes. Por ejemplo, puedes usar una CPU para procesar textos en una PC, 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 en los valores y vuelve a almacenar 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 la capacidad de procesamiento total de las CPUs. Esto se conoce a menudo como el cuello de botella de von Neumann.
Cómo funciona una GPU
Para aumentar la capacidad de procesamiento, las GPUs contienen miles de unidades lógicas aritméticas (ALU) en un solo procesador. Una GPU moderna suele contener entre 2,500 y 5,000 ALU. La gran cantidad de procesadores permite ejecutar miles de multiplicaciones y sumas de manera simultánea.
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 pedido de capacidad de procesamiento de magnitud mayor que una CPU.
Sin embargo, la GPU sigue siendo un procesador de propósito general que tiene que admitir muchas aplicaciones y software diferentes. Por lo tanto, las GPU presentan el mismo problema que las CPU. Para cada cálculo en los miles de ALU, una GPU debe acceder a los registros o a la memoria compartida para leer los operandos y almacenar los resultados del cálculo intermedio.
Cómo funciona una TPU
Google diseñó Cloud TPU como un procesador matricial especializado para cargas de trabajo de redes neuronales. Las TPU no pueden ejecutar procesadores de texto, controlar motores de cohetes ni ejecutar transacciones bancarias, pero pueden manejar operaciones de matrices masivas que se usan en redes neuronales a velocidades rápidas.
La tarea principal de las TPU es el procesamiento matricial, que es una combinación de operaciones de multiplicación y acumulación. Las TPU contienen miles de acumuladores multiplicadores que están conectados directamente entre sí para formar una gran matriz física. Esto se conoce como arquitectura de arreglo sistólico. Cloud TPU v3 contiene dos arreglos sistólicos de 128 x 128 ALU en un solo procesador.
El host de TPU transmite datos a una fila de infeed. La TPU carga datos de la fila de infeed y los almacena en la memoria HBM. Cuando se completa el procesamiento, la TPU carga los resultados en la cola de salida. Luego, el host de TPU lee los resultados de la cola de salida y los almacena en la memoria del host.
Para realizar las operaciones de matrices, la TPU carga los parámetros de la memoria HBM en la unidad de multiplicación de matrices (MXU).
Luego, la TPU carga datos de la memoria HBM. A medida que se ejecuta cada multiplicación, el resultado se pasa al siguiente acumulador de multiplicación. El resultado es la suma de todos los resultados de multiplicación entre los datos y los parámetros. No se requiere acceso a la memoria durante el proceso de multiplicación de matrices.
Como resultado, las TPU pueden lograr una alta capacidad de procesamiento 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 comunes del aprendizaje automático, consulta el Glosario de aprendizaje automático.
Si es la primera vez que usas Cloud TPU, consulta la página principal de la documentación de TPU.
Chip de TPU
Un chip de TPU contiene uno o más TensorCores. La cantidad de TensorCores depende de la versión del chip TPU. Cada TensorCore consta de una o más unidades de multiplicación de matrices (MXUs), una unidad vectorial y una unidad escalar. Para obtener más información sobre TensorCores, consulta Un supercomputador específico de dominio para el entrenamiento de redes neuronales profundas.
Una MXU se compone de acumuladores multiplicadores de 256 x 256 (TPU v6e) o 128 x 128 (versiones de TPU anteriores a la v6e) en un array sistólico. Las MXU proporcionan la mayor parte del poder de procesamiento en un TensorCore. Cada MXU puede realizar 16,000 operaciones de multiplicación y acumulación por ciclo. Todas las multiplicaciones toman entradas de bfloat16, pero todas las acumulaciones se realizan en formato de número FP32.
La unidad vectorial se usa para el procesamiento general, como las activaciones y el 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 TPU es un conjunto contiguo de TPU agrupadas en una red especializada. La cantidad de chips TPU en un pod de TPU depende de la versión de TPU.
Porción
Una porción es una colección de chips que se encuentran dentro del mismo pod de TPU y que están conectados por interconexiones entre chips (ICI) de alta velocidad. Las porciones se describen en términos de chips o TensorCores, según la versión de TPU.
Forma del chip y topología del chip también se refieren a las formas de las rebanadas.
Multislice frente a Single Slice
Multislice es un grupo de porciones que extiende la conectividad de las TPU más allá de las conexiones de interconexión entre chips (ICI) y aprovecha la red del centro de datos (DCN) para transmitir datos más allá de una porción. El ICI sigue transmitiendo los datos de cada fragmento. Con esta conectividad híbrida, Multislice habilita el paralelismo entre las porciones y te permite usar una mayor cantidad de núcleos de TPU para un solo trabajo de lo que puede admitir una sola porción.
Las TPU se pueden usar para ejecutar una tarea en una sola porción o en varias. Consulta la Introducción a la tomografía multislice para obtener más detalles.
Cubo de TPU
Una topología 4x4x4 de chips TPU interconectados. Esto solo se aplica a las topologías 3D (a partir de la TPU v4).
SparseCore
SparseCores son procesadores de flujo de datos que aceleran los modelos que se basan en incorporaciones que se encuentran en los modelos de recomendación. v5p incluye cuatro SparseCores por chip, y v6e incluye dos SparseCores por chip.
Resiliencia de ICI de Cloud TPU
La resiliencia de la ICI ayuda a mejorar la tolerancia a fallas de los vínculos ópticos y los conmutadores de circuitos ópticos (OCS) que conectan las TPU entre cubos. (Las conexiones de ICI dentro de un cubo usan vínculos de cobre que no se ven afectados). La resiliencia de ICI permite que las conexiones de ICI se enruten alrededor de las fallas de OCS y de ICI ópticas. Como resultado, mejora la disponibilidad de programación de las porciones de TPU, con la desventaja de una degradación temporal en el rendimiento de ICI.
En Cloud TPU v4 y v5p, la resiliencia de ICI está habilitada de forma predeterminada para las porciones que son de un cubo o más, por ejemplo:
- v5p-128 cuando se especifica el tipo de acelerador
- 4 × 4 × 4 cuando se especifica la configuración del acelerador
Versiones de TPU
La arquitectura exacta de un chip de TPU depende de la versión de TPU que uses. Cada versión de TPU también admite diferentes tamaños y configuraciones de rebanadas. Para obtener más información sobre la arquitectura del sistema y las configuraciones compatibles, consulta las siguientes páginas:
Arquitectura de nube de TPU
Google Cloud pone las TPU a disposición como recursos de procesamiento a través de VMs 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 la VM de TPU
La arquitectura de la VM de TPU te permite conectarte directamente a la VM conectada físicamente al dispositivo TPU mediante SSH. Una VM de TPU, también conocida como trabajador, es una máquina virtual que ejecuta Linux y tiene acceso a las TPU subyacentes. Tienes acceso raíz a la VM, por lo que puedes ejecutar código arbitrario. Puedes acceder a los registros de depuración y los mensajes de error del compilador y del entorno de ejecución.
Host único, multihost y subhost
Un host de TPU es una VM que se ejecuta en una computadora física conectada al hardware de TPU. Las cargas de trabajo de TPU pueden usar uno o más hosts.
Una carga de trabajo de host único se limita 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 del nodo TPU (obsoleto)
La arquitectura del nodo de TPU consta de una VM de usuario que se comunica con el host de TPU a través de gRPC. Cuando usas esta arquitectura, no puedes acceder directamente al host de TPU, lo que dificulta la depuración de errores de entrenamiento y TPU.
Cómo pasar de la arquitectura de nodo TPU a la de VM de TPU
Si tienes TPU que usan la arquitectura de nodos TPU, sigue los pasos que se indican a continuación para identificarlas, borrarlas 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 "VM de TPU", no es necesario que realices ninguna acción. Si la arquitectura es "Nodo de TPU", debes borrar la TPU y volver a aprovisionarla.
Borra la TPU y vuelve a aprovisionarla.
Consulta Administra TPUs para obtener instrucciones sobre cómo borrar y aprovisionar nuevamente TPUs.