Код ошибки: 0101

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

Тип: Среда выполнения

Пример журнала ошибок

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