Introducción a la TPU de Cloud
Las unidades de procesamiento de tensor (TPU) son circuitos integrados para aplicaciones específicas (ASIC) que se desarrollan de forma personalizada y se utilizan para agilizar cargas de trabajo de aprendizaje automático. Para obtener más información sobre el hardware de las TPU, consulta Arquitectura de las TPU. TPU de Cloud es un servicio web que ofrece TPUs como recursos de computación escalables en Google Cloud.
Las TPU entrenan tus modelos de forma más eficiente con hardware diseñado para realizar operaciones de matrices grandes que suelen encontrarse en los algoritmos de aprendizaje automático. Las TPUs tienen memoria de alto ancho de banda (HBM) en el chip, lo que te permite usar modelos y tamaños de lote más grandes. Las TPUs se pueden conectar en grupos llamados "slices" que escalan tus cargas de trabajo con pocos o ningún cambio en el código.
El código que se ejecuta en las TPUs debe compilarse con el compilador de álgebra lineal del acelerador (XLA). XLA es un compilador just-in-time que toma el gráfico emitido por una aplicación de framework de aprendizaje automático y compila los componentes de álgebra lineal, pérdida y gradiente del gráfico en código de máquina de TPU. El resto del programa se ejecuta en la máquina host de la TPU. El compilador XLA forma parte de la imagen de VM de TPU que se ejecuta en una máquina host de TPU.
Para obtener más información sobre las unidades de procesamiento de tensor, consulta Cómo usar las TPUs.
Cuándo usar las TPUs
Las TPUs de Cloud están optimizadas para cargas de trabajo específicas. En algunas situaciones, puede que quieras usar GPUs o CPUs en instancias de Compute Engine para ejecutar tus cargas de trabajo de aprendizaje automático. En general, puedes decidir qué hardware es el más adecuado para tu carga de trabajo según las directrices que se indican a continuación.
CPUs
- Prototipos rápidos que requieren la máxima flexibilidad
- Modelos sencillos que no tardan mucho en entrenarse
- Modelos pequeños con tamaños de lote pequeños y eficaces
- Modelos que contienen muchas operaciones personalizadas de TensorFlow escritas en C++
- Modelos que están limitados por las E/S disponibles o el ancho de banda de red del sistema host
GPUs
- Modelos con un número significativo de operaciones personalizadas de PyTorch o JAX que deben ejecutarse, al menos parcialmente, en CPUs
- Modelos con operaciones de TensorFlow que no están disponibles en las TPU de Cloud (consulta la lista de operaciones de TensorFlow disponibles)
- Modelos medianos o grandes con tamaños de lote eficaces más grandes
TPUs
- Modelos dominados por cálculos matriciales
- Modelos sin operaciones personalizadas de PyTorch o JAX en el bucle de entrenamiento principal
- Modelos que se entrenan durante semanas o meses
- Modelos grandes con tamaños de lote eficaces y de gran tamaño
- Modelos con inserciones de gran tamaño habituales en cargas de trabajo de ranking y recomendación avanzadas
Las TPUs de Cloud no son adecuadas para las siguientes cargas de trabajo:
- Programas de álgebra lineal que requieren ramificaciones frecuentes o que contienen muchas operaciones de álgebra elemento a elemento
- Cargas de trabajo que requieren aritmética de alta precisión
- Cargas de trabajo de redes neuronales que contienen operaciones personalizadas en el bucle de entrenamiento principal
TPUs en Google Cloud
Puedes usar las TPUs a través de las VMs de TPU de Cloud, Google Kubernetes Engine y Vertex AI. En la siguiente tabla se enumeran los recursos de cada servicio de Google Cloud.
| ServicioGoogle Cloud | Recursos |
|---|---|
| TPU de Cloud | Empezar a usar máquinas virtuales de TPU de Cloud |
| Google Kubernetes Engine | |
| Vertex AI |
Prácticas recomendadas para el desarrollo de modelos
Es probable que un programa cuyo cálculo esté dominado por operaciones que no sean de matrices, como add, reshape o concatenate, no alcance una alta utilización de MXUs. A continuación, se indican algunas directrices que te ayudarán a elegir y crear modelos adecuados para las TPU de Cloud.
Diseño
El compilador XLA realiza transformaciones de código, como dividir una multiplicación de matrices en bloques más pequeños, para ejecutar de forma eficiente los cálculos en la unidad de matriz (MXU). El compilador XLA usa la estructura del hardware de la MXU, una matriz sistólica de 128x128, y el diseño del subsistema de memoria de las TPU, que prefiere dimensiones que sean múltiplos de 8, para mejorar la eficiencia de la segmentación.
Por lo tanto, algunos diseños son más adecuados para acoplar ventanas, mientras que otros requieren que se cambie su forma antes de poder acoplarlos. Las operaciones de cambio de forma suelen estar limitadas por la memoria en las TPU de Cloud.
Formas
El compilador XLA compila un gráfico de aprendizaje automático justo a tiempo para el primer lote. Si los lotes posteriores tienen formas diferentes, el modelo no funcionará. Volver a compilar el gráfico cada vez que cambia la forma es demasiado lento. Por lo tanto, ningún modelo que tenga tensores con formas dinámicas es adecuado para las TPUs.
Relleno
Un programa de TPU de Cloud de alto rendimiento es aquel en el que la computación densa se puede dividir en bloques de 128x128. Cuando un cálculo de matrices no puede ocupar una MXU completa, el compilador rellena los tensores con ceros. El relleno tiene dos inconvenientes:
- Los tensores rellenados con ceros no aprovechan al máximo el núcleo de la TPU.
- El relleno aumenta la cantidad de almacenamiento de memoria en el chip que necesita un tensor y, en casos extremos, puede provocar un error de falta de memoria.
Aunque el compilador de XLA aplica el relleno automáticamente cuando es necesario, se puede determinar la cantidad de relleno que se aplica mediante la herramienta op_profile. Puedes evitar el relleno eligiendo dimensiones de tensor que se adapten bien a las TPUs.
Dimensiones
Elegir las dimensiones de tensor adecuadas es fundamental para extraer el máximo rendimiento del hardware de TPU, especialmente de la MXU. El compilador XLA intenta usar el tamaño del lote o una dimensión de la función para usar al máximo la MXU. Por lo tanto, uno de ellos debe ser un múltiplo de 128. De lo contrario, el compilador añadirá espacio a uno de ellos hasta 128. Lo ideal es que el tamaño de lote y las dimensiones de las características sean múltiplos de 8, lo que permite extraer un alto rendimiento del subsistema de memoria.
Empezar a usar Cloud TPU
- Configurar una Google Cloud cuenta
- Activa la API de Cloud TPU.
- Conceder acceso a las TPUs de Cloud a tus segmentos de Cloud Storage
- Realizar un cálculo básico en una TPU
- Entrenar un modelo de referencia en una TPU
- Analizar el modelo
Solicitar ayuda
Si necesitas ayuda, ponte en contacto con el equipo de Asistencia de Cloud TPU. Si tienes un Google Cloud proyecto activo, ten a mano la siguiente información:
- El Google Cloud ID de tu proyecto
- El nombre de tu TPU, si lo tienes
- Otra información que quieras proporcionar
Siguientes pasos
¿Quieres obtener más información sobre las TPU de Cloud? Los siguientes recursos pueden ayudarte: