Arquitectura del sistema

Las unidades de procesamiento tensorial (TPU) son aceleradores de AA diseñados por Google. Cloud TPU permite que las TPU estén disponibles como un recurso de nube de GCP escalable. Puedes ejecutar cargas de trabajo de aprendizaje automático en las Cloud TPU mediante frameworks de aprendizaje automático, como TensorFlow, Pytorch y JAX.

La unidad de multiplicación de matrices (MXU) se compone de 128 x 128 que se multiplican o acumulan en un array de presión sistólica. Las MXU proporcionan la mayor parte de la potencia de procesamiento en un TensorCore (núcleo de TPU). Cada MXU puede realizar 16,000 operaciones de multiplicación y acumulación en cada ciclo mediante el formato de número bfloat16.

La VPU se usa para el procesamiento general, como activaciones, softmax, etc. La unidad escalar se usa para el flujo de control, el cálculo de la dirección de memoria y otras operaciones de mantenimiento.

Versiones de TPU

El diseño exacto de una TPU depende de la versión de TPU que uses. Los detalles arquitectónicos y las características de rendimiento de TPU v2 y v3 están disponibles en el artículo sobre una supercomputadora específica del dominio para el entrenamiento de redes neuronales profundas.

TPU v4

La configuración más pequeña de TPU v4 contiene cuatro chips TPU y 32 GiB de HBM y 128 MiB de memoria compartida. Cada chip TPU contiene dos TensorCores. Cada TensorCore tiene cuatro MXU, una unidad de vector y una unidad escalar. En el siguiente diagrama, se muestra un chip TPU v4.

imagen

TPU v3

La configuración más pequeña de TPU v3 contiene cuatro chips TPU y 32 GiB de HBM. Cada chip TPU contiene dos TensorCores. Cada TensorCore tiene dos MXU, una unidad de vector y una unidad escalar. En el siguiente diagrama, se muestra un chip TPU v3.

imagen

TPU v2

La configuración de TPU v2 más pequeña contiene cuatro chips TPU y 16 GiB de HBM. Cada chip TPU contiene dos TensorCores. Cada TensorCore tiene una MXU, una unidad de vector y una unidad escalar. En el siguiente diagrama, se muestra un chip TPU v2.

imagen

Cloud TPU proporciona las siguientes configuraciones de TPU:

  • Un solo dispositivo de TPU
  • Un pod de TPU: un grupo de dispositivos TPU conectados por interconexiones de alta velocidad
  • Una porción de TPU: una subdivisión de un pod de TPU

Beneficios de rendimiento de la TPU v4 en comparación con la v3

Sistema de memoria:

  • Los chips TPU v4 tienen un espacio de memoria unificado de HBM de 32 GiB en todo el chip, lo que permite una mejor coordinación entre los dos TensorCore de chip.
  • Se mejoró el rendimiento de HBM con los estándares y las velocidades de memoria más recientes.
  • Perfil de rendimiento de DMA mejorado con compatibilidad nativa para lograr un alto rendimiento con un nivel de detalle de 512 B

TensorCores:

  • El doble de MXN y una velocidad de reloj más alta con 275 TFLOPS como máximo.
  • El ancho de banda de transmutación y permutación es 2 veces mayor.
  • Modelo de acceso a memoria de almacenamiento de carga para memoria común (Cmem).
  • Mayor ancho de banda de carga de MXU y compatibilidad con el modo de 8 bits para permitir menores lotes y una mejor latencia de inferencia.

Interconexión de chip:

6 vínculos de interconexión por chip para habilitar topologías de red que tienen diámetros de red más pequeños.

Otros:

  • Interfaz PCIE de 3a generación x16 para alojar (conexión directa)
  • Se mejoró el modelo de seguridad.
  • Mejor eficiencia energética.

Beneficios de rendimiento de TPU v3 en comparación con la v2

El aumento de los FLOPS por núcleo y de la capacidad de memoria en las configuraciones de TPU v3 puede mejorar el rendimiento de tus modelos de las siguientes maneras:

  • Las configuraciones de TPU v3 proporcionan importantes beneficios de rendimiento por núcleo para los modelos que dependen del cálculo. Los modelos que dependen de la memoria en las configuraciones de TPU v2 podrían no lograr el mismo rendimiento si también dependen de la memoria en las configuraciones de TPU v3.

  • En los casos en que los datos no caben en la memoria en las configuraciones de TPU v2, TPU v3 puede mejorar el rendimiento y un recálculo reducido de los valores intermedios (rematerialización).

  • Las configuraciones de TPU v3 pueden ejecutar nuevos modelos con tamaños de lote que no cabían en las configuraciones de TPU v2. Por ejemplo, TPU v3 podría permitir ResNets más profundos y también imágenes más grandes con RetinaNet.

Los modelos que dependen casi totalmente de entradas ("entrada") en TPU v2 debido a que los pasos de entrenamiento están esperando entradas también podrían depender de las entradas con Cloud TPU v3. La guía sobre rendimiento de la canalización puede ayudarte a resolver problemas con las entradas.

Configuraciones de TPU v4

Un Pod de TPU v4 está compuesto por 4,096 chips interconectados con vínculos reconfigurables de alta velocidad. Las redes flexibles de TPU v4 te permiten conectar los chips en una porción de mismo tamaño de varias formas. Las TPU v4 se configuran con AcceleratorConfig. AcceleratorConfig te permite especificar el tamaño de los fragmentos en términos de chips a través del campo de topología.

