Kod błędu: 1000

Kategoria: czas kompilacji: błąd braku pamięci HBM

Ten błąd oznacza, że program wymaga więcej pamięci o wysokiej przepustowości (HBM), niż jest fizycznie dostępna na urządzeniu TPU.

Przykładowe komunikaty o błędach:

RESOURCE_EXHAUSTED: XLA:TPU compile permanent error. Ran out of memory in memory space hbm. Used 49.34G of 32.00G hbm. Exceeded hbm capacity by 17.34G.

RESOURCE_EXHAUSTED: TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G

Backendy XLA: TPU

Przegląd

XLA sprawdza, czy łączny rozmiar wszystkich niezbędnych statycznych alokacji mieści się w pamięci HBM urządzenia.

Kompilator zarządza stałą pojemnością pamięci HBM na TPU w przypadku kilku typów alokacji:

• Dane wejściowe i wyjściowe programu: partie treningowe, stany optymalizatora itp.
• Pamięć tymczasowa TensorCore + SparseCore: pamięć dynamiczna wymagana do obliczeń pośrednich (np.aktywacji, gradientów itp.).
• Skompilowany kod binarny: kod maszynowy zarówno dla TensorCore (TC), jak i SparseCore (SC).
• Narzucone obciążenie systemu: zarezerwowane miejsce na środowisko wykonawcze XLA (np. bufory wejściowe w starszych generacjach TPU).
• Stałe: stałe wartości osadzone w HLO IR są przydzielane do HBM.
• Wewnętrzne działanie kompilatora: alokacje na poziomie programu i na poziomie HLO (np. informacje o routingu węzłów w siatce)

Ten błąd występuje, gdy kompilator XLA nie może zmieścić wszystkich powyższych alokacji w pamięci HBM urządzenia.

Debugowanie

Dokładnie przeanalizuj komunikat o błędzie i logi, aby określić, która z poniższych kategorii błędów HBM OOM najlepiej opisuje Twój problem:

• Wykorzystanie pamięci HBM przez rdzenie TensorCore (TC) i SparseCore (SC) przekracza limit: Jeśli w komunikacie o błędzie podano szczegółowe informacje o wykorzystaniu, np.: „TC Hbm usage: X, SC Hbm usage Y”. → Przejdź do sekcji 1. Równoważenie wykorzystania pamięci HBM w przypadku rdzeni TC i SC.
• Nieoczekiwanie duże przydziały: jeśli w komunikacie o błędzie widnieje „Ran out of memory in memory space HBM”, sprawdź dzienniki pod kątem wyliczenia największych przydziałów w pamięci HBM. Jeśli występuje co najmniej 1 nieoczekiwanie duży tensor (np. > 50% limitu HBM), przejdź do sekcji 2. Nieoczekiwanie duże przydziały.
• Aggregate Allocations Exceed HBM Limit: jeśli w dziennikach nie ma żadnych nieoczekiwanie dużych tensorów, a błąd brzmi „Ran out of memory in memory space HBM”, przejdź do sekcji 3. Łączne przydziały przekraczają limit HBM.

Sekcja 1. Równoważenie wykorzystania pamięci HBM przez rdzenie Tensor i rdzenie Scalar

Jeśli błąd wyraźnie wskazuje, jak wykorzystywane są dane, np. „TC Hbm usage: X, SC Hbm usage Y” – porównaj te 2 wartości, aby zidentyfikować wąskie gardło.

