Код ошибки: 0101

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

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

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

XlaRuntimeError: RESOURCE_EXHAUSTED: Error loading program 'jit_embedding_pipeline_step_fn': Attempting to reserve 29.49G at the bottom of memory. That was not possible. There are 147.64M free, 0B reserved, and 147.64M reservable. Scope: unknown..: while running replica 0 and partition 34 of a replicated computation (other replicas may have failed as well).

Бэкенды XLA: TPU

Обзор

Эта ошибка обычно вызвана одной из следующих причин:

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

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

Отладка

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

Инструменты

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