Guía de rendimiento de TPU de Cloud

El primer paso para solucionar problemas de rendimiento de la TPU es crear un perfil de tu modelo. Para obtener más información sobre cómo capturar un perfil de rendimiento, consulta el artículo Elaborar un perfil de modelo en Cloud TPU.

Rendimiento del modelo de TPU

En esta sección se describen los problemas generales que pueden reducir el rendimiento del modelo y cómo puedes solucionarlos.

1. El modelo está limitado por la entrada

   Las TPUs realizan cálculos muy rápido. Para asegurarte de que la TPU no esté inactiva, es importante que se cargue un flujo constante de datos en ella. 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 tiempo de ejecución de la TPU divide un lote entre 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 optimizar el uso de la memoria, utiliza 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 de lote debe ser divisible entre 8 o 128.

3. Ajuste de la gestión de la memoria

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

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

  Los siguientes ejemplos son valores válidos de TPU_PRE_MAPPED_BUFFER_SIZE.

      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 TPU, especialmente en 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 de host fijada, lo que reduce la memoria disponible para otros procesos.

  Tamaño del búfer

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

  "No se ha podido asignar el búfer de la región preasignada: RESOURCE_EXHAUSTED. Se ha intentado asignar allocation_size. No era posible. Hay available_size gratis".

  Si el búfer es demasiado grande, la inicialización de la TPU puede tardar mucho más (quizá más de 15 segundos), lo que puede dar la impresión de que la TPU se ha bloqueado.

  Para diagnosticar este problema, inspecciona los registros del tiempo de ejecución de la TPU. Estos registros detallan las operaciones que se están realizando, incluida la asignación previa de los 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 ajuste 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 preasignado es demasiado pequeño o demasiado grande, puedes definir manualmente el tamaño del búfer mediante 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:

   export TPU_PREMAPPED_BUFFER_SIZE=4294967296
  

  para definir 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.
  

• Recomendació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. Monitoriza el uso de la memoria del host y el rendimiento del modelo para encontrar un equilibrio óptimo. El valor predeterminado suele ser suficiente para la mayoría de los casos prácticos.

Optimizaciones del compilador XLA

XLA es un compilador para aprendizaje automático que puede generar archivos binarios para TPUs, CPUs, GPUs y otras plataformas. Aunque 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 de TPU de Cloud se traducen a un gráfico de XLA, que XLA compila en un ejecutable de TPU. Para obtener más información sobre XLA, consulta XLA: compilador de optimización para aprendizaje automático.

Relleno

Para usar la memoria de la TPU de forma eficiente, estructura los datos de forma que se puedan dividir en bloques de 128 x 8. Cuando los datos de un cálculo de matriz no llenan un bloque de 128x8 completo, el compilador de XLA rellena los tensores. El relleno tiene dos inconvenientes:

1. Los tensores con relleno no aprovechan al máximo el núcleo de la TPU.
2. El relleno aumenta la cantidad de almacenamiento de memoria en el chip que necesita un tensor y puede provocar un error de falta de memoria.

Aunque el compilador de XLA aplica el relleno automáticamente cuando es necesario, puedes determinar la cantidad de relleno que se aplica con la herramienta de visualización de memoria. Puedes evitar el relleno eligiendo dimensiones de tensor que se adapten bien a las TPUs.

Dimensiones de 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 de la versión de TPU que estés usando. El tamaño de MXU es de 256x256 en la versión 6e y de 128x128 en las versiones anteriores. Para obtener más información, consulta la arquitectura del sistema de las TPU de Cloud.

Tamaño del lote

El compilador XLA redondea los tamaños de los tensores almacenados en la memoria HBM de la TPU para realizar los cálculos de forma más eficiente. Este relleno se produce de forma transparente a nivel de hardware y no afecta a los resultados. Sin embargo, en algunos casos, el relleno puede provocar un aumento significativo del uso de memoria y del tiempo de ejecución.

El tiempo de ejecución de la TPU organiza los tensores en la memoria para maximizar la eficiencia computacional y minimizar el relleno. Para minimizar la sobrecarga de memoria y maximizar la eficiencia computacional, una de las siguientes afirmaciones debe ser verdadera:

1. El tamaño total del lote debe ser un múltiplo de 64 (8 por núcleo de TPU) y los tamaños de las dimensiones de las características deben ser un múltiplo de 128.

2. El tamaño total del lote debe ser un múltiplo de 1024 (128 por núcleo de TPU) y los tamaños de las dimensiones de las características deben ser un múltiplo de 8.

Si se usa un tamaño de lote de 1024 y dimensiones de características que sean múltiplos de 128, se obtendrá la mejor eficiencia, aunque puede que no sea posible en todos los modelos.

Nota: Dimensión de la característica hace referencia al tamaño oculto de una capa totalmente conectada o al número de canales de salida de una convolución. No todas las capas pueden cumplir esta regla, especialmente la primera y la última de la red. No pasa nada, la mayoría de los modelos requieren un cierto relleno.

Fusion

Fusion es una técnica general que usa el compilador XLA para optimizar programas. Una operación fusionada es la combinación de varias operaciones constitutivas que se van a ejecutar de forma combinada.

Por ejemplo, consideremos la siguiente serie de operaciones:

    tmp = tf.add(x, y)
    result = tf.multiply(tmp, z)

Este código es más o menos 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 a la matriz se producen al mismo tiempo:

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

En este ejemplo, se reduce el número de viajes de ida y vuelta de la memoria y XLA no necesita asignar espacio para "tmp".

La fusión es una optimización fundamental que beneficia a las TPU de Cloud de varias formas:

• Reduce las transferencias de memoria al eliminar la necesidad de almacenar resultados intermedios en la memoria principal, que es lenta.
• Permite un mayor aprovechamiento de las unidades de hardware que, de lo contrario, no se utilizarían.
• Puede reducir la utilización de memoria de un modelo, ya que se necesitan menos búferes activos al mismo tiempo.

Emisió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 difunda a la forma de la matriz. El resultado de la operación tiene la misma forma que la matriz. Para obtener más información, consulta la guía sobre difusión de arrays.

Aunque las emisiones a menudo se pueden fusionar con sus consumidores, forzar una emisión puede provocar un rendimiento deficiente y un mayor uso de memoria.

En el siguiente ejemplo, la emisión implícita en la suma 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)`