Код ошибки: 1000

Категория: Время компиляции: HBM OOM

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

Примеры сообщений об ошибках:

RESOURCE_EXHAUSTED: XLA:TPU compile permanent error. Ran out of memory in memory space hbm. Used 49.34G of 32.00G hbm. Exceeded hbm capacity by 17.34G.

RESOURCE_EXHAUSTED: TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G

XLA Backends: TPU

Обзор

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

Компилятор управляет фиксированной емкостью HBM процессора TPU для нескольких типов выделений памяти:

• Входные и выходные данные программы: обучающие пакеты, состояния оптимизатора и т. д.
• Временные переменные TensorCore + SparseCore: динамическая память, необходимая для промежуточных вычислений (например, активаций, градиентов и т. д.).
• Скомпилированный двоичный код: машинный код для TensorCore (TC) и SparseCore (SC).
• Системные накладные расходы: зарезервированное пространство для среды выполнения XLA (например, буферы ввода на более старых поколениях TPU).
• Константы: Постоянные значения, заложенные в HLO IR, выделяются в HBM.
• Внутреннее устройство компилятора: выделение памяти на уровне программы и для каждого HLO (например, информация о маршрутизации для узлов в сети).

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

Отладка

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

• Использование TensorCore (TC) + SparseCore (SC) HBM превышает лимит : Если в сообщении об ошибке явно указано использование, например, "Использование TC Hbm: X, использование SC Hbm: Y" . → Перейти к разделу 1. Балансировка использования TC и SC HBM .
• Неожиданно большие выделения памяти : Если ошибка гласит «Недостаточно памяти в пространстве памяти HBM» , проверьте журналы на наличие списка самых больших выделений памяти в HBM. В случае наличия одного или нескольких неожиданно больших тензоров (например, > 50% от лимита HBM) → Перейти к разделу 2. Неожиданно большие выделения памяти .
• Превышение лимита HBM при совокупном выделении памяти : Если в сообщении об ошибке указано «Исчерпана память в пространстве памяти HBM» , но в логах отсутствуют неожиданно большие тензоры → Перейти к разделу 3. Превышение лимита HBM при совокупном выделении памяти .

Раздел 1. Сбалансируйте использование HBM для TC и SC.

Если в сообщении об ошибке явно указано использование, например, "TC Hbm usage: X, SC Hbm usage Y", сравните эти два значения, чтобы определить узкое место.

• Высокая загрузка SparseCore:
  • Оптимизация использования стека HBM: потребление памяти стеком HBM масштабируется в зависимости от feature_width , max_unique_nz_per_row и logical_replica_count . Вы можете уменьшить пиковое использование стека, настроив флаг --xla_sc_num_serialized_tables_to_optimize_hbm , который сериализует обработку таблиц. Это происходит за счет снижения параллелизма.
  • Проверьте избыточность заполнения: SparseCore выравнивает таблицы встраивания по 32 байтам (8 чисел с плавающей запятой). Таблицы с малой шириной признаков (например, < 8 чисел с плавающей запятой) приводят к значительной избыточности заполнения, что нерационально использует HBM.
  • Снижение использования кучи: высокие значения параметра maximum_parallel_iterations увеличивают объем входных данных, предварительно загружаемых в кучу HBM. Снижение этого значения может освободить значительное количество памяти.
  • Проверка сегментирования: Убедитесь, что таблицы встраивания правильно сегментированы по всем чипам. См. раздел «Как ограничения преобразуются в таблицы» .
  • Ознакомьтесь с материалом SC: «Узкие места в производительности и памяти» для получения дополнительных идей.
• Высокая загрузка TensorCore:
  • Перейдите к разделу 2 .
• Сбалансированный
  • Если ни один из компонентов по отдельности не является избыточным, но их сумма слишком велика, значит, вы достигли предела возможностей чипа. Необходимо попытаться снизить использование обоих компонентов. Следуйте рекомендациям во всех трех разделах.

Раздел 2. Неожиданно большие ассигнования

Если в логах присутствует одно или несколько неожиданно больших выделений памяти (> 50% от лимита HBM), это почти никогда не связано с аппаратной емкостью. Обычно это ошибка конфигурации. Проверьте метку XLA (если она присутствует) больших выделений памяти, чтобы получить подсказки об их исходном коде JAX.

• Удаление артефактов отладки:
  • Использование jax.debug.print() в масштабируемых приложениях может заставить компилятор материализовать весь тензор в HBM для передачи его на ЦП, что нарушает слияние данных и увеличивает пиковое использование памяти. Удалите все оставшиеся вызовы jax.debug.print() .
