Introducción a Cloud TPU

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. Para obtener información más detallada sobre el hardware de TPU, consulta Arquitectura de TPU. Cloud TPU es un servicio web que pone las TPU a disposición como recursos de procesamiento escalables en Google Cloud.

Las TPU entrenan tus modelos de manera más eficiente con hardware diseñado para realizar operaciones de matrices grandes que suelen encontrarse en los algoritmos de aprendizaje automático. Las TPU tienen memoria de gran ancho de banda (HBM) en el chip, lo que te permite usar modelos y tamaños de lotes más grandes. Las TPU se pueden conectar en grupos llamados slices que escalan tus cargas de trabajo con pocos o sin cambios de código.

El compilador de álgebra lineal acelerada (XLA) debe compilar el código que se ejecuta en las TPU. XLA es un compilador justo a tiempo que toma el grafo emitido por una aplicación de framework de AA y compila los componentes de álgebra lineal, pérdida y gradiente del grafo en código máquina de TPU. El resto del programa se ejecuta en la máquina host de la TPU. El compilador de XLA forma parte de la imagen de la VM de TPU que se ejecuta en una máquina host de TPU.

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 GPUs o CPUs 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 contienen muchas 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 con una cantidad importante de operaciones personalizadas de PyTorch/JAX que deben ejecutarse al menos de manera parcial en CPUs
  • 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 personalizadas de PyTorch/JAX dentro del ciclo de entrenamiento principal
  • Modelos que se entrenan por semanas o meses
  • Modelos grandes con tamaños de lotes efectivos grandes
  • Modelos con incorporaciones ultragrandes comunes en cargas de trabajo de clasificación y recomendación avanzadas

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

  • Programas de álgebra lineal que requieren ramificaciones frecuentes o contienen muchas operaciones de álgebra a nivel de los elementos
  • Cargas de trabajo que requieren aritmética de alta precisión
  • Cargas de trabajo de redes neuronales que contienen operaciones personalizadas en el ciclo de entrenamiento principal

TPU en Google Cloud

Puedes usar las TPU a través de las VMs de Cloud TPU, Google Kubernetes Engine y Vertex AI. En la siguiente tabla, se enumeran los recursos de cada servicio de Google Cloud.

servicioGoogle Cloud Recursos
Cloud TPU Comienza a usar las VMs de Cloud TPU
Google Kubernetes Engine

Acerca de las TPUs en GKE

Ejecuta Ray en GKE con TPU
Vertex AI

Entrenamiento en Vertex AI con TPU

Cómo usar TPU para la predicción en línea en Vertex AI

Recomendaciones para el desarrollo del modelo

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.

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 AA 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, cualquier modelo que tenga tensores con formas dinámicas no es adecuado para las 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 automáticamente 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 una 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.

Cómo comenzar a usar Cloud TPU

Cómo solicitar ayuda

Comunícate con el equipo de asistencia de Cloud TPU. Si tienes un proyecto Google Cloud activo, prepárate para proporcionar la siguiente información:

  • El Google Cloud ID de tu proyecto
  • El nombre de tu TPU, si existe
  • Otra información que quieras proporcionar

Próximos pasos

¿Deseas aprender más sobre Cloud TPU? Los siguientes recursos pueden ayudarte: