Debugowanie błędów braku pamięci za pomocą XProf

Błędy braku pamięci (OOM) występują, gdy wyczerpie się pojemność pamięci o wysokiej przepustowości (HBM) akceleratora (GPU lub TPU). Niektóre typowe przyczyny problemów z brakiem pamięci i techniki debugowania zostały szczegółowo opisane w dokumentacji E1000 – Compile Time HBM OOM i dokumentacji JAX na temat przydzielania pamięci GPU.

Na tej stronie dowiesz się, jak używać narzędzia XProf Memory Viewer do wizualizacji wykorzystania pamięci przez program JAX, identyfikowania instancji szczytowego wykorzystania i debugowania błędów braku pamięci. Wymaga to wykonania tych czynności:

  1. Uruchom program z parametrem jax.profiler.trace aby zarejestrować profil.
  2. Uruchom XProf w tle i użyj narzędzia do wyświetlania pamięci, aby wyświetlić szczegóły wykorzystania pamięci.

Przykładowy program

Ten program JAX powoduje błąd braku pamięci:

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()

Na maszynie TPU ten program kończy się niepowodzeniem z tym komunikatem:

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>

(Na maszynie z GPU błąd wygląda tak: XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 214748364800 bytes.)

Uruchamianie XProf

Zainstaluj xprof (pip install xprof) i uruchom instancję XProf, podając katalog, w którym jest przechowywany profil:

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

Otwórz instancję (na komputerze lokalnym, pod adresem http://localhost:6006). W menu Narzędzia wybierz Przeglądarka pamięci, a w oknie narzędzia Przeglądarka pamięci w menu Typy pamięci wybierz HBM (zwykle jest to opcja domyślna).

Strona narzędzia XProf Memory Viewer dla powyższego przykładowego programu

W dokumentacji narzędzia XProf: Memory Viewer znajdziesz opis komponentów narzędzia i wyświetlanych informacji.

Skup się na sekcji Operacje HLO przy szczytowym przydziale pamięci, w której znajdują się 3 wykresy buforów w punkcie szczytowego wykorzystania pamięci. Bufor obejmuje: * dane wejściowe i wyjściowe programu: partie szkoleniowe, stany optymalizatora itp.; * dane tymczasowe TensorCore i SparseCore: pamięć dynamiczna wymagana do obliczeń pośrednich (takich jak aktywacje, gradienty itp.).

Aby uzyskać więcej informacji o operacji, np. jej rozmiarze, kształcie czy typie przydziału, najedź kursorem na wykresy buforowania. Pomoże Ci to zidentyfikować i ocenić operacje, które mogą mieć duże lub długotrwałe zmienne tymczasowe, duże tensory wejściowe/pośrednie/wyjściowe z nieefektywnym dopełnieniem itp., które przyczyniają się do osiągania szczytowego zużycia pamięci i wymagają dostosowania lub optymalizacji.

Poznaj konkretne techniki debugowania w artykule E1000: Debugowanie.