Arquitectura del sistema

En este documento, se describe la arquitectura de todos los componentes de hardware y software del sistema de Cloud TPU.

Descripción general

Las unidades de procesamiento tensorial (TPU) son circuitos integrados personalizados específicos de aplicaciones (ASIC) de Google que se utilizan para acelerar las cargas de trabajo de aprendizaje automático. Estas TPU están diseñadas desde cero con el beneficio de la gran experiencia y el liderazgo de Google en el aprendizaje automático.

Puedes usar Cloud TPU y TensorFlow para ejecutar tus propias cargas de trabajo de aprendizaje automático en el hardware acelerador de TPU de Google. Cloud TPU se ha diseñado para ofrecer el mayor rendimiento y flexibilidad, y así ayudar a los investigadores, desarrolladores y empresas a compilar clústeres de procesamiento de TensorFlow que puedan utilizar CPU, GPU y TPU. Las API de Tensorflow de alto nivel facilitan la ejecución de los modelos replicados en el hardware de Cloud TPU.

Tus aplicaciones de TensorFlow pueden acceder a nodos TPU desde contenedores, instancias o servicios en Google Cloud. La aplicación requiere una conexión al nodo TPU a través de tu red de VPC.

Versiones de TPU

Cada versión de TPU define las características específicas de hardware de un dispositivo de TPU. La versión de TPU define la arquitectura para cada núcleo de TPU, la cantidad de memoria de gran ancho de banda (HBM) para cada núcleo de TPU, las interconexiones entre los núcleos de cada dispositivo de la TPU y las interfaces de red disponibles para la comunicación entre dispositivos. Por ejemplo, cada versión de TPU tiene las siguientes características:

  • TPU v2:
    • 8 GiB de HBM para cada núcleo de TPU
    • Una MXU para cada núcleo de TPU
    • Hasta 512 núcleos de TPU totales y 4 TiB de memoria total en un Pod de TPU
  • TPU v3:
    • 16 GiB de HBM para cada núcleo de TPU
    • Dos MXU para cada núcleo de TPU
    • Hasta 2048 núcleos de TPU totales y 32 TiB de memoria total en un Pod de TPU

image

Cada núcleo de TPU tiene unidades escalares, de vectores y matrices (MXU). La MXU proporciona la mayor parte del poder de procesamiento en un chip TPU. Cada MXU puede realizar 16,000 operaciones de multiplicación y acumulación en cada ciclo a una precisión de bfloat16 reducida. Bfloat16 es la representación de un punto flotante de 16 bits que otorga un mejor entrenamiento y mayor exactitud del modelo que la representación IEEE de media precisión.

Cada uno de los núcleos en un dispositivo de TPU puede ejecutar cálculos de usuario (operaciones de XLA) de manera independiente. Las interconexiones con un alto ancho de banda permiten que los chips se comuniquen de manera directa entre sí en el dispositivo de la TPU. En una configuración de pod de TPU, las interfaces de red de alta velocidad dedicadas conectan varios dispositivos de TPU a fin de proporcionar una mayor cantidad de núcleos de TPU y un grupo de memoria de TPU más grande para tus cargas de trabajo de aprendizaje automático.

Beneficios de rendimiento de TPU v3

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.

Para determinar si el rendimiento de tu modelo mejorará con TPU v3, puedes ejecutar comparativas en diferentes versiones de TPU y supervisar el rendimiento por medio de las herramientas de TensorBoard.

Configuraciones de TPU

En un centro de datos de Google, los dispositivos de la TPU están disponibles en las siguientes configuraciones para TPU v2 y TPU v3:

  • TPU con dispositivo único, que son dispositivos de TPU individuales que no están conectados entre sí en una red dedicada de alta velocidad. No puedes combinar varios tipos de TPU con dispositivo único para que colaboren en una carga de trabajo única.
  • Pods de TPU, que son clústeres de dispositivos de TPU que están conectados entre sí en redes dedicadas de alta velocidad.

TPU con dispositivo único

Una configuración de TPU con dispositivo único en un centro de datos de Google consiste en un dispositivo de TPU sin conexiones de red dedicadas de alta velocidad a otros dispositivos de TPU. Tu nodo TPU se conecta solo a este dispositivo único.

image

En el caso de las TPU con dispositivo único, los chips están interconectados en el dispositivo, por lo que la comunicación entre los chips no requiere la CPU del host o los recursos de red del host.

Cuando creas un nodo TPU, especificas el tipo de TPU. Por ejemplo, puedes especificar v2-8 o v3-8 para configurar tu nodo TPU con un dispositivo único. Las TPU con dispositivo único no forman parte de las configuraciones del pod de TPU y no ocupan una parte del pod de TPU. Lee la página sobre tipos de TPU y consulta qué configuraciones de TPU con dispositivo único están disponibles para tus nodos TPU.

