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 un conjunto de dispositivos de TPU conectados mediante 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 de alto rendimiento basada en CPU para tareas de carga y procesamiento previo. Para aprovechar al máximo grandes cantidades de TPU, debes ajustar varios parámetros de tarea de entrenamiento. En este documento, se explican algunos problemas comunes, cambios que debes realizar en tus modelos, y prácticas recomendadas para reducir o evitar fallas de pods.

Escala el tamaño del lote y los pasos del entrenamiento

Para lograr un escalamiento lineal en tipos de TPU más grandes, mantén el mismo tamaño de 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. De esta manera, se utiliza por completo el hardware de TPU. Puedes usar tamaños del lote más pequeños, pero tu entrenamiento no se escalará linealmente si lo haces.

Muchos modelos incluyen una marca train_steps en la que un paso corresponde a procesar un solo lote de datos. Cuando aumentes el tamaño del lote, reduce la escala de la cantidad de pasos de entrenamiento para que la cantidad total de ejemplos de entrenamiento permanezca igual.

Por ejemplo, si tienes un tamaño de lote de 1,000 para los 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, tendrías que ajustar la cantidad de pasos a 25 para procesar los mismos 100,000 ejemplos. Si tu modelo usa una marca epochs, no necesitas escalar la cantidad 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.

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 porque las tasas de transferencia de datos son mayores 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 separada para cada nodo TPU. De forma predeterminada, a cada VM de TPU se le asigna un ID de usuario (UID) diferente. Esto genera problemas cuando intentas acceder al mismo directorio NFS desde varios nodos: el mismo directorio tendría diferentes propietarios en distintos nodos y los permisos estándar de Linux no se aplican a todos los nodos. Por ejemplo, un proceso en un nodo no puede escribir en el directorio de registro que creó 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 para un proyecto. Para obtener más información, consulta la sección sobre cómo configurar 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, luego, iterar de manera progresiva para los tamaños de TPU más grandes. Comienza con la TPU más pequeña (por ejemplo, v2-8 o v3-8).

  • Prueba la carga de trabajo para la funcionalidad
  • Pruebe y valide el rendimiento con las herramientas de rendimiento

Una vez que tu carga de trabajo sea funcional y alcance tus objetivos de rendimiento, escala verticalmente hasta v2-32 o v3-32. Aumenta de forma gradual y iterativa el tamaño de la TPU, a la vez que validas la escalabilidad (funcionalidad y el rendimiento) hasta alcanzar el tamaño de TPU deseado.