Kod błędu: 0101

Kategoria: środowisko wykonawcze: błąd przydzielania programu

Ten błąd oznacza, że środowisko wykonawcze XLA na urządzeniu TPU nie może załadować skompilowanego pliku wykonywalnego programu XLA do pamięci HBM TPU.

Przykładowy komunikat o błędzie:

XlaRuntimeError: RESOURCE_EXHAUSTED: Error loading program 'jit_embedding_pipeline_step_fn': Attempting to reserve 29.49G at the bottom of memory. That was not possible. There are 147.64M free, 0B reserved, and 147.64M reservable. Scope: unknown..: while running replica 0 and partition 34 of a replicated computation (other replicas may have failed as well).

Backendy XLA: TPU

Przegląd

Ten błąd jest zwykle spowodowany jedną z tych przyczyn:

• Program Size Exceeds Available HBM: skompilowany program XLA, w tym jego instrukcje, dane statyczne i wszystkie osadzone stałe, jest większy niż łączna ilość wolnej pamięci HBM dostępnej obecnie na konkretnych rdzeniach TPU, na których program jest ładowany.
• Fragmentacja pamięci HBM: chociaż łączna ilość wolnej pamięci HBM na urządzeniu może być wystarczająca, nie jest ona dostępna w postaci jednego ciągłego bloku o rozmiarze wystarczającym do pomieszczenia całego programu.

Warto wiedzieć, jak środowisko wykonawcze TPU traktuje priorytetowo pamięć. Alokacje bufora mają wyższy priorytet niż załadowane programy. Jeśli przydzielenie bufora się nie powiedzie, środowisko wykonawcze usunie z pamięci HBM już załadowane programy, aby zwolnić miejsce. Może to doprowadzić do sytuacji, w której program, który wcześniej wczytywał się bez problemu, teraz kończy się błędem braku pamięci, ponieważ pamięć HBM jest teraz zajęta przez większą liczbę buforów danych.

Debugowanie

• Zmniejszanie ilości pamięci bufora: zwolnienie pamięci używanej przez bufory danych pozostawi więcej miejsca na sam program:
  • Zmniejsz rozmiar partii: to jeden z najskuteczniejszych sposobów na zmniejszenie ilości pamięci używanej na potrzeby aktywacji.
  • Podział parametrów: w przypadku bardzo dużych modeli użyj równoległości modelu lub technik podziału (takich jak FSDP lub Megascale), aby rozdzielić parametry i obliczenia modelu na wiele rdzeni TPU lub hostów.
  • Skróć długość sekwencji lub kontekstu: w przypadku modeli przetwarzających dane sekwencyjne (np. modele NLP), zmniejszenie długości sekwencji może znacznie zmniejszyć zużycie pamięci.
  • Buforowanie darowizn: korzystaj z funkcji platformy (np. jax.jit(..., donate_argnums=...)), aby umożliwić XLA ponowne wykorzystanie pamięci buforów wejściowych do przechowywania danych wyjściowych, co zmniejsza maksymalne wykorzystanie pamięci.
• Zmniejsz wymagania programu dotyczące pamięci tymczasowej:
  • Zmniejsz zużycie pamięci przez programy na potrzeby zmiennych tymczasowych, używając flagi tpu_shared_memory_percent. Pamiętaj, że może to negatywnie wpłynąć na skuteczność.
• Optymalizacja strategii realizacji lub zmniejszenie obciążenia wyświetlania:
  • Zarządzanie ładowaniem programu: jeśli kompilujesz wiele funkcji w trybie JIT, pamiętaj, że każda funkcja może spowodować załadowanie programu. Spróbuj tak zorganizować pracę, aby zminimalizować liczbę jednocześnie wczytywanych programó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 będzie generować szczegółowy raport o wszystkich bieżących alokacjach w przypadku błędu braku pamięci. Może to być bardzo przydatne podczas debugowania.
• Profilowanie programu: użyj profilera pamięci JAX lub profilera TensorFlow, aby uzyskać szczegółowy wgląd w wykorzystanie pamięci przez program w czasie. Może to pomóc w identyfikowaniu nieoczekiwanych skoków zużycia pamięci.