• Duże wykorzystanie SparseCore:
  • Optymalizacja wykorzystania stosu HBM: zużycie pamięci stosu HBM skaluje się wraz z feature_width, max_unique_nz_per_row i logical_replica_count. Możesz zmniejszyć szczytowe wykorzystanie stosu, dostosowując flagę --xla_sc_num_serialized_tables_to_optimize_hbm, która serializuje przetwarzanie tabel. Ogranicza to jednak równoległość.
  • Sprawdź narzut na dopełnienie: SparseCore wyrównuje tabele z osadzaniem do 32 B (8 liczb zmiennoprzecinkowych). tabele z małą szerokością kolumn (np. < 8 wartości zmiennoprzecinkowych) powodują znaczny narzut związany z wypełnieniem, co prowadzi do marnowania pamięci HBM.
  • Zmniejsz użycie sterty: wysokie wartości parametru maximum_parallel_iterations zwiększają ilość danych wejściowych wstępnie pobieranych do sterty HBM. Obniżenie tej wartości może zwolnić znaczną ilość pamięci.
  • Sprawdź podział na fragmenty: upewnij się, że tabele z osadzonymi danymi są prawidłowo dzielone na fragmenty za pomocą operacji modulo na wszystkich układach. Zobacz Jak limity przekładają się na tabele.
  • Więcej pomysłów znajdziesz w artykule SC: Wąskie gardła wydajności i pamięci.
• Wysokie wykorzystanie Tensor Core:
  • Przejdź do sekcji 2.
• Zrównoważona
  • Jeśli żadna z nich nie jest nadmierna, ale ich suma jest zbyt wysoka, oznacza to, że osiągnięto maksymalną wydajność procesora. Musisz spróbować ograniczyć użycie obu komponentów. Postępuj zgodnie z rekomendacjami we wszystkich 3 sekcjach.

Sekcja 2. Nieoczekiwanie duże przydziały

Jeśli w logach występuje co najmniej 1 nieoczekiwanie duża alokacja (powyżej 50% limitu HBM), prawie nigdy nie jest to problem z pojemnością sprzętu. Zwykle jest to błąd konfiguracji. Sprawdź etykietę XLA (jeśli jest obecna) dużych alokacji, aby uzyskać wskazówki dotyczące ich kodu źródłowego JAX.

• Usuwanie artefaktów debugowania:
  • Używanie funkcji jax.debug.print() w przypadku uruchomień na dużą skalę może wymusić na kompilatorze zmaterializowanie pełnego tensora w pamięci HBM w celu przeniesienia go do procesora, co spowoduje przerwanie fuzji i zwiększenie maksymalnego wykorzystania pamięci. Usuń pozostałe znaki jax.debug.print().
• Naprawianie nieefektywnych kształtów siatki lub dzielenia na mniejsze części:
  • Nieprawidłowe kształty siatki lub brak adnotacji o podziale mogą spowodować, że kompilator domyślnie użyje replikacji, co zmusi go do próby dopasowania bardzo dużych tensorów do jednego układu.
  • Sprawdź kształty dużych przydziałów i upewnij się, że XLA prawidłowo określa i propaguje dzielenie na fragmenty.

Sekcja 3. Łączne przydziały przekraczają limit HBM

Jeśli program wyczerpie zasoby z powodu przekroczenia limitu HBM przez łączną sumę przydzielonych zasobów, często warto wizualizować profil pamięci, aby zidentyfikować konkretne bufory, które przyczyniają się do szczytowego wykorzystania. Więcej informacji o identyfikowaniu elementów, które w największym stopniu przyczyniają się do zużycia pamięci, znajdziesz w tym przewodniku.

Gdy zidentyfikujesz najważniejszych współtwórców, wykonaj te czynności, aby zoptymalizować wykorzystanie pamięci.

A. Sprawdzanie dopełnienia i wyrównania tensora

Nieefektywne kształty tensorów są częstą, cichą przyczyną błędów OOM na TPU. Aby uzyskać maksymalną wydajność na TPU, XLA uzupełnia wymiary tensorów – zwykle do wielokrotności 128 w przypadku najmniejszego wymiaru i 8 w przypadku drugiego najmniejszego. To dopełnienie wpływa zarówno na tablice wejściowe, jak i na tensory pośrednie (tymczasowe HLO), co może znacznie zwiększyć zużycie pamięci, zwłaszcza w przypadku małych rozmiarów wymiarów. Zobacz układy tablic.

