Introducción a Cloud TPU
Las unidades de procesamiento tensorial (TPU) son los circuitos integrados específicos de la aplicación (ASIC) personalizados de Google que se usan para acelerar las cargas de trabajo de aprendizaje automático. Para obtener más información sobre el hardware de las TPU, consulta Arquitectura de TPU. Cloud TPU es un servicio web que pone a disposición las TPU 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 se encuentran a menudo 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 porciones que escalan verticalmente tus cargas de trabajo con pocos cambios o ninguno en el código.
El compilador acelerador de álgebra lineal (XLA) debe compilar el código que se ejecuta en las TPU. XLA es un compilador justo a tiempo que toma el grafo que emite una aplicación de framework de AA y compila los componentes de álgebra lineal, pérdida y gradiente del grafo en código de máquina de TPU. El resto del programa se ejecuta en la máquina anfitrión de la TPU. El compilador XLA forma parte de la imagen de la VM de TPU que se ejecuta en una máquina anfitrión de TPU.
Para obtener más información sobre las unidades de procesamiento tensorial, consulta Cómo pensar en las 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 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 eficaces y pequeños
- Modelos que contienen muchas operaciones personalizadas de TensorFlow escritas en C++
- Modelos que están limitados por la 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 en parte en CPU
- 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 eficaces más grandes
TPU
- Modelos dominados por cálculos de matrices
- Modelos sin operaciones de PyTorch ni JAX personalizadas dentro del bucle de entrenamiento principal
- Modelos que se entrenan por semanas o meses
- Modelos grandes con tamaños de lotes eficaces grandes
- Modelos con embeddings ultragrandes que son comunes en cargas de trabajo avanzadas de clasificación y recomendación
Las Cloud TPU no son apropiadas para las siguientes cargas de trabajo:
- Programas de álgebra lineal que requieren ramificaciones frecuentes o contienen muchas operaciones algebraicas a nivel de los elementos
- Las cargas de trabajo que requieren aritmética de alta precisión
- Cargas de trabajo de entrenamiento de redes neuronales que contienen operaciones personalizadas en el bucle principal de entrenamiento
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 para cada servicio de Google Cloud.
| Servicio deGoogle Cloud | Recursos |
|---|---|
| Cloud TPU | Comienza a usar las VMs de Cloud TPU |
| Google Kubernetes Engine | |
| Vertex AI |
Prácticas recomendadas para el desarrollo del modelo
En cambio, es probable que un programa cuyos cálculos estén dominados por operaciones sin matrices, como suma, cambio de forma o concatenación, no logre un uso alto de 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 XLA realiza transformaciones de código, incluso dividir una multiplicación de matriz en bloques más pequeños, para ejecutar cálculos en la unidad de matriz (MXU) de manera eficiente. Para garantizar la eficiencia de la división, el compilador XLA emplea la estructura del hardware de la MXU, un array sistólico de 128x128, 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 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, todo modelo que tenga tensores de formas dinámicas será inapropiado para el uso con TPU.
Relleno
Un programa de Cloud TPU de buen rendimiento será aquel cuyo procesamiento denso se divida con facilidad en fragmentos de 128x128. 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 rellenados 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 XLA ejecuta operaciones de relleno de forma automática cuando es necesario, se puede determinar la cantidad de veces que se realizan con la herramienta op_profile. Puedes evitar el relleno eligiendo 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 la MXU. El compilador XLA trata de usar el tamaño del lote o una dimensión de los atributos para usar la 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. En lo posible, 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.
Comienza a usar Cloud TPU
- Configura una cuenta de Google Cloud .
- Activa la API de Cloud TPU.
- Otorga acceso a Cloud TPU a tus buckets de Cloud Storage.
- Ejecuta un cálculo básico en una TPU.
- Entrena un modelo de referencia en una TPU.
- Analiza tu modelo.
Cómo solicitar ayuda
Para obtener ayuda, comunícate con el equipo de asistencia de Cloud TPU. Si tienes un proyecto de Google Cloud activo, prepárate para proporcionar la siguiente información:
- El ID de tu proyecto de Google Cloud
- Tu nombre de TPU, si existe uno
- Otra información que quieras proporcionar
¿Qué sigue?
¿Quieres aprender más sobre Cloud TPU? Los siguientes recursos pueden ayudar: