Guía de rendimiento de Cloud TPU
El primer paso para solucionar problemas de rendimiento de la 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 crear 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. 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 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 X 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. Modelos para Cloud TPU se traducen en un grafo de XLA, que XLA luego 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 forma eficiente, estructura los datos para que se puedan dividir en mosaicos. bloques 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 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.
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 padding si eliges dimensiones de tensor que se ajusten bien a las 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 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 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 la TPU) y los tamaños de las dimensiones de las funciones deben ser un múltiplo 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 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. Para obtener más detalles, consulta la guía sobre
de transmisión de arrays.
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)`