• Sprawdzanie kształtów dużych buforów: (na TPU v5 z domyślnymi układami)
  • Po najechaniu kursorem na bufor w Xprof Memory Viewer wyświetli się karta ze szczegółami bufora, która zawiera informacje o buforze, w tym informacje o wypełnieniu.
  • Przykład: kształt (129, 1024) może zostać dopełniony do (256, 1024), co spowoduje zmarnowanie prawie 50% pamięci.
  • Poprawka: kształt (128, 1024) nie wymaga dopełnienia i powoduje 0% marnotrawstwa pamięci.
• Wyrównaj wymiary: upewnij się, że wszystkie duże wymiary tensora (rozmiar partii, wymiar osadzania, rozmiar ukryty) są wielokrotnościami liczby 128.

B. Dostosowywanie konfiguracji

Problemy z brakiem pamięci można często rozwiązać, wprowadzając te zmiany w konfiguracji:

• Zmniejsz rozmiar partii: pamięć potrzebna na aktywacje pośrednie i gradienty jest wprost proporcjonalna do rozmiaru partii. Zmniejszenie rozmiaru partii może często pomóc w ograniczeniu zużycia pamięci.
• Przekazywanie buforów wejściowych: jeśli używasz jax.jit, określ donate_argnums dla parametrów modelu. Dzięki temu XLA może zastąpić pamięć wejściową danymi wyjściowymi.
• Włączanie precyzji mieszanej (bfloat16): używaj bfloat16 lub kwantyzacji (int8 itp.) w przypadku największych tensorów w programie, jeśli architektura modelu i wymagania dotyczące jakości na to pozwalają.

C. Optymalizacja architektury i dzielenia na fragmenty

Jeśli zmiany konfiguracji są niewystarczające, topologia modelu może być zbyt duża w stosunku do bieżącej konfiguracji sprzętowej.

• Używaj nowszych generacji TPU: nowsze TPU mają zwykle więcej pamięci HBM na chip. Jeśli to możliwe, przejdź na nowsze generacje TPU.
• Uruchomienie na większej topologii układów: jeśli wagi modelu są zbyt duże dla istniejącej topologii, możesz spróbować podzielić je na więcej układów.
• Wdróż zaawansowane techniki dzielenia na partycje:
  • Poznaj bardziej zaawansowane podejścia do równoległości danych, tensorów lub potoków.
  • Określ wskazówki dotyczące dzielenia dla wartości pośrednich i wyników.
• Używaj funkcji Host Offloading w JAX: przenoś duże tensory do pamięci procesora hosta, np. przenoszenie aktywacji i przenoszenie stanu optymalizatora.

D. Dostosuj flagi XLA wpływające na pamięć klucza:

Kluczowe flagi pamięci można dostosować, aby zwiększyć wydajność kosztem mniejszego zużycia pamięci. Należy jednak stosować je w ostateczności, ponieważ mogą negatywnie wpłynąć na skuteczność.

E. Dostrajanie XLA Rematerialization Pass / ręczne tworzenie punktów kontrolnych

Jeśli model jest bliski zmieszczenia się w pamięci, możesz wymusić przekazanie XLA::Rematerialization, aby nadać priorytet oszczędzaniu pamięci, potencjalnie kosztem wolniejszej kompilacji:

Możesz też użyć dekoratora jax.checkpoint z parametrem jax.grad, aby ręcznie określać, które wartości pośrednie mają być zapisywane w trakcie przejścia w przód, a które mają być ponownie obliczane w trakcie przejścia w tył. W ten sposób możesz wymieniać cykle obliczeniowe na pamięć HBM.

F. Korzystanie z zaawansowanych narzędzi do profilowania

W artykule Debugowanie błędów braku pamięci za pomocą XProf znajdziesz samouczek dotyczący korzystania z przeglądarki pamięci XProf do wizualizacji sposobu wykorzystania pamięci HBM z perspektywy kompilatora.

To narzędzie umożliwia sprawdzenie maksymalnej alokacji pamięci i czasu życia bufora, co jest kluczowe, aby dokładnie zrozumieć, co zużywa pamięć HBM w momencie maksymalnego wykorzystania. Ogólne informacje o konfiguracji profilowania znajdziesz w artykułach Pierwsze kroki z Xprof i Profilowanie w TensorBoard.