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á 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 cargue en la TPU. Cómo se hace esto depende de cómo cargas y procesas previamente 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 con un tamaño de lote de 16 (128 / 8).
Para lograr un uso óptimo de la memoria, utiliza el tamaño de lote más grande que se adapte a las TPU. memoria. Cada núcleo de TPU usa registros vectoriales de 8 × 128 bidimensionales para procesar multiplicaciones de matrices. En general, el tamaño del lote debería ser divisible por 8 o 128 de manera uniforme.
Optimizaciones del compilador XLA
XLA es un compilador para modelos de máquina que puede producir objetos 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. Modelos para Cloud TPU se traducen en un grafo de XLA, que XLA luego compila en una TPU ejecutable. Si deseas obtener más información sobre XLA, consulta XLA: Optimizing Compiler for Machine Learning.
Relleno
Para usar la memoria de TPU de forma eficiente, estructura los datos para que se puedan dividir en mosaicos. bloques de 128 x 8. Es cuando los datos para el cálculo de una matriz no llenan un Bloque de 128 x 8, el compilador XLA rellena los tensores. El relleno presenta dos desventajas:
- Los tensores rellenos no usan el núcleo TPU lo suficiente.
- El relleno aumenta la cantidad de almacenamiento de memoria en el chip que se requiere para un tensor y puede generar un error de memoria insuficiente.
Si bien el compilador de XLA ejecuta rellenos automáticamente cuando es necesario, puedes puedes determinar la cantidad de relleno que se realiza con la herramienta Visualizador de memoria. Puedes y evita el relleno; para ello, elige las dimensiones del tensor que sean adecuadas para la TPU.
Dimensiones de 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 ocurre de forma transparente en la de 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 del procesamiento. y minimizan el relleno. Para minimizar la sobrecarga de la memoria y maximizar de eficiencia, debe ser verdadera una de las siguientes opciones:
El tamaño total del lote debe ser un múltiplo de 64 (8 por núcleo de la TPU) y debe tener los tamaños de las dimensiones deben ser múltiplos 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 las dimensiones del atributo deben ser múltiplos de 8.
Con un tamaño de lote de 1,024 y dimensiones de atributos que sean múltiplos de 128 ofrece la mejor eficiencia, aunque esto no sea posible en todos los modelos.
Fusión
La Fusion es una técnica general que usa el compilador de XLA para optimizar programas. R una operación fusionada es la combinación de múltiples operaciones constituyentes que se que se ejecuten en combinación.
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 recorridos de ida y vuelta de memoria se reduce y XLA no deberá asignar cualquier 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. Para obtener más detalles, consulta la guía sobre
de transmisión de arrays.
Si bien las transmisiones a menudo se pueden fusionar con sus consumidores, lo que obliga a una lo que puede generar un rendimiento deficiente y un aumento en el 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)`