Код ошибки: E1001

Категория: Время компиляции: ОШИБКА В памяти 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) в памяти, ограниченной областью действия инструкции, возникает, когда выделение памяти, ограниченное областью действия инструкции, превышает лимит выделения памяти для этой инструкции. Этот лимит контролируется.

и

Эти ошибки обычно вызваны внутренней ошибкой компилятора или превышением лимита выделения памяти пользовательским ядром.

Допустимые лимиты {#usable_limits}

Различные аппаратные компоненты TPU имеют разные размеры Vmem. Кроме того, различное оборудование может резервировать разные объемы Vmem для различных задач, тем самым ограничивая объем, который может быть выделен для Scoped Vmem.

Ниже приведены типичные ограничения для различного оборудования.

Аппаратное обеспечение Ограничение (удобное для человека) Лимит (КиБ) Лимит (байты)
ТПУ v2 16 МБ 16384 16777216
ТПУ v3 16 МБ 16384 16777216
TPU v4 16 МБ 16384 16777216
TPU v4i 16 МБ 16384 16777216
TPU v5e 128 МБ 131072 134217728
ТПУ v5p 63,9 МБ 65472 67043328
TPU v6e 127,9 МБ 131008 134152192
ТПУ 7x 63,9 МБ 65472 67043328

Отладка

Внимательно проанализируйте сообщение об ошибке, чтобы определить, вызвана ли она пользовательским ядром или стандартным 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 OOM) в пользовательском ядре : Если ошибка указывает на пользовательское ядро ​​→ Перейти к перенастройке ядра .
  • Проблемы с Vmem, не связанные с ядром : Если ошибка нехватки памяти (OOM) возникает из-за операции, не связанной с пользовательским ядром, это, вероятно, внутренняя ошибка компилятора. Пожалуйста, отправьте отчет об ошибке в XLA с дампом HLO.

Перенастройте ядро

Если ошибка вызвана пользовательским ядром, используйте следующие методы для уменьшения требований ядра к памяти:

  • Настройка размеров блоков: Уменьшите размеры блоков (размеры плиток) в конфигурации ядра, чтобы снизить использование виртуальной памяти в рамках области видимости.
  • Установите ограничения на использование виртуальной памяти для каждого ядра: явно запросите необходимый объем памяти для конкретного ядра, используя параметр vmem_limit_bytes.
  • Изменение цветовой маркировки памяти: Явно задайте цветовую маркировку/ограничение ввода/вывода ядра для Vmem с помощью pallas.tpu.with_memory_space_constraint . Будьте осторожны, не закрашивайте слишком много входных и выходных сигналов для Vmem, так как это может привести к общей ошибке нехватки памяти (OOM) для Vmem.
  • Настройка ограничения Vmem: Если перенастройка для конкретного ядра затруднена или проблема затрагивает множество ядер, вы можете настроить глобальное ограничение Vmem с помощью флага --xla_tpu_scoped_vmem_limit_kib .