Código de error: 0101

Categoría: Tiempo de ejecución: Error en la asignación del programa

Este error indica que el tiempo de ejecución de XLA en un dispositivo TPU no pudo cargar un ejecutable de programa XLA compilado en la HBM de la TPU.

Ejemplo de mensaje de error:

XlaRuntimeError: RESOURCE_EXHAUSTED: Error loading program 'jit_embedding_pipeline_step_fn': Attempting to reserve 29.49G at the bottom of memory. That was not possible. There are 147.64M free, 0B reserved, and 147.64M reservable. Scope: unknown..: while running replica 0 and partition 34 of a replicated computation (other replicas may have failed as well).

Back-ends de XLA: TPU

Descripción general

Por lo general, este error se debe a uno de los siguientes motivos:

• El tamaño del programa supera la HBM disponible: El programa XLA compilado, incluidas sus instrucciones, datos estáticos y cualquier constante incorporada, es más grande que la cantidad total de HBM libre disponible actualmente en los núcleos de TPU específicos en los que se está cargando el programa.
• Fragmentación de la HBM: Si bien la HBM libre total del dispositivo podría ser suficiente en conjunto, no está disponible en un solo bloque contiguo lo suficientemente grande como para que quepa todo el programa.

Es importante comprender cómo el tiempo de ejecución de la TPU prioriza la memoria. Las asignaciones de búfer tienen prioridad sobre los programas cargados. Si falla la asignación de un búfer, el tiempo de ejecución expulsará de la HBM los programas ya cargados para liberar espacio. Esto puede generar una situación en la que un programa que se cargó correctamente antes ahora falla con un error de OOM, ya que la HBM ahora está ocupada con más búferes de datos.

Depuración

• Reduce la huella de memoria del búfer: Liberar la memoria que usan los búferes de datos dejará más espacio para el programa en sí:
  • Disminuye el tamaño del lote: Esta es una de las formas más eficaces de reducir la cantidad de memoria que se usa para las activaciones.
  • Fragmentación de parámetros: Para modelos muy grandes, usa técnicas de paralelismo o fragmentación de modelos (como FSDP o Megascale) para distribuir los parámetros y el procesamiento del modelo en varios hosts o núcleos de TPU.
  • Acortar la longitud de la secuencia o el contexto: Para los modelos que procesan datos secuenciales (p.ej., En el caso de los modelos de PNL, reducir la longitud de la secuencia puede disminuir significativamente el uso de memoria.
  • Donación de búfer: Usa funciones del framework (p.ej., jax.jit(..., donate_argnums=...)) para permitir que XLA reutilice la memoria de los búferes de entrada para almacenar la salida, lo que reduce el uso máximo de memoria.
• Reduce los requisitos de memoria del programa para los temporales:
  • Reduce el uso de memoria de los programas para los temporales con la marca tpu_shared_memory_percent. Ten en cuenta que esto podría afectar negativamente el rendimiento.
• Optimiza la estrategia de ejecución y reduce la carga de publicación:
  • Administra la carga del programa: Si compilas varias funciones con JIT, ten en cuenta que cada función puede generar la carga de un programa. Intenta estructurar tu carga de trabajo para minimizar la cantidad de programas cargados de forma simultánea.
• Asegúrate de que no haya fugas de memoria:
  • Asegúrate de que las referencias a objetos jax.Array no se mantengan durante más tiempo del previsto. Mantener objetos jax.Array puede impedir la desasignación automática incluso después de que se complete la compilación del programa.

Herramientas

• Habilita la marca tpu_log_allocations_on_oom para la que el asignador volcará un informe detallado de todas las asignaciones actuales cuando se produzca un error de OOM, lo que puede ser muy valioso para la depuración.
• Genera perfiles de tu programa: Usa el generador de perfiles de memoria de JAX o el de TensorFlow para obtener una vista detallada del uso de memoria de tu programa a lo largo del tiempo. Esto puede ayudar a identificar picos inesperados en el consumo de memoria.