Código de error: 0100

Categoría: Tiempo de ejecución: Error de asignación de búfer

Este error indica que el asignador de memoria del tiempo de ejecución de XLA:TPU no pudo encontrar un bloque de memoria adecuado en la HBM del acelerador para la asignación solicitada.

Ejemplo de mensaje de error:

ValueError: RESOURCE_EXHAUSTED: Error allocating device buffer: Attempting to allocate 8.00M. That was not possible. There are 6.43M free.; (0x0x1_HBM0)

Back-ends de XLA: TPU

Descripción general

Este error se arroja en

• fallas en la asignación de búferes iniciada por el usuario a través de jax.device_put o
• fallas en las asignaciones de salida de los programas programados por el usuario.

Por lo general, estas fallas se deben a los siguientes motivos:

• Memoria insuficiente (OOM): El usuario intenta asignar un fragmento de memoria que es más grande que la cantidad total de memoria libre disponible en la HBM de la TPU.
• Fragmentación de la memoria: La asignación falla porque no hay un solo bloque libre contiguo en el espacio de memoria que sea lo suficientemente grande como para satisfacer el tamaño solicitado. La cantidad total de memoria libre es suficiente para la asignación, pero está dispersa en el espacio de memoria en bloques pequeños y no contiguos.

El tiempo de ejecución de la TPU tiene varios mecanismos para reintentar las fallas de asignación, incluidos los siguientes:

• Si hay desasignaciones en cola, el tiempo de ejecución reintenta las asignaciones fallidas.
• En los errores de OOM causados por una fragmentación, el tiempo de ejecución puede activar automáticamente una desfragmentación y un reintento.
• El tiempo de ejecución de la TPU prioriza las asignaciones de búfer sobre el mantenimiento de los programas cargados. Si falla la asignación de un búfer debido a que la HBM es insuficiente, el sistema desalojará los programas de TPU cargados hasta que haya suficiente memoria disponible para el búfer.

Por lo tanto, los errores que se producen después de las mitigaciones anteriores suelen requerir la acción del usuario.

Depuración

• Reduce la huella de memoria de tu modelo:
  • Disminuye el tamaño del lote: Reducir el tamaño del lote disminuye directamente el uso de memoria.
  • Fragmentación de parámetros: Para modelos muy grandes, usa técnicas como el paralelismo de modelos o la fragmentación para distribuir los parámetros en la HBM de varios núcleos o hosts de TPU.
  • Acorta la longitud de la secuencia o el contexto: En el caso de los modelos que operan en secuencias (como los modelos de lenguaje), reducir la longitud de la secuencia de entrada puede disminuir significativamente el alcance de memoria.
  • Donación de búfer: Utiliza funciones del framework (como jax.jit(..., donate_argnums=...)) para indicarle a XLA que ciertos búferes de entrada se pueden reemplazar y reutilizar para las salidas.
  • Estrategia de optimización de puntos de control: En lugar de guardar el estado completo del modelo de una vez, considera guardar solo los pesos del modelo o usar una estrategia de puntos de control fragmentados.
• Diseño y relleno de la memoria de direcciones:
  • La memoria de la TPU se asigna en fragmentos, y el relleno puede aumentar el tamaño real de los tensores.
• 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.