Kod błędu: 0100

Kategoria: Runtime: Buffer Allocation Failure

Ten błąd oznacza, że moduł przydzielania pamięci środowiska wykonawczego XLA:TPU nie mógł znaleźć odpowiedniego bloku pamięci 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

  • niepowodzenia przydzielania bufora zainicjowanego przez użytkownika za pomocą funkcji jax.device_put lub
  • niepowodzenia związane z przydziałami danych wyjściowych programów zaplanowanych przez użytkownika.

Przyczyną tych niepowodzeń jest zwykle kilka czynników:

  • Brak pamięci: użytkownik próbuje przydzielić blok pamięci większy niż łączna ilość wolnej pamięci dostępnej w pamięci HBM TPU.
  • Fragmentacja pamięci: alokacja nie powiodła się, ponieważ w przestrzeni pamięci nie ma pojedynczego, ciągłego bloku wolnego miejsca o rozmiarze wystarczającym do spełnienia żądania. Łączna ilość wolnej pamięci jest wystarczająca do przydzielenia, ale jest rozproszona w przestrzeni pamięci w małych, nieciągłych blokach.

Środowisko wykonawcze TPU ma kilka mechanizmów ponawiania prób przydzielania, w tym:

  • Jeśli istnieją kolejki zwalniania pamięci, środowisko wykonawcze ponawia nieudane przydziały pamięci.
  • W przypadku błędów braku pamięci spowodowanych fragmentacją środowisko wykonawcze może automatycznie wywołać defragmentację i ponowić próbę.
  • Środowisko wykonawcze TPU nadaje priorytet przydzielaniu buforów nad utrzymywaniem załadowanych programów. Jeśli przydzielenie bufora nie powiedzie się z powodu niewystarczającej ilości pamięci HBM, system usunie załadowane programy TPU, dopóki nie będzie wystarczająco dużo pamięci na bufor.

Błędy, które wystąpią po zastosowaniu powyższych środków, zwykle wymagają działania użytkownika.

Debugowanie

  • Zmniejsz wykorzystanie pamięci przez model:
    • Zmniejsz rozmiar partii: zmniejszenie rozmiaru partii bezpośrednio obniża zużycie 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 pamięć HBM wielu rdzeni TPU lub hostów.
    • Skróć długość sekwencji lub kontekstu: w przypadku modeli działających na sekwencjach (np. modeli językowych) zmniejszenie długości sekwencji wejściowej może znacznie zmniejszyć zajętość pamięci.
    • Bufor darowizny: wykorzystaj funkcje platformy (np. jax.jit(..., donate_argnums=...)), aby zasygnalizować XLA, że określone bufory wejściowe można zastąpić i ponownie wykorzystać na potrzeby danych wyjściowych.
    • Strategia optymalizacji punktu kontrolnego: zamiast zapisywać cały stan modelu naraz, rozważ zapisywanie tylko wag modelu lub używanie strategii dzielonego punktu kontrolnego.
  • Układ pamięci adresów i wypełnienie:
    • Pamięć TPU jest przydzielana w blokach, a dopełnienie może zwiększyć rzeczywisty rozmiar tensorów.
  • Sprawdź, czy nie występują wycieki pamięci:
    • Sprawdź, czy odwołania do obiektów jax.Array nie są przechowywane dłużej niż zamierzono. Przechowywanie obiektów jax.Array może uniemożliwić automatyczne zwolnienie pamięci nawet po zakończeniu kompilacji programu.

Narzędzia

Włącz flagę tpu_log_allocations_on_oom, dla której alokator wygeneruje szczegółowy raport o wszystkich bieżących alokacjach, gdy wystąpi błąd braku pamięci. Może to być bardzo przydatne podczas debugowania.