Kod błędu: E0100

Kategoria: środowisko wykonawcze: nieudana alokacja bufora

Ten błąd oznacza, że alokator pamięci środowiska wykonawczego XLA:TPU nie może znaleźć odpowiedniego bloku pamięci w HBM akceleratora na potrzeby żądanej alokacji.

Przykładowy komunikat o błędzie:

ValueError: RESOURCE_EXHAUSTED: Error allocating device buffer: Attempting to allocate 8.00M. That was not possible. There are 6.43M free.; (0x0x1_HBM0)

Backendy XLA: TPU

Przegląd

Ten błąd jest zgłaszany w tych przypadkach:

  • Nieudane alokacje bufora zainicjowane przez użytkownika za pomocą funkcji jax.device_put

lub

  • Nieudane alokacje wyjściowe programu zaplanowane przez użytkownika.

Te błędy są zwykle spowodowane kilkoma przyczynami:

  • Brak pamięci (OOM): użytkownik próbuje przydzielić fragment pamięci, który jest większy niż łączna ilość wolnej pamięci dostępnej w HBM TPU.
  • Fragmentacja pamięci: alokacja nie powiodła się, ponieważ żaden pojedynczy ciągły wolny blok w przestrzeni pamięci nie jest wystarczająco duży, aby spełnić żądany rozmiar. Łączna ilość wolnej pamięci jest wystarczająca do alokacji, ale jest rozproszona w przestrzeni pamięci w małych, nieciągłych blokach.

Środowisko wykonawcze TPU ma kilka mechanizmów ponawiania nieudanych alokacji, w tym:

  • Jeśli są kolejkowane deallokacje, środowisko wykonawcze ponawia nieudane alokacje.
  • W przypadku błędów braku pamięci spowodowanych fragmentacją środowisko wykonawcze może automatycznie uruchomić defragmentację i ponowić próbę.
  • Środowisko wykonawcze TPU nadaje priorytet alokacjom buforów przed utrzymywaniem załadowanych programów. Jeśli alokacja bufora nie powiedzie się z powodu niewystarczającej ilości HBM, system usunie załadowane programy TPU, dopóki nie będzie dostępna wystarczająca ilość pamięci dla bufora.

Błąd, który wystąpi po zastosowaniu powyższych środków zaradczych, zwykle wymaga działania użytkownika.

Debugowanie

  • Zmniejsz rozmiar pamięci modelu
    • Zmniejsz wielkość wsadu: zmniejszenie wielkości wsadu bezpośrednio zmniejsza wykorzystanie pamięci.
    • Dzielenie parametrów: w przypadku bardzo dużych modeli używaj technik takich jak równoległość modelu lub dzielenie, aby rozdzielać parametry między HBM wielu rdzeni TPU lub hostów.
    • Skróć długość sekwencji/długość kontekstu: w przypadku modeli działających na sekwencjach (takich jak modele językowe) zmniejszenie długości sekwencji wejściowej może znacznie zmniejszyć wykorzystanie pamięci.
    • Przekazywanie bufora: używaj funkcji platformy (takich jak jax.jit(..., donate_argnums=...)), aby zasygnalizować XLA, że niektóre bufory wejściowe można zastąpić i ponownie wykorzystać na potrzeby danych wyjściowych. Więcej informacji znajdziesz w artykule Przekazywanie bufora.
    • Zoptymalizuj strategię punktów kontrolnych: zamiast zapisywać cały stan modelu naraz, rozważ zapisywanie tylko wag modelu lub używanie strategii dzielenia punktów kontrolnych.
  • Rozwiąż problemy z układem pamięci i dopełnieniem
    • Pamięć TPU jest przydzielana w blokach, a dopełnienie może zwiększyć rzeczywisty rozmiar tensorów.
  • Upewnij się, że nie ma wycieków pamięci
    • Upewnij się, że odwołania do obiektów jax.Array nie są przechowywane dłużej niż zamierzone. Przechowywanie obiektów jax.Array może uniemożliwić automatyczne zwolnienie pamięci nawet po zakończeniu kompilacji programu.

Inne strategie, których możesz użyć, aby zmniejszyć ilość pamięci używanej przez każdy program, znajdziesz też w artykule Kod błędu: E1000 for.

Narzędzia

Włącz flagę tpu_log_allocations_on_oom, aby alokator generował szczegółowy raport o wszystkich bieżących alokacjach w przypadku wystąpienia błędu braku pamięci. Może to być bardzo przydatne podczas debugowania.