Guía de rendimiento de Cloud TPU
El primer paso para solucionar problemas de rendimiento de TPU es generar un perfil del modelo. Si deseas 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á vinculado a 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 carguen en la TPU. La forma en que esto se hace depende de cómo cargues y proceses tu conjunto de datos de forma previa. 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 (la 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 del lote global de 128, cada núcleo recibe un tamaño de lote de 16 (128 / 8).
Para un uso óptimo de la memoria, usa el tamaño de lote más grande que se adapte a la memoria de la TPU. Cada núcleo de TPU usa registros vectoriales de 8 x 128 bidimensionales para procesar las multiplicaciones de matrices. En general, el tamaño del lote debe ser divisible por 8 o 128 de manera uniforme.
Optimizaciones del compilador XLA
XLA es un compilador para el aprendizaje automático que puede producir objetos binarios para TPU, CPU, GPU y otras plataformas. Si bien XLA es parte de la base de código estándar de TensorFlow, también se puede usar en los modelos PyTorch y JAX. Los modelos de Cloud TPU se traducen a un grafo XLA, que 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 TPU de manera eficiente, estructura los datos de modo que se puedan dividir en 128 x 8 fragmentos. Cuando los datos para un cálculo de matriz no ocupan un fragmento completo de 128 x 8, el compilador de XLA rellena los tensores. El relleno presenta dos desventajas:
- Los tensores con padding no usan mucho el núcleo TPU.
- El relleno aumenta la cantidad de almacenamiento en la memoria del chip que se requiere para un tensor y puede generar un error por falta de memoria.
Si bien el compilador de XLA ejecuta operaciones de relleno automáticamente cuando es necesario, puedes determinar la cantidad que se realiza con la herramienta de visualización de memoria. Puedes evitar el relleno si eliges dimensiones de tensor que sean adecuadas para TPU.
Dimensiones de tensor
El compilador de XLA redondea los tamaños de los tensores almacenados en la memoria HBM de TPU para realizar cálculos de manera más eficiente. Este relleno ocurre 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 dispone los tensores en la memoria para maximizar la eficiencia computacional y minimizar el relleno. Para minimizar la sobrecarga de memoria y maximizar la eficiencia del procesamiento, debe cumplirse una de las siguientes condiciones:
El tamaño total del lote debe ser un múltiplo de 64 (8 por núcleo de TPU) y los tamaños de dimensión 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 TPU) y los tamaños de la dimensión de las funciones deben ser un múltiplo de 8.
El uso de un tamaño de lote de 1,024 y dimensiones de atributos que sean múltiplos de 128 da como resultado la mejor eficiencia, aunque es posible que esto no sea posible para todos los modelos.
Fusión
La fusión es una técnica general que usa el compilador de XLA a fin de optimizar programas. Una operación fusionada es la combinación de varias 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, se reduce la cantidad de recorridos de ida y vuelta en la memoria, y XLA no necesita asignar ningún 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 para transmitir arreglos.
Si bien las transmisiones a menudo se pueden fusionar con sus consumidores, forzar una transmisión puede generar un rendimiento deficiente 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)`