La topología de TPU v4 se especifica mediante una tupla triple que describe la cantidad, las formas y las interconexiones entre los chips TPU. En las siguientes ilustraciones, se muestran algunas topologías comunes de TPU v4.

imagen

Las porciones más grandes se pueden crear a partir de uno o más cubos de 4 x 4 × 4.

Las secciones de TPU de un número determinado de chips se pueden configurar de diferentes maneras. Por ejemplo, una porción de TPU con el acelerador de tipo v4-1024 ahora se puede configurar como 4x4x32, 4x8x16 y 8x8x8. Una TPU de v4-4096 ofrece aún más flexibilidad: 4x4x128, 4x8x64, 4x16x32, 8x16x16, 8x16x16_twisted. Para obtener más información, consulta Topologías y tipos de TPU.

Configuraciones de TPU v2 y v3

Las TPU están disponibles en las siguientes configuraciones:

  1. Una sola placa de TPU
  2. Un pod de TPU
  3. Una porción de pod de TPU

Placa de TPU única

Una configuración de TPU con una sola placa es una placa independiente con 4 chips TPU (8 núcleos de TensorCore) sin conexiones de red a otras placas de TPU. Las TPU de placa única no forman parte de una configuración de pod de TPU y no ocupan una parte de un pod de TPU.

Pods y Slices de TPU

En un fragmento de TPU o de un pod de TPU, los chips TPU se conectan a través de una interconexión de alta velocidad, y cada chip TPU se comunica directamente con los demás chips del dispositivo de TPU. El software de TPU controla automáticamente la distribución de datos a cada núcleo de TPU en un Pod o Slice. Las porciones de pod están disponibles con 32, 128, 512, 1,024 o 2,048 núcleos.

Arquitecturas de VM de Cloud TPU

Las TPU solo pueden realizar operaciones de matriz, por lo que cada placa de TPU está conectada a una máquina anfitrión basada en CPU para realizar operaciones que no se pueden ejecutar en la TPU. Las máquinas host son responsables de cargar datos desde Cloud Storage, procesar previamente los datos y enviarlos a la TPU.

imagen

En un pod de TPU, hay un host de TPU para cada placa de TPU.

imagen

La forma en la que interactúas con el host de TPU (y la placa TPU) depende de la arquitectura de VM de TPU que uses: nodos de TPU o VM de TPU.

Nodos TPU

La arquitectura del nodo 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. Esto puede hacer que sea difícil depurar los errores de entrenamiento y TPU.

imagen

VM de TPU

La arquitectura de VM de TPU elimina la necesidad de la VM de usuario y puedes establecer una conexión SSH directamente a la VM que está conectada físicamente al dispositivo de TPU. Tienes acceso raíz a la VM, por lo que puedes ejecutar código arbitrario. Puedes acceder a los registros y mensajes de error del compilador y el entorno de ejecución.

imagen

Frameworks compatibles

Los frameworks como JAX, PyTorch y TensorFlow acceden a las TPU mediante una biblioteca compartida llamada libtpu que está presente en todas las VM de TPU. Esta biblioteca incluye el compilador de XLA que se usa para compilar los programas de TPU, el software de TPU que se usa para ejecutar los programas compilados y el controlador de TPU que usa el entorno de ejecución para el acceso de bajo nivel a la TPU.

imagen

TensorFlow

VM de TPU

Con las VM de TPU, en lugar de tu código de Python que se ejecuta en una VM de usuario, se puede ejecutar directamente en el host de TPU.

imagen

Para obtener más información sobre TensorFlow y Cloud TPU, consulta Cómo ejecutar modelos de TensorFlow en Cloud TPU.

Nodos TPU

La arquitectura del sistema de nodos de Cloud TPU se creó originalmente para TensorFlow. El usuario no puede acceder a los hosts de TPU y ejecuta una copia sin interfaz gráfica del servidor de TensorFlow. No ejecutan Python ni ningún código de usuario que no se represente como un grafo de TensorFlow. El código de usuario se ejecuta en una VM de usuario diferente que se comunica con los hosts de TPU a través de la red de gRPC.

imagen

PyTorch

VM de TPU

Con las VM de TPU, tu código de PyTorch se ejecuta directamente en los hosts de TPU.

imágenes

Para obtener más información sobre PyTorch y Cloud TPU, consulta Ejecuta modelos de PyTorch en Cloud TPU.

Nodos TPU

PyTorch se ejecuta en la arquitectura de nodos de Cloud TPU con una biblioteca llamada XRT, que permite enviar gráficos de XLA y las instrucciones del entorno de ejecución a través de las conexiones de gRPC de TensorFlow y ejecutarlas en los servidores de TensorFlow. Se requiere una VM de usuario para cada host de TPU.

imagen

JAX

VM de TPU

Con las VM de TPU, no es necesario utilizar VM de usuario, ya que puedes ejecutar tu código directamente en los hosts de TPU.

imagen

Para obtener más información sobre cómo ejecutar JAX en Cloud TPU, consulta la guía de inicio rápido de JAX.

Nodos TPU

JAX en nodos de Cloud TPU se ejecuta de manera similar a PyTorch, ya que se requiere una VM de usuario diferente para cada VM host.

imagen

¿Qué sigue?