Cloud Tensor Processing Unit (TPU)

Las unidades de procesamiento tensorial (TPU) son los circuitos integrados personalizados específicos de aplicaciones (ASIC) de Google que se usan 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 en aprendizaje automático de Google.

Cloud TPU te permite ejecutar tus cargas de trabajo de aprendizaje automático en el hardware acelerador de TPU de Google mediante TensorFlow. Cloud TPU está diseñada para alcanzar el mayor rendimiento y flexibilidad, y así ayudar a los investigadores, desarrolladores y negocios a compilar clústeres de procesamiento de TensorFlow que puedan aprovechar al máximo las CPU, GPU y TPU. Las API de TensorFlow de alto nivel te ayudan a ejecutar los modelos en el hardware de Cloud TPU.

Ventajas de las TPU

Los recursos de Cloud TPU aceleran el rendimiento de los cálculos de álgebra lineal, muy utilizados en las aplicaciones de aprendizaje automático. Las TPU minimizan el tiempo necesario para alcanzar la exactitud cuando entrenas modelos de redes neuronales complejas y extensas. Un modelo que antes tardaba semanas en entrenarse con otras plataformas de hardware ahora puede converger con las TPU en cuestión de horas.

Disponibilidad

Las Cloud TPU están disponibles en las siguientes zonas:

EE.UU.

Tipo de TPU (v2) Núcleos de TPU v2 Memoria total de TPU Zonas disponibles
v2-8 8 64 GiB us-central1-b
us-central1-c
(us-central1-fsolo para TFRC)
v2-32 (Beta) 32 256 GiB us-central1-a
v2-128 (Beta) 128 1 TiB us-central1-a
v2-256 (Beta) 256 2 TiB us-central1-a
v2-512 (Beta) 512 4 TiB us-central1-a
Tipo de TPU (v3) Núcleos de TPU v3 Memoria total de TPU Zonas disponibles
v3-8 8 128 GiB us-central1-a
us-central1-b
(us-central1-fsolo para TFRC)

Europa

Tipo de TPU (v2) Núcleos de TPU v2 Memoria total de TPU Zonas disponibles
v2-8 8 64 GiB europe-west4-a
v2-32 (Beta) 32 256 GiB europe-west4-a
v2-128 (Beta) 128 1 TiB europe-west4-a
v2-256 (Beta) 256 2 TiB europe-west4-a
v2-512 (Beta) 512 4 TiB europe-west4-a
Tipo de TPU (v3) Núcleos de TPU v3 Memoria total de TPU Zonas disponibles
v3-8 8 128 GiB europe-west4-a
v3-32 (Beta) 32 512 GiB europe-west4-a
v3-64 (Beta) 64 1 TiB europe-west4-a
v3-128 (Beta) 128 2 TiB europe-west4-a
v3-256 (Beta) 256 4 TiB europe-west4-a
v3-512 (Beta) 512 8 TiB europe-west4-a
v3-1024 (Beta) 1,024 16 TiB europe-west4-a
v3-2048 (Beta) 2,048 32 TiB europe-west4-a

Asia-Pacífico

Tipo de TPU (v2) Núcleos de TPU v2 Memoria total de TPU Zonas disponibles
v2-8 8 64 GiB asia-east1-c

Modelo de programación de Cloud TPU

Las Cloud TPU realizan cálculos de vector y matrices densos muy rápidamente. La transferencia de datos entre Cloud TPU y la memoria del host es lenta en comparación con la velocidad del procesamiento. Esto se debe a que la velocidad del bus PCIe es mucho menor que la de la interconexión de Cloud TPU y la de la memoria de gran ancho de banda (HBM) en el chip. Esto significa que una compilación parcial de un modelo, en el que la ejecución “va y viene” entre el host y el dispositivo, usa este último de forma muy poco eficiente, ya que permanecería inactivo la mayor parte del tiempo, a la espera de que lleguen los datos por medio del bus PCIe. A fin de solucionar este problema, el modelo de programación de Cloud TPU está diseñado para ejecutar gran parte del entrenamiento en la TPU, idealmente el ciclo de entrenamiento completo.

A continuación, se detallan algunas características importantes del modelo de programación implementado por TPUEstimator:

  • Se guardan todos los parámetros del modelo en la memoria de gran ancho de banda en el chip.
  • La ejecución de muchos pasos de entrenamiento en un mismo ciclo permite amortizar el costo de iniciar cálculos en Cloud TPU.
  • Se transmiten los datos de entrenamiento de entrada a una cola de entrada en Cloud TPU. Un programa que se ejecuta en Cloud TPU obtiene los lotes desde esas colas durante cada paso de entrenamiento.
  • El servidor de TensorFlow que se ejecuta en la máquina host (la CPU adjunta al dispositivo de Cloud TPU) obtiene los datos y realiza procesamientos previos antes de ingresarlos en el hardware de Cloud TPU.
  • Paralelismo de datos: Los núcleos de Cloud TPU ejecutan un programa idéntico que reside en sus respectivas HBM de manera síncrona. Se realiza una operación de reducción al final de cada paso de red neuronal en todos los núcleos.

