Код ошибки: E0101

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

Эта ошибка указывает на то, что среда выполнения 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 может препятствовать автоматическому освобождению памяти даже после завершения компиляции программы.

См. также код ошибки E1000, где описаны другие стратегии, позволяющие уменьшить объем памяти, используемый каждой программой.

Инструменты

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