Código de error: E1001

Categoría: Tiempo de compilación: Vmem OOM con alcance

Este error indica que el programa requiere más memoria vectorial con alcance (Vmem) de la que se asignó.

Mensajes de error de muestra:

RESOURCE_EXHAUSTED: Ran out of memory in memory space vmem while allocating on stack for %my-custom-kernel = bf16[2048,4096]{1,0:T(8,128)(2,1)} custom-call(...) ...

Backends de XLA: TPU

Descripción general

Las TPU tienen memoria vectorial (VMEM), que es una memoria de borrador local que usa exclusivamente TensorCore (TC). El compilador administra Vmem para diferentes tipos de asignaciones:

  • Asignaciones con alcance de instrucción: Almacenamiento temporal en Vmem mientras se ejecuta una sola instrucción HLO. Esto incluye el búfer de intervalo de operandos (p.ej., para el almacenamiento en búfer doble) y los desbordamientos de registros.
  • Asignaciones con alcance de programa: Asignaciones que existen más allá del alcance de una sola instrucción HLO. Por lo general, son resultados intermedios y temporales de HLO que son entradas o salidas de instrucciones HLO.

Se produce un error de Vmem OOM con alcance de tiempo de compilación cuando las asignaciones con alcance de instrucción superan el límite de asignación para esa instrucción. Este límite se controla

y

Por lo general, estos errores se deben a un error interno del compilador o a que un kernel personalizado supera su límite de asignación.

Depuración

Analiza cuidadosamente el mensaje de error para identificar si el error proviene de un kernel personalizado o de un HLO estándar. Un error debido a un kernel personalizado debe tener la siguiente firma:

Ran out of memory in memory space vmem while allocating on stack for %my-custom-call = <output-shape> custom-call(<params>), custom_call_target="tpu_custom_call" ...
  • Vmem OOM con alcance de kernel personalizado: Si el error apunta a un kernel personalizado, ve a Reajusta el kernel.
  • Problemas de Vmem sin kernel: Si el error de Vmem OOM se produce debido a una operación que no es de kernel personalizado, es probable que se trate de un error interno del compilador. Envía un informe de errores en XLA con un volcado de HLO.

Reajusta el kernel

Si el error se origina en un kernel personalizado, usa las siguientes técnicas para reducir el requisito de memoria del kernel:

  • Ajusta los tamaños de bloque: Reduce los tamaños de bloque (tamaños de mosaico) en la configuración del kernel para disminuir el uso de Vmem con alcance.
  • Establece límites de Vmem con alcance por kernel: Solicita explícitamente la cantidad de memoria requerida para ese kernel específico con el parámetro vmem_limit_bytes
  • Modifica la coloración de la memoria: Colorea o restringe explícitamente las entradas o salidas del kernel a Vmem con pallas.tpu.with_memory_space_constraint. Ten cuidado de no colorear demasiadas entradas o salidas a Vmem, ya que eso podría causar un error general de Vmem OOM.
  • Ajusta el límite de Vmem: Si el reajuste específico del kernel es difícil o el problema afecta a muchos kernels, puedes ajustar el límite global de Vmem con la marca --xla_tpu_scoped_vmem_limit_kib.