Guía de rendimiento de Cloud TPU

El primer paso para solucionar problemas de rendimiento de TPU es crear un perfil de tu modelo. Para obtener más información sobre cómo capturar un perfil de rendimiento, consulta Genera un perfil 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á 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 se le cargue un flujo constante de datos. La forma de hacerlo depende de cómo cargues y preproceses tu 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 (división de lotes entre núcleos)

   El entorno de ejecución de la 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 uno de 16 (128/8).

   Para obtener un uso óptimo de la memoria, usa el tamaño de lote más grande que quepa en la memoria de la TPU. Cada núcleo de TPU usa registros de vectores bidimensionales de 8 x 128 para procesar multiplicaciones de matrices. En general, el tamaño del lote debe ser divisible de forma uniforme por 8 o 128.

3. Ajuste de la administración de la memoria

   Puedes usar las variables de entorno TPU_PREMAPPED_BUFFER_SIZE para ajustar los comportamientos del entorno de ejecución de bajo nivel.

• Descripción: TPU_PREMAPPED_BUFFER_SIZE establece el tamaño del búfer de memoria del host (en bytes) que se asigna previamente y se fija para que el entorno de ejecución de la TPU lo use en las transferencias de datos (por ejemplo, DMA). El valor predeterminado es 4,294,967,296 bytes. El valor debe ser un múltiplo de 2 ^ 12 (4 KB = 4 * 1,024 bytes = 4,096 = 2 ^ 12).

  Los siguientes ejemplos son valores TPU_PRE_MAPPED_BUFFER_SIZE válidos.

      17179869184 = 2^34 = 2^22 * 2^12 (2^22 4KB pages will be premapped).
      40000000000 = 5^10 * 2^12 = (5^10 4KB pages will be premapped).
  

• Impacto: Aumentar este tamaño puede mejorar el rendimiento de la transferencia de datos entre el host y el dispositivo de TPU, especialmente para las cargas de trabajo con tensores grandes o comunicación frecuente entre el host y el dispositivo. Sin embargo, también aumenta la cantidad de memoria del host fijada, lo que reduce la memoria disponible para otros procesos.

  Tamaño del búfer

  Si la región del búfer asignada previamente no es lo suficientemente grande para asignar una memoria durante el entorno de ejecución del programa, la carga de trabajo fallará y devolverá un error RESOURCE_EXHAUSTED, similar al siguiente:

  "No se pudo asignar el búfer desde la región asignada previamente con el siguiente error: RESOURCE_EXHAUSTED: Se intentará asignar allocation_size. Eso no fue posible. Hay available_size gratis".

  Si el búfer es demasiado grande, la inicialización de la TPU puede tardar mucho más (posiblemente, más de 15 segundos), lo que puede hacer que parezca que la TPU se detuvo.

  Para diagnosticar este problema, inspecciona los registros del entorno de ejecución de la TPU. En estos registros, se detallan las operaciones que se realizan, incluida la asignación previa de búferes. Puedes encontrar los registros en /tmp/tpu_logs/tpu_driver.INFO o imprimirlos directamente en la consola configurando la variable de entorno TPU_STDERR_LOG_LEVEL=0. Este parámetro de configuración generará un resultado similar al siguiente:

   I0604 12:45:24.926233   62136 tpu_hal.cc:214] Starting premapped memory manager initialization...
   I0604 12:45:29.411218   62136 system.cc:1059] tpu::System initialized, current host id: 0, logical device ids: 0
   I0604 12:45:29.411244   61600 tfrt_tpu_system_state.cc:216] CreateTpuSystemState: TPU initialization is successful and it took 5.583190661s
   I0604 12:45:29.411267   61600 tfrt_tpu_system_state.cc:220] CreateTpuSystemState: using TPU host premapped buffer of size: 4294967296
   ```
  
  This output will tell you how long it took to initialize the TPU and
  the size of the premapped buffer.
  

• Uso: Si el búfer asignado previamente es demasiado pequeño o grande, puedes establecer su tamaño de forma manual con las siguientes variables de entorno.

  TPU_PREMAPPED_BUFFER_SIZE: Sets the total size (in bytes) of the
  pre-mapped buffer region.
  TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES: Sets the maximum size of
  a single buffer that can be allocated from the pre-mapped region.
  

  Por ejemplo, puedes hacer lo siguiente:

   export TPU_PREMAPPED_BUFFER_SIZE=4294967296
  

  para establecer el tamaño del búfer y:

   export TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES
   ```
   to enable it.
  
   This export sets the size to the default.
  

• Orientación: Ajusta el valor de TPU_PREMAPPED_BUFFER_SIZE si sospechas que la transferencia de datos entre el host y el dispositivo es un cuello de botella. Supervisa el uso de la memoria del host y el rendimiento del modelo para encontrar un equilibrio óptimo. Por lo general, el valor predeterminado es suficiente para la mayoría de los casos de uso.

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 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 traducen a un grafo XLA que, a su vez, XLA compila en una TPU ejecutable. Para obtener más información sobre XLA, consulta XLA: Optimiza el 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 para un cálculo de matriz no llenan un fragmento de 128 x 8 completo, el compilador XLA rellena los tensores. El relleno presenta dos desventajas:

1. Los tensores con relleno no usan el núcleo TPU lo suficiente.
2. 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 XLA ejecuta operaciones de relleno de forma automática cuando es necesario, puedes determinar la cantidad de estas operaciones que se realizan con la herramienta del visualizador de memoria. Puedes evitar el relleno eligiendo dimensiones de tensor que se ajusten bien a las TPU.

Dimensiones del tensor

Para alcanzar el máximo de flops, las dimensiones de la multiplicación de matrices deben ser mayores que el tamaño de la MXU para la versión de TPU que usas. El tamaño de la MXU es de 256 x 256 para v6e, y de 128 x 128 para las versiones anteriores a v6e. Para obtener más información, consulta Arquitectura del sistema de Cloud TPU.

Tamaño del lote

El compilador XLA redondea los tamaños de los tensores almacenados en la memoria HBM de la TPU para realizar cálculos de manera más eficiente. 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 la TPU distribuye los tensores en la memoria para maximizar la eficiencia del procesamiento y minimizar el relleno. Para hacer esto y no sobrecargar la memoria, una de las siguientes condiciones debe ser verdadera:

1. El tamaño total del lote debe ser un múltiplo de 64 (8 por núcleo de la TPU) y las dimensiones de las funciones 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 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.

Nota: La dimensión de las funciones se refiere al tamaño oculto de una capa completamente conectada o a la cantidad de canales de salida en una convolución. No todas las capas pueden cumplir con esta regla, en especial la primera y última capa de la red. Esto es correcto, ya que la mayoría de los modelos requieren cierta cantidad de relleno.

Fusión

La fusión es una técnica general que utiliza el compilador XLA para optimizar programas. Una operación fusionada es la combinación de múltiples operaciones constituyentes que se ejecutarán 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[i] = tmp[i] * z[i];
    }

Con la fusión, los accesos al array suceden al mismo tiempo:

    for (i = 0; i < element_count; i++) {
      result[i] = (x[i] + y[i]) * z[i];
    }

En este ejemplo, la cantidad de idas y vueltas a la memoria se reduce, y XLA no necesita asignar más espacio para "tmp".

La fusión es una optimización fundamental 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 emisió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 emita 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 Emite arrays.

Si bien las emisiones suelen fusionarse con sus consumidores, forzar una emisión puede provocar un rendimiento bajo y un mayor uso de memoria.

En el siguiente ejemplo, la emisió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 emisión materializada:

`tf.argmax(tf.add(vector, zero_matrix), axis=0)`