Cuándo conviene usar las TPU

Las Cloud TPU están optimizadas para ciertas cargas de trabajo en particular. En algunas situaciones, tal vez quieras usar GPU o CPU en las instancias de Compute Engine para ejecutar tus cargas de trabajo de aprendizaje automático. En general, puedes decidir qué hardware se ajusta mejor a tu carga de trabajo según los siguientes lineamientos:

  • CPU

    • Prototipado rápido que requiere máxima flexibilidad
    • Modelos simples que no tardan mucho en entrenarse
    • Modelos pequeños con tamaños de lotes efectivos pequeños
    • Modelos que están dominados por operaciones de TensorFlow personalizadas escritas en C++
    • Modelos que están limitados por las E/S disponibles o el ancho de banda de red del sistema host
  • GPU

    • Modelos que no están escritos o no pueden escribirse en TensorFlow
    • Modelos para los que no existe la fuente o cuya fuente es muy costosa para cambiarla
    • Modelos con una cantidad importante de operaciones personalizadas de TensorFlow que deben ejecutarse en CPU al menos de manera parcial
    • Modelos con operaciones de TensorFlow que no están disponibles en Cloud TPU (consulta la lista de operaciones de TensorFlow disponibles)
    • Modelos medianos a grandes con tamaños de lotes efectivos más grandes
  • TPU

    • Modelos dominados por cálculos de matrices
    • Modelos sin operaciones de TensorFlow personalizadas dentro del ciclo de entrenamiento principal
    • Modelos que se entrenan por semanas o meses
    • Modelos grandes o muy grandes con tamaños de lotes efectivos muy grandes

Las Cloud TPU no son apropiadas para las siguientes cargas de trabajo:

  • Programas de álgebra lineal que requieren ramificaciones frecuentes o dominados por el álgebra a nivel de los elementos. Las TPU están optimizadas para realizar multiplicaciones de matrices voluminosas con rapidez, por lo que una carga de trabajo que no esté dominada por multiplicaciones de matrices probablemente no tenga un buen rendimiento en TPU, en comparación con otras plataformas.
  • Las cargas de trabajo que acceden a la memoria de manera dispersa podrían no estar disponibles en TPU.
  • Las cargas de trabajo que requieren aritmética de alta precisión. Por ejemplo, la aritmética de doble precisión no es adecuada para las TPU.
  • Cargas de trabajo de redes neuronales que contienen operaciones personalizadas de TensorFlow escritas en C++. De manera más específica, las operaciones personalizadas en el cuerpo del ciclo de entrenamiento principal no son adecuadas para las TPU.

Las cargas de trabajo de redes neuronales deben ser capaces de ejecutar múltiples iteraciones del ciclo de entrenamiento completo en la TPU. Aunque no es un requisito fundamental de las TPU en sí, es una de las limitaciones actuales del ecosistema del software de TPU que se requiere para alcanzar la eficiencia.

Diferencias con el entrenamiento convencional

Un grafo de entrenamiento de TensorFlow típico consiste en varios subgrafos superpuestos, los cuales proporcionan una variedad de funcionalidades, incluidas las que se mencionan a continuación:

  • Operaciones de E/S para leer datos de entrenamiento
  • Etapas de procesamiento previo a la entrada, a menudo conectadas a través de colas
  • Variables del modelo
  • Código de inicialización de esas variables
  • El modelo en sí
  • Funciones de pérdida
  • Código de gradiente (usualmente generado automáticamente)
  • Operaciones de resumen para supervisar el entrenamiento
  • Operaciones de Guardar o Restablecer para establecer puntos de control

En Cloud TPU, los programas de TensorFlow se compilan con el compilador justo a tiempo XLA. Cuando entrenas un modelo en Cloud TPU, el único código que puede compilarse y ejecutarse en el hardware es el que corresponde a las partes densas del modelo, los subgrafos de pérdida y gradiente. Todas las demás partes del programa de TensorFlow se ejecutan en las máquinas host (servidor de Cloud TPU) como parte de una sesión de TensorFlow distribuida regular. En general, esta consiste en operaciones de E/S que leen los datos de entrenamiento, códigos de procesamiento previo (por ejemplo, para decodificar imágenes comprimidas, realizar muestras o recortes aleatorios, o reunir minilotes de entrenamiento) y todas las partes de mantenimiento del grafo, como las tareas de guardar/restablecer puntos de control.

