Entrenar en slices de TPU
Las TPUs están diseñadas para escalarse a un pod de TPU. Un pod de TPU es un conjunto de dispositivos de TPU conectados mediante interfaces de red específicas de alta velocidad. Un pod de TPU te permite distribuir la carga de procesamiento entre varias TPUs. Cada placa de TPU está conectada a una máquina host de alto rendimiento basada en CPU para tareas como la carga y el preprocesamiento de datos. Para aprovechar al máximo un mayor número de TPUs, debes ajustar varios parámetros de la tarea de entrenamiento.
La configuración del entrenamiento con pods de TPU es diferente para cada framework. Utiliza los siguientes enlaces para consultar información detallada sobre la formación en Pods con cada framework:
En las siguientes secciones se explican algunos problemas habituales, los cambios que debes hacer en tus modelos y las prácticas recomendadas para reducir o evitar los fallos de los pods.
Aumentar el tamaño del lote y los pasos de entrenamiento
Para conseguir un escalado lineal en los tipos de TPU más grandes, mantén el mismo tamaño de lote por núcleo.
Por ejemplo, si usas un tamaño de lote de 1024 en una v6e-8, usa un tamaño de lote de 4096 (4 * 1024) en una v6e-32. De esta forma, se aprovecha al máximo el hardware de la TPU. Puedes usar tamaños de lote más pequeños, pero el entrenamiento no se escalará de forma lineal si lo haces.
Algunos modelos incluyen una marca train_steps en la que un paso corresponde al procesamiento de un solo lote de datos. Si aumentas el tamaño del lote, reduce el número de pasos de entrenamiento para que el número total de ejemplos de entrenamiento siga siendo el mismo.
Por ejemplo, si tienes un tamaño de lote de 1000 para 100 pasos, se procesarán 100.000 ejemplos durante el entrenamiento. Si ahora tienes 4 trabajadores y un tamaño de lote efectivo de 4000, tendrías que ajustar el número de pasos a 25 para procesar los mismos 100.000 ejemplos. Si tu modelo usa una marca epochs, no tienes que ajustar el número de pasos.
Los tamaños de 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.
Usar segmentos de Cloud Storage regionales en la misma región que el pod de TPU
En general, la práctica recomendada para el entrenamiento con TPU es usar siempre recursos de la misma región. La región de los recursos es especialmente importante cuando se usan pods de TPU, ya que las velocidades de transferencia de datos son más altas cuando el segmento de Cloud Storage y la TPU se encuentran en la misma región.
Asegúrate de usar un segmento de 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 para el desarrollo de flujos de trabajo en pods de TPU
Cuando desarrolles una nueva carga de trabajo de TPU, suele ser óptimo empezar con las TPUs más pequeñas e ir iterando progresivamente hasta llegar a tamaños de TPU más grandes. Empieza usando una versión pequeña de TPU (por ejemplo, v6e-8).
- Probar la carga de trabajo para comprobar que se comporta como se espera
- Probar y validar el rendimiento con las herramientas de rendimiento
Una vez que tu carga de trabajo funcione y alcance los objetivos de rendimiento, aumenta la escala a un tipo de TPU más grande, como v6e-32. Aumenta el tamaño de la TPU de forma gradual e iterativa mientras validas la escalabilidad (funcionalidad y rendimiento) hasta alcanzar el tamaño de TPU que quieras.