Категория: Ошибка выделения буфера — TPU
Тип: Среда выполнения
Пример журнала ошибок
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 не удалось найти подходящий блок памяти в HBM ускорителя для запрошенной операции. Эти операции обычно представляют собой выделение буфера пользователем через jax.device_put или выделение памяти для вывода программы. Эти сбои происходят по нескольким причинам: - Недостаточно памяти (OOM) - Пользователь пытается выделить фрагмент памяти, превышающий общий объём свободной памяти, доступной в HBM TPU. - Фрагментация памяти - Выделение не удаётся, поскольку ни один непрерывный свободный блок в пространстве памяти не достаточно велик для удовлетворения запрошенного размера. Общий объём свободной памяти достаточен для выделения, но она разбросана по пространству памяти небольшими, несмежными блоками.
Среда выполнения TPU имеет ряд механизмов для повторных попыток выделения памяти в случае сбоя, включая: - При наличии очередных освобождений памяти среда выполнения повторяет попытки выделения памяти, завершившиеся сбоем. - При возникновении OOM, вызванных фрагментацией, среда выполнения может автоматически запустить дефрагментацию и повторную попытку. - Среда выполнения TPU отдает приоритет выделению буфера, а не сохранению загруженных программ. Если выделение буфера не удается из-за недостаточного количества HBM, система вытесняет загруженные программы TPU до тех пор, пока для буфера не будет достаточно памяти.
Таким образом, ошибка, обнаруженная после принятия вышеуказанных мер, обычно требует действий пользователя.
Как пользователь может исправить свою программу, если такое произошло?
- Уменьшите объем памяти, занимаемый вашей моделью:
- Уменьшение размера пакета: уменьшение размера пакета напрямую снижает использование памяти.
- Разделение параметров: для очень больших моделей используйте такие методы, как параллелизм моделей или разделение, чтобы распределить параметры по HBM нескольких ядер TPU или хостов.
- Сокращение длины последовательности/контекста: для моделей, которые работают с последовательностями (например, языковых моделей), сокращение длины входной последовательности может значительно уменьшить объем памяти.
- Донорство буфера: используйте возможности фреймворка (например:
jax.jit(..., donate_argnums=...)), чтобы сигнализировать XLA о том, что определенные входные буферы могут быть перезаписаны и повторно использованы для выходных данных. - Оптимизируйте стратегию контрольных точек: вместо того, чтобы сохранять все состояние модели сразу, рассмотрите возможность сохранения только весовых коэффициентов модели или использования стратегии сегментированных контрольных точек.
- Расположение адресов памяти и заполнение:
- Память TPU выделяется порциями, а заполнение может увеличить фактический размер тензоров.
- Убедитесь в отсутствии утечек памяти:
- Убедитесь, что ссылки на объекты
jax.Arrayне хранятся дольше положенного времени. Сохранение объектовjax.Arrayможет помешать автоматическому освобождению памяти даже после завершения компиляции программы.
- Убедитесь, что ссылки на объекты
Как пользователь может устранить эти сбои?
Включите флаг tpu_log_allocations_on_oom при котором распределитель будет выдавать подробный отчет обо всех текущих выделениях при возникновении OOM, что может быть бесценно для отладки.