Отладка ошибок нехватки памяти с помощью XProf

Ошибки Out Of Memory (OOM) возникают, когда исчерпана емкость высокоскоростной памяти (HBM) ускорителя (GPU или TPU). Некоторые распространенные причины проблем с OOM и методы отладки подробно описаны в документации E1000 — Compile Time HBM OOM и документации JAX по выделению памяти GPU .

На этой странице описано, как использовать инструмент Memory Viewer в XProf для визуализации использования памяти вашей JAX-программы, выявления пиковых значений использования и отладки ошибок нехватки памяти (OOM). Для этого необходимо выполнить следующие шаги:

  1. Запустите свою программу с помощью jax.profiler.trace , чтобы получить профиль.
  2. Запустите XProf в фоновом режиме и используйте инструмент «Просмотр памяти» , чтобы просмотреть подробную информацию об использовании памяти.

Пример программы

Следующая программа JAX приводит к ошибке нехватки памяти (OOM):

import jax
from jax import random
import jax.numpy as jnp


@jax.profiler.trace("/tmp/xprof")
@jax.jit
def oom():
    a = random.normal(random.PRNGKey(1), (327680, 327680), dtype=jnp.bfloat16)
    return a @ a


if __name__ == "__main__":
    oom()

На TPU-машине эта программа завершается с ошибкой:

XlaRuntimeError: RESOURCE_EXHAUSTED: Allocation (size=107374182400) would exceed memory (size=17179869184) :: #allocation7 [shape = 'u8[327680,327680]{1,0:T(8,128)(4,1)}', space=hbm, size = 0xffffffffffffffff, tag = 'output of xor_convert_fusion@{}'] :: <no-hlo-instruction>

(На машине с графическим процессором ошибка выглядит так: XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 214748364800 bytes. )

Запустите XProf

Установите xprof ( pip install xprof ) и запустите экземпляр XProf, указав каталог, где хранится профиль:

xprof --logdir=/tmp/xprof/ --port=6006

Перейдите к экземпляру (на локальном компьютере по http://localhost:6006 ). В раскрывающемся списке «Инструменты» выберите «Просмотр памяти» , а в окне инструментов «Просмотр памяти» выберите HBM в раскрывающемся списке «Типы памяти» (обычно выбрано по умолчанию).

Страница программы просмотра памяти XProf для приведенного выше примера программы.

В документации к инструменту XProf: Memory Viewer описаны его компоненты и предоставляемая информация.

Обратите внимание на раздел «Операции HLO при пиковом распределении памяти» , где показаны три графика буфера в точке пикового использования памяти. Буфер включает в себя: * Входные и выходные данные программы: обучающие пакеты, состояния оптимизатора и т. д. * Временные данные TensorCore и SparseCore: динамическая память, необходимая для промежуточных вычислений (например, активаций, градиентов и т. д.)

Наведя курсор на графики буферов, вы получите более подробную информацию об операции, такую ​​как ее размер, форма, тип выделения памяти и многое другое. Это поможет вам выявить и оценить операции, которые могут иметь большие или длительные временные буферы, любые большие входные/промежуточные/выходные тензоры с неэффективным заполнением и т. д., которые способствуют пиковому потреблению памяти и нуждаются в корректировке или оптимизации.

Изучите конкретные методы отладки в курсе E1000: Отладка .