Guía de rendimiento de Cloud TPU

El primer paso para solucionar problemas de rendimiento de TPU es generar un perfil del 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.

  1. El modelo está vinculado a la entrada

    Las TPU realizan cálculos con mucha rapidez. Para garantizar que la TPU no esté inactiva, es importante asegurarse de que se cargue un flujo constante de datos en la TPU. La forma de hacerlo depende de cómo cargues y proceses el conjunto de datos. Por ejemplo, puedes leer archivos de datos en paralelo con tf.data.TFRecordset() y el parámetro num_parallel_reads.

  2. El tamaño del lote es demasiado pequeño debido a la fragmentación (dividir los lotes entre núcleos)

    El entorno de ejecución de TPU divide un lote en los 8 núcleos de un dispositivo 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 lograr un uso óptimo de la memoria, usa el tamaño de lote más grande que quepa en la memoria de TPU. Cada núcleo de TPU usa registros vectoriales de dos dimensiones de 8 x 128 para procesar multiplicaciones de matrices. En general, el tamaño del lote debe ser divisible por 8 o 128.

Optimizaciones del compilador XLA

XLA es un compilador de aprendizaje automático 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 PyTorch y JAX. Los modelos de Cloud TPU se traducen a un grafo de XLA, que, luego, 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 TPU de forma eficiente, estructura tus datos a fin de que puedan dividirse en fragmentos de 128 x 8. Cuando los datos para un cálculo de matriz no completan un fragmento completo de 128 x 8, el compilador de XLA rellena los tensores. El relleno presenta dos desventajas:

  1. Los tensores rellenos tienen un uso inferior del núcleo TPU.
  2. El relleno aumenta el volumen de memoria en el chip que requiere un tensor y puede llevar a un error de falta de memoria.

Si bien el compilador de XLA ejecuta operaciones de relleno automáticamente cuando es necesario, puedes determinar la cantidad de relleno que se realiza con la herramienta del visor de memoria. Puedes evitar el relleno si eliges dimensiones de tensor que sean adecuadas para TPU.

Dimensiones de tensores

El compilador de XLA redondea los tamaños de los tensores almacenados en la memoria TPU HBM para realizar cálculos de manera más eficiente. Este relleno ocurre con transparencia 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 de procesamiento y minimizar el relleno. Para minimizar la sobrecarga de la memoria y maximizar la eficiencia del procesamiento, se debe cumplir con una de las siguientes condiciones:

  1. 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 los atributos deben ser un múltiplo de 128.

  2. 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 atributos que sean múltiplos de 128, aunque es posible que no sea posible para todos los modelos.

Fusión

La fusión es una técnica general que utiliza el compilador de XLA a fin de optimizar programas. Una operación combinada es la combinación de varias operaciones constituyentes que se ejecutarán de manera combinada.

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 el número de viajes de ida y vuelta de 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. Para obtener más detalles, consulta la guía sobre cómo transmitir arreglos.

Si bien las transmisiones a menudo pueden fusionarse con sus consumidores, forzar una transmisión puede generar un rendimiento deficiente y un aumento en el uso de la 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)`