Код ошибки: 0100

Категория: Выполнение: Ошибка выделения буфера

Эта ошибка указывает на то, что распределитель памяти среды выполнения XLA:TPU не смог найти подходящий блок памяти на HBM ускорителя для запрошенного выделения.

Пример сообщения об ошибке:

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

Бэкенды XLA: TPU

Обзор

Эта ошибка возникает в

  • сбои при инициированном пользователем выделении буфера через jax.device_put или
  • сбои при выделении выходных ресурсов программой, запланированной пользователем.

Эти сбои обычно вызваны несколькими причинами:

  • Ошибка нехватки памяти (OOM): Пользователь пытается выделить фрагмент памяти, размер которого превышает общий объем свободной памяти, доступной в памяти HBM процессора TPU.
  • Фрагментация памяти: выделение памяти не удается, поскольку ни один отдельный непрерывный свободный блок в адресном пространстве памяти не достаточно велик, чтобы удовлетворить запрошенный размер. Общий объем свободной памяти достаточен для выделения, но он разбросан по адресному пространству памяти в виде небольших, несмежных блоков.

В среде выполнения TPU реализован ряд механизмов для повторной попытки выделения памяти в случае сбоя, в том числе:

  • Если в очереди находятся операции освобождения памяти, среда выполнения повторяет попытки выполнения неудачных операций выделения памяти.
  • При ошибках нехватки памяти, вызванных фрагментацией, среда выполнения может автоматически запустить дефрагментацию и повторную попытку.
  • В среде выполнения TPU приоритет отдается выделению буферов, а не поддержанию загруженных программ. Если выделение буфера не удается из-за недостатка памяти HBM, система вытеснит загруженные программы TPU до тех пор, пока не будет достаточно памяти для буфера.

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

Отладка

  • Уменьшите объем памяти, занимаемый вашей моделью:
    • Уменьшение размера пакета: Уменьшение размера пакета напрямую снижает потребление памяти.
    • Распределение параметров: Для очень больших моделей используйте такие методы, как параллельная обработка моделей или распределение параметров, чтобы распределить параметры по HBM нескольких ядер или хостов TPU.
    • Сокращение длины последовательности/контекста: Для моделей, работающих с последовательностями (например, языковых моделей), сокращение длины входной последовательности может значительно уменьшить объем используемой памяти.
    • Передача буферов: Используйте возможности фреймворка (например, jax.jit(..., donate_argnums=...) ) для того, чтобы сообщить XLA, что определенные входные буферы могут быть перезаписаны и повторно использованы для выходных данных.
    • Оптимизация стратегии контрольных точек: вместо сохранения всего состояния модели сразу, рассмотрите возможность сохранения только весов модели или использования стратегии сегментированного сохранения контрольных точек.
  • Расположение и заполнение адресов в памяти:
    • Память TPU выделяется блоками, а добавление отступов может увеличить фактический размер тензоров.
  • Убедитесь в отсутствии утечек памяти:
    • Убедитесь, что ссылки на объекты jax.Array не удерживаются дольше, чем предполагалось. Удержание объектов jax.Array может препятствовать автоматическому освобождению памяти даже после завершения компиляции программы.

Инструменты

Включите флаг tpu_log_allocations_on_oom который позволит распределителю памяти выводить подробный отчет обо всех текущих выделениях памяти при возникновении ошибки нехватки памяти (OOM), что может быть крайне полезно для отладки.