Categoría: Tiempo de compilación: OOM de Vmem con alcance
Este error indica que el programa requiere más memoria vectorial con alcance (Vmem) de la que se asignó.
Ejemplos de mensajes de error:
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(...) ...
Back-ends de XLA: TPU
Descripción general
Las TPU tienen memoria vectorial (VMEM), que es una memoria de trabajo local que usa exclusivamente el TensorCore (TC). El compilador administra la Vmem para diferentes tipos de asignaciones:
- Asignaciones con alcance de instrucción: Almacenamiento temporal en Vmem durante la ejecución de una sola instrucción de HLO. Esto incluye el búfer de intervalo de operandos (p.ej., para el doble búfer) y las escrituras en registros.
- Asignaciones con alcance de programa: Son asignaciones que van más allá del alcance de una sola instrucción de HLO. Por lo general, se trata de resultados intermedios y temporales de HLO que son entradas o salidas de las instrucciones de HLO.
Se produce un error de OOM de Vmem 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 está controlado
- de forma global para todo el programa a través de la marca --xla_tpu_scoped_vmem_limit_kib
- por kernel personalizado a través del parámetro vmem_limit_bytes
Por lo general, estos errores se deben a un error interno del compilador o a un kernel personalizado que supera su límite de asignación.
Depuración
Analiza cuidadosamente el mensaje de error para identificar si este 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" ...
- Custom Kernel Scoped Vmem OOM: Si el error apunta a un kernel personalizado, ve a Retune the Kernel.
- Problemas de Vmem que no son del kernel: Si el error de OOM de Vmem se produce debido a una operación que no es del kernel personalizado, es probable que se trate de un error interno del compilador. Notifica un error en XLA con un volcado de HLO.
Cómo volver a ajustar 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 de forma explícita 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 de forma explícita las entradas o salidas del kernel a la VMEM con pallas.tpu.with_memory_space_constraint. Sin embargo, ten cuidado de no colorear demasiadas entradas y salidas en Vmem, ya que esto podría provocar un error general de OOM de VMEM.
- Si es difícil volver a ajustar un kernel específico o el problema afecta a muchos kernels, puedes ajustar el límite global de Vmem con la marca --xla_tpu_scoped_vmem_limit_kib.