Guía de rendimiento de Cloud TPU
El primer paso para solucionar problemas de rendimiento de las TPU es generar un perfil de tu modelo. Para obtener más información sobre cómo capturar un perfil de rendimiento, consulta Cómo generar perfiles de tu modelo en Cloud TPU.
Rendimiento del modelo de TPU
En esta sección, se describen problemas generales que pueden reducir el rendimiento del modelo y cómo puedes abordarlos.
El modelo está limitado por la entrada
Las TPU realizan cálculos muy rápido. Para garantizar que la TPU no esté inactiva, es importante asegurarse de que haya un flujo constante de datos que se cargue en la TPU. La forma de hacerlo depende de cómo cargues y proceses tu conjunto de datos. Por ejemplo, puedes leer archivos de datos en paralelo con tf.data.TFRecordset() y el parámetro
num_parallel_reads.El tamaño del lote es demasiado pequeño debido a la fragmentación (división de lotes entre núcleos)
El entorno de ejecución de TPU divide un lote en los 8 núcleos de un dispositivo de TPU (por ejemplo, v2-8 o v3-8). Si especificas un tamaño de lote global de 128, cada núcleo recibe un tamaño de lote de 16 (128 / 8).
Para obtener un uso óptimo de la memoria, usa el tamaño de lote más grande que se ajuste a la memoria de la TPU. Cada núcleo de TPU usa registros vectoriales de 8 × 128 bidimensionales para procesar multiplicaciones de matrices. En general, el tamaño del lote debe poder dividirse por 8 o 128.
Optimizaciones del compilador XLA
XLA es un compilador para el aprendizaje automático que puede producir binarios para TPU, CPU, GPU y otras plataformas. Si bien XLA forma parte de la base de código estándar de TensorFlow, también se puede usar en modelos de PyTorch y JAX. Los modelos para Cloud TPU se convierten en un grafo XLA, que, a su vez, XLA compila en una TPU ejecutable. Para obtener más información sobre XLA, consulta XLA: Optimización del compilador para el aprendizaje automático.
Relleno
Para usar la memoria de la TPU de manera eficiente, estructura tus datos de modo que se puedan dividir en fragmentos de 128 x 8. Cuando los datos de un cálculo de matriz no ocupan un fragmento completo de 128 × 8, el compilador de XLA rellena los tensores. El padding presenta dos inconvenientes:
- Los tensores rellenos no usan el núcleo TPU lo suficiente.
- El padding aumenta la cantidad de almacenamiento en la memoria del chip que necesita un tensor y puede llevar a un error de falta de memoria.
Aunque el compilador de XLA ejecuta operaciones de relleno automáticamente cuando es necesario, puedes determinar la cantidad de operaciones de relleno que se realizan con la herramienta de visor de memoria. Puedes evitar el relleno si eliges dimensiones de tensor que se ajusten bien a las TPU.
Dimensiones del tensor
El compilador de XLA redondea los tamaños de los tensores almacenados en la memoria HBM de la TPU para realizar cálculos de forma más eficaz. Este relleno sucede de manera transparente en el nivel del hardware y no afecta los resultados. No obstante, en ciertos casos el relleno puede provocar un aumento significativo del uso de la memoria y del tiempo de ejecución.
El entorno de ejecución de TPU distribuye los tensores en la memoria para maximizar la eficiencia del cálculo y minimizar el relleno. Para minimizar la sobrecarga de la memoria y maximizar la eficiencia del cálculo, una de las siguientes condiciones debe ser verdadera:
El tamaño total del lote debe ser un múltiplo de 64 (8 por núcleo de la TPU) y los tamaños de las dimensiones de las funciones deben ser un múltiplo de 128.
El tamaño total del lote debe ser un múltiplo de 1,024 (128 por núcleo de la TPU) y los tamaños de las dimensiones de las funciones deben ser un múltiplo de 8.
Se logra una mejor eficiencia con un tamaño de lote de 1,024 y dimensiones de las funciones que sean múltiplos de 128, aunque esto puede no ser posible para todos los modelos.
Fusión
La Fusion es una técnica general que utiliza el compilador de XLA para optimizar programas. Una operación fusionada es la combinación de múltiples operaciones constituyentes que se ejecutarán de forma conjunta.
Por ejemplo, considera las siguientes series de operaciones:
tmp = tf.add(x, y)
result = tf.multiply(tmp, z)
Este código es, aproximadamente, equivalente al siguiente pseudocódigo:
for (i = 0; i < element_count; i++) {
tmp[i] = x[i] + y[i];
}
for (i = 0; i < element_count; i++) {
result = tmp[i] * z[i];
}
Con la fusión, los accesos al arreglo suceden al mismo tiempo:
for (i = 0; i < element_count; i++) {
result = (x[i] + y[i]) * z[i];
}
En este ejemplo, la cantidad de viajes ida y vuelta a la memoria se reduce, y XLA no necesita asignar más espacio para 'tmp'.
La fusión es una optimización crítica y beneficia a Cloud TPU de diferentes maneras:
- Reduce las transferencias de memoria, ya que quita la necesidad de almacenar resultados inmediatos en la memoria principal, lo cual es lento.
- Permite una mejor utilización de unidades de hardware, que, de otra manera, no se hubieran utilizado.
- Puede reducir la utilización de memoria de un modelo, ya que se necesitan menos búferes al mismo tiempo.
Transmisión
La transmisión se produce implícitamente cuando se combinan dos tensores con formas diferentes pero compatibles.
Por ejemplo, tf.add(vector, matrix) requiere que el vector se transmita a la forma de la matriz. El resultado de la operación tiene la misma forma que la matriz. Si deseas obtener más detalles, consulta la guía sobre transmisión de arreglos.
Si bien las transmisiones, normalmente, pueden fusionarse con sus consumidores, cuando se obliga a que una transmisión se materialice, se obtiene como resultado un rendimiento bajo y un mayor uso de memoria.
En el siguiente ejemplo, la transmisión implícita en la adición de un vector y una matriz no se puede fusionar con el argmax, lo que da como resultado una transmisión materializada:
`tf.argmax(tf.add(vector, zero_matrix), axis=0)`