Código de error: E1000

Categoría: Tiempo de compilación: HBM OOM

Este error indica que el programa requiere más memoria de gran ancho de banda (HBM) de la que está disponible físicamente en el dispositivo TPU.

Mensajes de error de muestra:

RESOURCE_EXHAUSTED: TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G

RESOURCE_EXHAUSTED: XLA:TPU compile permanent error. Ran out of memory in memory space hbm. Used 49.34G of 32.00G hbm. Exceeded hbm capacity by 17.34G.

Backends de XLA: TPU

Descripción general

XLA realiza verificaciones para garantizar que el tamaño agregado de todas las asignaciones estáticas necesarias quepa en la HBM del dispositivo.

El compilador administra la capacidad fija de HBM de la TPU para varios tipos de asignaciones:

• Entradas y salidas del programa: Lotes de entrenamiento, estados del optimizador, etcétera
• Temporales de TPU: Memoria dinámica necesaria para cálculos intermedios (p. ej., activaciones, gradientes, etcétera)
• Objeto binario compilado: El código máquina para TensorCore (TC) y SparseCore (SC)
• Sobrecarga del sistema: Espacio reservado para el tiempo de ejecución de XLA (p. ej., búferes de entrada en generaciones anteriores de TPU)
• Constantes: Los valores constantes incorporados en la IR de HLO se asignan en HBM
• Internos del compilador: Asignaciones a nivel del programa y por HLO (p. ej., información de enrutamiento para nodos en la malla)

Este error ocurre cuando el compilador de XLA no puede ajustar todas las asignaciones anteriores en la HBM del dispositivo.

Depuración

Analiza cuidadosamente el mensaje de error y los registros para determinar qué categoría de HBM OOM a continuación describe mejor tu error:

• "TC Hbm usage: X, SC Hbm usage Y": Si el error desglosa explícitamente el uso de HBM, el uso agregado de TensorCore (TC) + SparseCore (SC) supera el límite de HBM. → Ir a la situación 1. Equilibra el uso de HBM de TC y SC.
• "Ran out of memory in memory space HBM": Consulta los registros para obtener una enumeración de las asignaciones más grandes en HBM.
  • En caso de que haya uno o más tensores inesperadamente grandes (p.ej., > 50% del límite de HBM) presentes → Ir a la situación 2. No hay memoria suficiente debido a asignaciones inesperadamente grandes.
  • Si no hay tensores inesperadamente grandes en los registros → Ir a la situación 3. No hay memoria suficiente debido a asignaciones agregadas.

Situación 1: Equilibra el uso de HBM de TC y SC

Si el error desglosa explícitamente el uso, p.ej., "TC Hbm usage: X, SC Hbm usage Y", significa que el uso agregado de TensorCore (TC) + SparseCore (SC) supera el límite de HBM. Compara los dos valores para identificar el cuello de botella:

• Uso alto de SparseCore
  • Optimiza el uso de la pila de HBM: El consumo de memoria de la pila de HBM se ajusta con feature_width, max_unique_nz_per_row y logical_replica_count. Puedes reducir el uso máximo de la pila ajustando la marca --xla_sc_num_serialized_tables_to_optimize_hbm, que serializa el procesamiento de tablas. Esto se produce a costa de un paralelismo reducido.
  • Verifica la sobrecarga de relleno: SparseCore alinea las tablas de incorporación a 32B (8 flotantes). Las tablas con anchos de características pequeños (p.ej., < 8 flotantes) generan una sobrecarga de relleno significativa, lo que desperdicia HBM.
  • Reduce el uso del montón: Los valores altos para maximum_parallel_iterations aumentan la cantidad de datos de entrada que se capturan previamente en el montón de HBM. Si disminuyes este valor, puedes liberar una cantidad significativa de memoria.
  • Verifica la fragmentación: Asegúrate de que las tablas de incorporación estén correctamente fragmentadas en todos los chips. Consulta Cómo se traducen los límites a las tablas.
  • Consulta SC: Cuellos de botella de rendimiento y memoria para obtener más ideas.
• Uso alto de TensorCore
  • Continúa con la situación 2.
• Equilibrado
  • Si ninguno es excesivo de forma individual, pero la suma es demasiado alta, estás en la capacidad del chip. Debes intentar reducir el uso de ambos componentes. Sigue las recomendaciones en las tres secciones.

Situación 2: No hay memoria suficiente debido a asignaciones inesperadamente grandes

Si observas el mensaje de error "Ran out of memory in memory space HBM" y hay una o más asignaciones inesperadamente grandes en los registros (> 50% del límite de HBM ), casi nunca es un problema de capacidad de hardware. Por lo general, es un error de configuración. Verifica la etiqueta XLA (si está presente) de las asignaciones grandes para obtener sugerencias sobre su código fuente de JAX.

• Quita los artefactos de depuración
  • El uso de jax.debug.print() en ejecuciones a gran escala puede obligar al compilador a materializar el tensor completo en HBM para transferirlo a la CPU, lo que interrumpe la fusión y aumenta el uso máximo de memoria. Quita cualquier jax.debug.print() restante.
• Corrige las formas de malla o la fragmentación ineficientes
  • Las formas de malla incorrectas o las anotaciones de fragmentación faltantes pueden hacer que el compilador use la replicación de forma predeterminada, lo que obliga al compilador a intentar ajustar tensores muy grandes en un solo chip.
  • Verifica las formas de las asignaciones grandes y comprueba que XLA especifique y propague correctamente la fragmentación.