Prácticas recomendadas para el desarrollo del modelo

Cada chip individual de Cloud TPU contiene 2 núcleos, cada uno con varias unidades de matriz (MXU) diseñadas para acelerar los programas dominados por convoluciones y multiplicaciones de matrices densas (consulta esta página sobre la arquitectura del hardware). Los programas que pasan una parte considerable de su tiempo de ejecución realizando multiplicaciones de matrices suelen ajustarse bien a Cloud TPU. En cambio, es probable que un programa cuyos cálculos están dominados por operaciones no matriciales, como suma, cambio de forma o concatenación, no logre un uso alto de las MXU. A continuación, se proporcionan algunos lineamientos que te ayudarán a elegir y compilar modelos apropiados para Cloud TPU.

Réplicas

Cada dispositivo de Cloud TPU contiene cuatro chips, cada uno de los cuales tiene dos núcleos de TPU. Por lo tanto, para lograr una utilización eficiente de Cloud TPU, un programa debería usar cada uno de los ocho núcleos. TPUEstimator proporciona un operador de grafos para compilar y ejecutar un cálculo replicado. Cada réplica es en esencia una copia del grafo de entrenamiento que se ejecuta en cada núcleo y entrena un minilote que tiene 1/8 del tamaño del lote total.

Diseño

El compilador de XLA realiza transformaciones de código, incluso dividir una multiplicación de matriz en bloques más pequeños a fin de ejecutar cálculos en la unidad matriz (MXU) de manera eficiente. Para garantizar la eficiencia de la división, el compilador de XLA emplea la estructura del hardware de las MXU, un arreglo sistólico de 128 x 128, y el diseño del subsistema de memoria de las TPU, que prefiere dimensiones que sean múltiplos de 8. En consecuencia, algunos diseños son más propicios para la división, mientras que otros requieren que se realicen cambios de forma antes de que puedan dividirse. Las operaciones de cambio de forma suelen depender de la memoria en Cloud TPU.

Formas

El compilador de XLA compila un grafo de TensorFlow justo a tiempo antes del primer lote. Si los lotes subsiguientes tienen formas diferentes, el modelo no funciona. (Volver a compilar el grafo cada vez que la forma cambia es demasiado lento). Por lo tanto, todo modelo que tenga tensores de formas dinámicas que cambien en el entorno de ejecución será inapropiado para el uso con TPU.

Relleno

Un programa de Cloud TPU de buen rendimiento será aquel cuyo procesamiento denso se divida fácilmente en fragmentos de 128 x 128. Cuando un cálculo de matriz no puede ocupar una MXU entera, el compilador rellena los tensores con ceros. El relleno presenta dos desventajas:

  • Los tensores rellenos con ceros no usan el núcleo TPU lo suficiente.
  • El relleno aumenta el tamaño del almacenamiento en la memoria del chip que necesita un tensor y puede llevar a un error de falta de memoria en casos extremos.

Aunque el compilador de XLA ejecuta operaciones de relleno de forma automática cuando es necesario, se puede determinar la cantidad de operaciones de relleno que se realizan mediante la herramienta op_profile. Puedes evitar el relleno si eliges dimensiones de tensor que se ajusten bien a las TPU.

Dimensiones

La elección de dimensiones de tensor adecuadas contribuye de manera considerable a extraer el máximo rendimiento del hardware de TPU, en particular las MXU. El compilador de XLA trata de usar el tamaño del lote o la dimensión de los atributos para usar las MXU al máximo. Por lo tanto, uno de estos debe ser un múltiplo de 128. Si no es así, el compilador rellenará uno de estos hasta que llegue a 128. Idealmente, tanto el tamaño del lote como la dimensión de los atributos deberían ser múltiplos de 8, lo que posibilita lograr un alto rendimiento en el subsistema de la memoria.

Operaciones

Consulta la lista de operaciones de TensorFlow disponibles.

Edge TPU

Los modelos de inteligencia artificial (IA) entrenados en la nube necesitan ejecutarse en el perímetro cada vez más. A los efectos de este instructivo, "en el perímetro" se refiere a dispositivos que se ejecutan en los perímetros de la Internet de las cosas (IoT). Entre estos dispositivos, se incluye una gran variedad de sensores y dispositivos inteligentes que recogen datos en tiempo real y los comunican a otros dispositivos y a la nube.

Edge TPU aumenta Cloud TPU y Cloud IoT para proporcionar una infraestructura de extremo a extremo (de hardware y software, de la nube al perímetro) con el fin de facilitar la implementación de soluciones basadas en la IA.

Por el momento, Edge TPU está en fase de acceso anticipado. Para obtener información acerca de este programa, consulta la página de acceso anticipado.

Próximos pasos