• Исправление неэффективных форм сетки или сегментирования:
  • Неправильная форма сетки или отсутствие аннотаций сегментирования могут привести к тому, что компилятор по умолчанию будет использовать режим репликации , заставляя его пытаться разместить очень большие тензоры на одном чипе.
  • Проверьте структуру больших выделенных ресурсов и убедитесь, что сегментирование задано и распространяется XLA корректно.

Раздел 3. Совокупные ассигнования превышают лимит HBM.

Если программа испытывает нехватку памяти из-за превышения суммарного объема выделений лимита HBM, часто бывает полезно визуализировать профиль использования памяти, чтобы определить конкретные буферы, влияющие на пиковое потребление. См. раздел «Отладка ошибок нехватки памяти с помощью XProf» для пошагового руководства по определению источников пикового потребления памяти.

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

А. Проверьте заполнение и выравнивание тензора.

Неэффективные формы тензоров являются распространенной, скрытой причиной ошибок нехватки памяти (OOM) на TPU. Для достижения максимальной производительности на TPU, XLA дополняет размеры тензоров — обычно до значений, кратных 128 для самого младшего измерения и 8 для второго младшего. Это дополнение влияет как на входные массивы, так и на промежуточные тензоры (временные HLO), потенциально значительно увеличивая использование памяти, особенно при малых размерах измерений. См. раздел «Расположение массивов» .

• Проверка структуры больших буферов: (На TPU v5 со стандартной компоновкой)
  • При наведении курсора на буфер в программе Xprof Memory Viewer появляется карточка с подробными сведениями о буфере, включая информацию о заполнении.
  • Пример : Формат (129, 1024) может быть дополнен до (256, 1024) , что приведет к почти 50% неэффективному использованию памяти.
  • Исправление: Форма (128, 1024) не требует заполнения и не приводит к 0% расходованию памяти.
• Выравнивание размеров: Убедитесь, что все большие размеры тензора (размер пакета, размерность встраивания, размер скрытого слоя) кратны 128.

B. Настройка конфигурации

Часто проблему нехватки памяти можно решить с помощью следующих настроек конфигурации:

• Уменьшите размер пакета: объем памяти, необходимый для промежуточных активаций и градиентов, прямо пропорционален размеру пакета. Уменьшение размера пакета часто помогает снизить потребление памяти.
• Передача входных буферов: При использовании jax.jit укажите donate_argnums для параметров вашей модели. Это позволит XLA перезаписывать входную память выходными данными.
• Включить смешанную точность (bfloat16): Используйте bfloat16 или квантование (int8 и т. д.) для самых больших тензоров в программе, если это позволяют архитектура модели и требования к качеству.

C. Оптимизация архитектуры и сегментирования.

Если изменений в конфигурации недостаточно, топология модели может оказаться слишком большой для текущей аппаратной конфигурации.

• Используйте более новые поколения TPU: более новые TPU, как правило, предлагают больший объем памяти HBM на чип; переходите на более новые поколения TPU, если они доступны.
• Запуск на более крупной топологии чипов: Если веса модели слишком велики для существующей топологии, можно попробовать распределить их по большему количеству чипов.
• Внедрите передовые методы сегментирования:
  • Изучите более продвинутые подходы к параллельной обработке данных, тензоров или конвейеров.
  • Укажите подсказки по сегментированию для промежуточных значений и выходных данных.
• Используйте JAX Host Offloading: переносите обработку больших тензоров в память ЦП хоста. Например, разгрузка активации и разгрузка состояния оптимизатора .

D. Настройка памяти клавиш, влияющая на флаги XLA:

Ключевые параметры памяти можно настроить таким образом, чтобы обеспечить компромисс между производительностью и меньшим использованием памяти. Однако их следует использовать только в крайнем случае, поскольку это может негативно сказаться на производительности.

E. Tune XLA Повторная материализация / Контрольная точка, созданная вручную

Если модель почти помещается в память, можно принудительно установить приоритет экономии памяти для прохода XLA::Rematerialization , что потенциально может привести к замедлению компиляции:

В качестве альтернативы, используйте декоратор jax.checkpoint с jax.grad , чтобы вручную управлять тем, какие промежуточные значения сохраняются при прямом проходе, а какие пересчитываются при обратном проходе, обменивая вычислительные циклы на HBM.

F. Используйте передовые инструменты профилирования.

В руководстве по отладке ошибок нехватки памяти (OOM) с помощью XProf представлен обзор использования памяти XProf для визуализации данных об использовании HBM, отображаемых компилятором.

Этот инструмент позволяет отслеживать пиковые значения выделения памяти и время жизни буферов, что крайне важно для точного понимания того, что именно потребляет HBM в момент пиковой нагрузки. Общие сведения о настройке профилирования см. в разделе «Начало работы с Xprof и профилированием TensorBoard» .