Código de error: 1000

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

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

Ejemplos de mensajes de error:

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.

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

Back-ends 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 la HBM de la TPU para varios tipos de asignaciones:

• Entradas y salidas del programa: Estados del optimizador, lotes de entrenamiento, etcétera
• Variables temporales de TensorCore y SparseCore: Memoria dinámica requerida para cálculos intermedios (p. ej., activaciones, gradientes, etc.).
• Binario compilado: Es 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 de TPU anteriores).
• Constantes: Los valores constantes incorporados en el código intermedio de HLO se asignan en la 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 se produce 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 OOM de HBM que se indica a continuación describe mejor tu error:

• TensorCore (TC) + SparseCore (SC) HBM Usage Exceeds Limit: Si el error desglosa explícitamente el uso, p.ej., "Uso de HBM de TC: X, uso de HBM de SC: Y". → Ir a la Sección 1 Balancea el uso de HBM de TC y SC.
• Asignaciones inesperadamente grandes: Si el error indica "Se agotó la memoria en el espacio de memoria HBM", revisa los registros para ver 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., más del 50% del límite de HBM), ve a la sección 2. Asignaciones Inesperadamente Grandes.
• Aggregate Allocations Exceed HBM Limit: Si el error indica "Ran out of memory in memory space HBM", pero no hay tensores inesperadamente grandes en los registros, ve a la Sección 3. Las asignaciones agregadas superan el límite de HBM.

Sección 1: Balancea el uso de la HBM de TC y SC

Si el error desglosa explícitamente el uso, p.ej., "TC Hbm usage: X, SC Hbm usage Y" compara los dos valores para identificar el cuello de botella.

• Uso elevado 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 las tablas. Esto se logra a costa de reducir el paralelismo.
  • Verifica la sobrecarga de padding: SparseCore alinea las tablas de incorporación con 32 B (8 números de punto flotante). Tablas con anchos de columna pequeños (p.ej., < 8 flotantes) generan una sobrecarga de padding significativa, lo que desperdicia HBM.
  • Reduce Heap Usage: Los valores altos de maximum_parallel_iterations aumentan la cantidad de datos de entrada que se obtienen previamente en el montón de HBM. Reducir este valor puede liberar una cantidad significativa de memoria.
  • Verifica la fragmentación: Asegúrate de que las tablas de incorporación se fragmenten correctamente en todos los chips. Consulta Cómo se aplican los límites a las tablas.
  • Consulta SC: Cuellos de botella de rendimiento y memoria para obtener más ideas.
• Uso elevado de TensorCore:
  • Continúa con la Secció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 de las tres secciones.

Sección 2: Asignaciones inesperadamente grandes

Si hay una o más asignaciones inesperadamente grandes en los registros (más del 50% del límite de HBM), casi nunca se trata de un problema de capacidad de hardware. Por lo general, se trata de 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:
  • Usar jax.debug.print() en ejecuciones a gran escala puede obligar al compilador a materializar el tensor completo en la HBM para transferirlo a la CPU, lo que interrumpe la fusión y aumenta el uso máximo de memoria. Quita los jax.debug.print() restantes.
• Corrige las formas de malla o el fragmentado ineficientes:
  • Las formas de malla incorrectas o las anotaciones de fragmentación faltantes pueden hacer que el compilador use Replication de forma predeterminada, lo que lo obliga 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 el sharding.

Sección 3: Las asignaciones agregadas exceden el límite de HBM

Si el programa se queda sin capacidad debido a que la suma agregada de las asignaciones supera el límite de HBM, 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 principales factores que contribuyen al uso máximo de memoria.

Una vez que hayas identificado a algunos de los principales colaboradores, sigue estos pasos para optimizar el uso de memoria.

A. Verifica el relleno y la alineación del tensor

Las formas de tensor ineficientes son una causa común y silenciosa de errores de OOM en las TPU. Para obtener el máximo rendimiento en las TPU, XLA agrega relleno a las dimensiones del tensor, por lo general, a múltiplos de 128 para la dimensión más secundaria y de 8 para la segunda más secundaria. Este padding afecta tanto a los arrays de entrada como a 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 matrices.

• Audita las formas de los búferes grandes: (en TPU v5 con diseños predeterminados)
  • Si te desplazas sobre un búfer en Xprof Memory Viewer, aparecerá la tarjeta de detalles del búfer, que contiene información sobre el búfer, incluidos los datos de padding.
  • Ejemplo: Una forma de (129, 1024) podría rellenarse hasta (256, 1024), lo que generaría un desperdicio de memoria de casi el 50%.
  • Corrección: Una forma de (128, 1024) no requiere padding y genera un desperdicio de memoria del 0%.
• Alinear 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.

B. Ajusta la configuración

A menudo, puedes resolver los errores de OOM con los siguientes 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 a menudo puede ayudar a reducir el uso de memoria.
• Donate Input Buffers: 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 la cuantificación (int8, etc.) para los tensores más grandes del programa si la arquitectura del modelo y los requisitos de calidad lo permiten.

C. Optimiza la arquitectura y el sharding

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 chips 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 avanzadas de fragmentación:
  • Explora enfoques más avanzados de paralelismo de datos, tensores o canalizaciones.
  • Especifica sugerencias de fragmentación para los valores y resultados intermedios.
• Usa la descarga del host de JAX: Descarga tensores grandes en la memoria de la CPU del host, p. ej., descarga de la activación y descarga del estado del optimizador.

D. Ajusta las marcas de XLA clave que afectan la memoria:

Las marcas de memoria clave se pueden ajustar para intercambiar rendimiento por un menor uso de memoria. Sin embargo, se deben usar como último recurso, ya que pueden afectar negativamente el rendimiento.

E. Ajusta el pase de rematerialización de XLA / puntos de control manuales

Si el modelo está cerca de caber en la memoria, puedes forzar el paso XLA::Rematerialization para priorizar el ahorro de memoria, posiblemente a costa de compilaciones más lentas:

Como alternativa, usa el decorador jax.checkpoint con jax.grad para controlar de forma manual qué resultados intermedios se guardan en el pase hacia adelante en comparación con los que se vuelven a calcular en el pase hacia atrás, intercambiando ciclos de procesamiento por HBM.

F. Usa herramientas avanzadas de generación de perfiles

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

Esta herramienta te permite ver la asignación máxima de memoria y la vida útil del búfer, lo que es fundamental para comprender exactamente qué consume HBM en el punto de utilización máxima. Para obtener información general sobre la configuración de la creación de perfiles, consulta Comienza a usar Xprof y TensorBoard Profiling.