Pods de TPU

Una configuración de pod de TPU en un centro de datos de Google tiene múltiples dispositivos de TPU conectados entre sí en una conexión de red dedicada de alta velocidad. Los hosts de tu nodo TPU distribuyen sus cargas de trabajo de aprendizaje automático en todos los dispositivos de TPU.

image

En un pod de TPU, los chips TPU están interconectados en el dispositivo, por lo que la comunicación entre los chips no requiere la CPU del host ni los recursos de red del host. Además, todos los dispositivos de TPU de un pod de TPU se conectan entre sí en redes dedicadas de alta velocidad que tampoco requieren CPU del host o recursos de red del host.

Cuando creas un nodo TPU, debes especificar que deseas un tipo de TPU que ocupe el pod de TPU completo, o bien una fracción menor de ese pod de TPU. Por ejemplo, un tipo de TPU v2-512 ocupa un pod de TPU v2 completo y un tipo de TPU v2-128 ocupa solo un cuarto de un pod de TPU v2. Lee la página sobre tipos de TPU y consulta qué configuraciones de pod de TPU están disponibles para tus nodos TPU.

image

El tipo de pod de TPU v2 proporciona una configuración máxima de 64 dispositivos para un total de 512 núcleos TPU v2 y 4 TiB de memoria de TPU.

image

El tipo de pod de TPU v3 proporciona una configuración máxima de 256 dispositivos para un total de 2048 núcleos TPU v3 y 32 TiB de memoria de TPU.

Consulta la sección sobre versiones de TPU para obtener más información acerca de las diferencias en la arquitectura entre las distintas versiones de TPU.

Puedes usar la API de Cloud TPU a fin de automatizar la administración de la TPU para tus nodos TPU, cualquiera que sea su tamaño. Como resultado, es muy fácil escalar a clústeres de procesamiento masivo, ejecutar tus cargas de trabajo y, luego, reducir la escala esos clústeres cuando se hayan completado tus cargas de trabajo. La compatibilidad de hardware integrada en los chips ofrece un rendimiento lineal efectivo que escala por un amplio rango de cargas de trabajo de aprendizaje profundo. En la práctica, la pila de software de Cloud TPU quita la complejidad de generar, ejecutar y alimentar los programas en la nube de la TPU.

Arquitectura de software

Cuando ejecutas tu aplicación, TensorFlow genera un gráfico de computación y lo envía a un nodo TPU mediante gRPC. El tipo de TPU que seleccionas para tu nodo TPU determina cuántos dispositivos están disponibles en tu carga de trabajo. El nodo TPU compila el grafo de cálculo justo a tiempo y envía el objeto binario del programa a uno o más dispositivos de TPU para su ejecución. Las entradas del modelo se suelen almacenar en Cloud Storage. El nodo TPU transmite las entradas a uno o más dispositivos de TPU para su consumo.

En el siguiente diagrama de bloques, se muestra la arquitectura de software de Cloud TPU, que consiste del modelo de red neuronal, el estimador de TPU y el cliente de TensorFlow, el servidor de TensorFlow y el compilador XLA.

image

Estimador de TPU

Los estimadores de TPU son un conjunto de API de alto nivel creadas en estimadores que simplifican la compilación de modelos para Cloud TPU y que obtienen el máximo rendimiento de la TPU. Cuando escribas un modelo de red neuronal que utilice Cloud TPU, debes utilizar las API de TPU Estimator.

Cliente de TensorFlow

Los estimadores de TPU traducen tus programas en operaciones de TensorFlow, que luego un cliente de TensorFlow convierte en un grafo de procesamiento. Un cliente de TensorFlow comunica el grafo de procesamiento a un servidor de TensorFlow.

Servidor de TensorFlow

Un servidor de TensorFlow se ejecuta en un servidor de Cloud TPU. Cuando el servidor recibe un grafo de procesamiento de parte del cliente de TensorFlow, realiza estas acciones:

  1. Carga entradas desde Cloud Storage.
  2. Divide el grafo en partes que se pueden ejecutar en una Cloud TPU y otras que se deben ejecutar en una CPU.
  3. Genera operaciones de XLA que corresponden al subgrafo que se ejecutará en Cloud TPU.
  4. Invoca el compilador XLA.

Compilador XLA

El XLA es un compilador justo a tiempo que toma como entradas las operaciones del optimizador de alto nivel (HLO) que produce el servidor de TensorFlow. XLA genera código binario que se ejecuta en Cloud TPU, que incluye la organización de los datos de la memoria en el chip a las unidades de ejecución del hardware y la comunicación entre los chips. El código binario generado se carga en Cloud TPU a través de una conectividad PCIe entre el servidor de Cloud TPU y la Cloud TPU y, luego, se inicia para la ejecución.

Pasos siguientes