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 puede tener hasta 2,048 núcleos de TPU, lo que 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. En este documento, 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.

Varios 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.

Usa un NFS para el almacenamiento de datos

Cuando creas un pod de TPU, se crea una VM independiente para cada nodo TPU. De forma predeterminada, a cada VM de TPU se le asigna un ID de usuario (UID) diferente. Esto provoca problemas cuando intentas acceder al mismo directorio NFS desde varios nodos: el mismo directorio tendría diferentes propietarios en nodos diferentes y no se aplican los permisos estándar de Linux entre nodos. Por ejemplo, un proceso en un nodo no puede escribir en el directorio de registro creado por otro nodo.

Para solucionar este problema, puedes habilitar el Acceso al SO. El Acceso al SO se puede configurar para una instancia de VM específica o un proyecto. Para obtener más información, consulta Configura el Acceso al SO.

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 la TPU más pequeña (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 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.