Entrena en pods de TPU
Descripción general
Las TPU se diseñaron para escalar horizontalmente a un Pod de TPU. Un pod de TPU es una colección de dispositivos de TPU conectados por interfaces de red dedicadas de alta velocidad. Un pod de TPU te permite distribuir la carga de procesamiento en varias TPU. Cada placa de TPU está conectada a una máquina anfitrión basada en CPU de alto rendimiento para tareas como cargar y procesar datos con anterioridad. Para aprovechar al máximo las cantidades más grandes de TPU, debes ajustar varios parámetros de la tarea de entrenamiento.
La configuración para el entrenamiento con pods de TPU es diferente para cada framework. Usa los siguientes vínculos para ver información detallada sobre el entrenamiento en Pods con cada framework:
En las siguientes secciones, se explican algunos problemas habituales, cambios que debes realizar en tus modelos y prácticas recomendadas para reducir o evitar fallas de Pod.
Escala el tamaño del lote y los pasos del entrenamiento
Para lograr el escalamiento lineal en los tipos de TPU más grandes, conserva el mismo tamaño del lote por núcleo.
Por ejemplo, si usas un tamaño del lote de 1,024 en una v2-8, usa un tamaño del lote de 4,096 (4 * 1,024) en un v2-32. Esto usa el hardware de TPU por completo. Puedes usar tamaños del lote más pequeños, sin embargo, tu entrenamiento no se escalará de forma lineal si lo haces.
Algunos modelos incluyen una marca train_steps en la que un paso corresponde a un solo lote de datos. Cuando aumentas el tamaño del lote, reduce la escala verticalmente la cantidad de pasos de entrenamiento, para que la cantidad total de ejemplos de entrenamiento siga siendo la misma.
Por ejemplo, si tienes un tamaño de lote de 1,000 para 100 pasos, se procesan 100,000 ejemplos durante el entrenamiento. Si ahora tienes 4 trabajadores y un tamaño de lote efectivo de 4,000, deberías ajustar la cantidad de pasos a 25 para procesar los mismos 100,000 ejemplos. Si el modelo usa una marca epochs, no es necesario que escales la cantidad de pasos.
Los tamaños del lote más grandes pueden cambiar el comportamiento de convergencia del modelo, por lo que también puedes ajustar algunos hiperparámetros, como la tasa de aprendizaje.
Usa depósitos de Google Cloud Storage regionales en la misma región que el pod de TPU
En general, la práctica recomendada para el entrenamiento de la TPU es usar siempre recursos en la misma región. La región de recursos es muy importante cuando se usan pods de TPU, ya que las tasas de transferencia de datos son más altas cuando el bucket de Google Cloud Storage y la TPU están en la misma región.
Asegúrate de usar un bucket de Google Cloud Storage regional en la misma región que la TPU para los conjuntos de datos de entrenamiento y los puntos de control.
Prácticas recomendadas del flujo de trabajo para el desarrollo en pods de TPU
Cuando se desarrolla una carga de trabajo de TPU nueva, suele ser óptimo comenzar el desarrollo en las TPU más pequeñas y, de manera progresiva, iterar a tamaños de TPU más grandes. Comienza con el uso de una versión pequeña de TPU (por ejemplo, v2-8 o v3-8).
- Prueba la carga de trabajo para ver la funcionalidad
- Prueba y valida el rendimiento con las herramientas de rendimiento
Una vez que tu carga de trabajo sea funcional y alcance tus objetivos de rendimiento, escala verticalmente a un tipo de TPU más grande, como v2-32 o v3-32. De manera gradual y también iterativa, aumenta el tamaño de TPU mientras validas la escalabilidad (funcionalidad y rendimiento) hasta que alcances el tamaño de TPU deseado.