Guía de rendimiento de Cloud TPU

El primer paso para solucionar problemas de rendimiento de la TPU es crear un perfil de tu modelo. Si deseas obtener más información para capturar un perfil de rendimiento, consulta Genera un perfil de tu modelo en Cloud TPU.

Rendimiento del modelo de la 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 las TPU, las CPU, las 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)`

Recomendaciones de rendimiento para la arquitectura de doble chiplet Ironwood

El modelo de programación de Ironwood te permite acceder a dos dispositivos TPU en lugar de la arquitectura de un solo núcleo lógico (también conocido como MegaCore) que se usaba en generaciones anteriores (TPU v4 y v5p). Este cambio mejora la rentabilidad y la eficiencia de la fabricación del chip. Si bien esto representa un cambio estructural, el nuevo diseño garantiza que puedas volver a usar los modelos de software existentes con cambios mínimos.

Para lograr el mejor rendimiento con la arquitectura de doble chiplet, recomendamos los siguientes enfoques:

• Usa el paralelismo de tensor en los chiplets: La interfaz D2D de alto ancho de banda está diseñada para un paralelismo de tensor eficiente. Recomendamos dividir los tensores entre los dos dispositivos integrados en el chip.

• Usa colectivos jerárquicos: Para maximizar la eficiencia de la comunicación, aprovecha la jerarquía de red de dos niveles: la vinculación D2D ultrarrápida entre los chiplets en el chip y las vinculaciones ICI rápidas dentro de una porción. Cuando se usa el paralelismo automático con SPMD (programa único, datos múltiples), el compilador de XLA se encarga de esto por ti generando automáticamente operaciones colectivas jerárquicas. Cuando particiones tu modelo de forma manual, también debes diseñar tus patrones de comunicación en torno a esta jerarquía. Prioriza la comunicación entre los dos dispositivos en el mismo chip antes de comunicarse con los dispositivos en otros chips.

• Superposición de la comunicación con el procesamiento: Para maximizar el uso del hardware, descarga las operaciones de comunicación colectiva, como all-reduce, en los SparseCores. Estas operaciones, que no están vinculadas a la unidad de multiplicación de matrices (MXU), se pueden ejecutar en los SparseCores de forma simultánea mientras los TensorCores continúan con su cálculo. Esta técnica puede recuperar algunos de los beneficios de rendimiento inherentes a las operaciones fusionadas en la arquitectura anterior de MegaCore.

• Descarga en SparseCore para las incorporaciones: En el diseño de doble chiplet, las tablas de incorporación se podrían particionar en la HBM de ambos chiplets. Para evitar la degradación del rendimiento debido a esta falta de memoria compartida, descarga las operaciones de recopilación de incorporaciones en SparseCore. Esta estrategia usa la interconexión D2D de alta velocidad para transferir de manera eficiente los vectores de incorporación entre los chiplets. Para obtener más información sobre SparseCore y los modelos de embedding, consulta Análisis detallado de SparseCore para modelos de embedding grandes (LEM).

Para obtener más información sobre la arquitectura de Ironwood en TPU7x, consulta TPU7x (Ironwood).