Situación 3: No hay memoria suficiente debido a asignaciones agregadas

Si observas el mensaje de error "Ran out of memory in memory space HBM" y no hay tensores inesperadamente grandes en los registros, el programa se queda sin capacidad debido a que la suma agregada de las asignaciones supera el límite de HBM. En este caso, suele ser útil visualizar el perfil de memoria para identificar los búferes específicos que contribuyen al uso máximo. Consulta Cómo depurar errores de OOM con XProf para obtener una guía paso a paso para identificar los colaboradores de memoria máxima.

Una vez que hayas identificado algunos de los principales colaboradores, sigue los pasos que se indican a continuación para optimizar el espacio de memoria.

Situación 3.A Ajusta la configuración

A menudo, puedes resolver los errores de OOM con estos ajustes de configuración:

• Reduce el tamaño del lote: La memoria necesaria para las activaciones y los gradientes intermedios es directamente proporcional al tamaño del lote. Reducir el tamaño del lote puede ayudar a reducir el uso de memoria.
• Dona búferes de entrada: Cuando uses jax.jit, especifica donate_argnums para los parámetros de tu modelo. Esto permite que XLA sobrescriba la memoria de entrada con la salida.
• Habilita la precisión mixta (bfloat16): Usa bfloat16 o cuantización (int8, etc.) para los tensores más grandes del programa si la arquitectura del modelo y los requisitos de calidad lo permiten. Ten en cuenta que este cambio puede afectar el comportamiento del modelo y debe considerarse cuidadosamente.

Situación 3.B Optimiza la arquitectura y la fragmentación

Si los cambios de configuración no son suficientes, es posible que la topología del modelo sea demasiado grande para la configuración de hardware actual.

• Usa generaciones de TPU más recientes: Por lo general, las TPU más recientes ofrecen más HBM por chip. Cambia a generaciones de TPU más recientes si están disponibles.
• Ejecuta en una topología de chip más grande: Si los pesos del modelo son demasiado grandes para la topología existente, puedes intentar fragmentarlos en más chips.
• Implementa técnicas de fragmentación avanzadas:
  • Explora enfoques más avanzados de paralelismo de datos, tensores o canalizaciones.
  • Especifica sugerencias de fragmentación para valores y salidas intermedios.
• Usa la descarga del host de JAX: Las técnicas de descarga del host permiten que el usuario descargue tensores grandes a la memoria de la CPU del host (p.ej., descarga de activación y descarga de estado del optimizador).

Situación 3.C Verifica el relleno y la alineación de tensores

Las formas de tensores ineficientes son una causa común y silenciosa de los errores de OOM en las TPU. Para obtener el máximo rendimiento en las TPU, XLA rellena las dimensiones de los tensores, por lo general, a múltiplos de 128 para la dimensión más pequeña y 8 para la segunda más pequeña. Este relleno afecta los arrays de entrada y los tensores intermedios (temporales de HLO), lo que puede aumentar significativamente el uso de memoria, en especial con tamaños de dimensión pequeños. Consulta Diseños de arrays.

• Audita las formas de los búferes grandes: (En TPU v5 con diseños predeterminados)
  • Si te desplazas sobre un búfer en el lector de memoria de Xprof , aparecerá la tarjeta de detalles del búfer, que contiene detalles del búfer, incluida la información de relleno.
  • Ejemplo: Una forma de (129, 1024) se puede rellenar a (256, 1024), lo que genera casi un 50% de desperdicio de memoria.
  • Corrección: Una forma de (128, 1024) no requiere relleno y genera un 0% de desperdicio de memoria.
• Alinea las dimensiones: Asegúrate de que todas las dimensiones de tensores grandes (tamaño del lote, dimensión de incorporación, tamaño oculto) sean múltiplos de 128. Ten en cuenta que este cambio puede afectar el comportamiento del modelo y debe considerarse cuidadosamente.

Situación 3.D Ajusta las marcas de XLA clave que afectan la memoria

Se pueden ajustar las marcas de memoria clave para compensar el rendimiento por un uso de memoria más bajo. Sin embargo, esta estrategia debe usarse como último recurso, ya que puede afectar negativamente el rendimiento.

Situación 3.E Ajusta el pase de rematerialización de XLA o el punto de control manual

Si el modelo está cerca de ajustarse a la memoria, puedes usar el jax.checkpoint con jax.grad para controlar manualmente qué elementos intermedios se guardan en el pase directo en comparación con los que se vuelven a calcular en el pase inverso, lo que intercambia ciclos de procesamiento por HBM.

Como alternativa, puedes forzar el pase XLA::Rematerialization para priorizar el ahorro de memoria, posiblemente a costa de compilaciones más lentas:

Ten en cuenta que los cambios en las marcas de XLA deben usarse como último recurso, ya que pueden afectar negativamente el rendimiento.

Situación 3.F Usa herramientas de generación de perfiles avanzadas

Cómo depurar errores de OOM con XProf proporciona un instructivo sobre el uso del lector de memoria de XProf para visualizar la vista del compilador del uso de HBM.

Esta herramienta te permite ver la asignación máxima de memoria y los ciclos de vida del búfer, lo que es fundamental para comprender exactamente qué consume HBM en el punto de uso máximo. Para la configuración general de la generación de perfiles, consulta Comienza a usar Xprof y la generación de perfiles de TensorBoard.