Категория: Время компиляции: ОШИБКА В памяти Vmem в пределах области видимости
Эта ошибка указывает на то, что программе требуется больше памяти Scoped Vector Memory (Vmem), чем было выделено.
Примеры сообщений об ошибках:
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(...) ...
XLA Backends: TPU
Обзор
TPU обладают векторной памятью (VMEM), которая представляет собой локальную временную память, используемую исключительно ядром TensorCore (TC). Компилятор управляет Vmem для различных типов выделения памяти:
- Выделение памяти, ограниченное областью действия инструкции: временное хранение в Vmem во время выполнения одной инструкции HLO. Это включает буфер диапазона операндов (например, для двойной буферизации) и выгрузку регистров.
- Выделение памяти, ограниченное областью действия программы: выделение памяти, выходящее за рамки одной инструкции HLO. Обычно это временные ресурсы HLO и промежуточные результаты, являющиеся входными и/или выходными данными инструкций HLO.
Ошибка нехватки памяти (OOM) в памяти, ограниченной областью действия инструкции, возникает, когда выделение памяти, ограниченное областью действия инструкции, превышает лимит выделения памяти для этой инструкции. Этот лимит контролируется.
- Глобально для всей программы с помощью флага --xla_tpu_scoped_vmem_limit_kib и
- для каждого пользовательского ядра через параметр vmem_limit_bytes .
Эти ошибки обычно вызваны внутренней ошибкой компилятора или превышением лимита выделения памяти пользовательским ядром.
Отладка
Внимательно проанализируйте сообщение об ошибке, чтобы определить, вызвана ли она пользовательским ядром или стандартным HLO. Ошибка, вызванная пользовательским ядром, должна иметь следующую сигнатуру:
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) в пользовательском ядре : Если ошибка указывает на пользовательское ядро → Перейти к перенастройке ядра .
- Проблемы с Vmem, не связанные с ядром : Если ошибка нехватки памяти (OOM) возникает из-за операции, не связанной с пользовательским ядром, это, вероятно, внутренняя ошибка компилятора. Пожалуйста, сообщите об ошибке в XLA, приложив дамп HLO.
Перенастройте ядро
Если ошибка вызвана пользовательским ядром, используйте следующие методы для уменьшения требований ядра к памяти:
- Настройка размеров блоков: Уменьшите размеры блоков (размеры плиток) в конфигурации ядра, чтобы снизить использование виртуальной памяти в рамках области видимости.
- Установите ограничения на использование виртуальной памяти для каждого ядра: явно запросите необходимый объем памяти для конкретного ядра, используя параметр vmem_limit_bytes.
- Изменение цветовой маркировки памяти: Явно задайте цветовую маркировку/ограничение ввода/вывода ядра для VMEM с помощью pallas.tpu.with_memory_space_constraint . Однако будьте осторожны, не закрашивайте слишком много входов и выходов для Vmem, так как это может привести к общей ошибке нехватки памяти (OOM) для VMEM.
- Если настройка для конкретного ядра затруднена или проблема затрагивает множество ядер, вы можете отрегулировать глобальное ограничение Vmem с помощью флага --xla_tpu_scoped_vmem_limit_kib .