Kod błędu: E1000

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: TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G

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.

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 szkoleniowe, stany optymalizatora itp.
  • Pamięć tymczasowa TPU: pamięć dynamiczna wymagana do obliczeń pośrednich (np.aktywacji, gradientów itp.).
  • Skompilowany plik binarny: kod maszynowy zarówno dla rdzeni Tensor Core (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 na podstawie HBM.
  • Wewnętrzne działanie kompilatora: alokacje na poziomie programu i na poziomie HLO (np. informacje o routingu dla 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:

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

Jeśli błąd zawiera szczegółowe informacje o użyciu, np. „TC Hbm usage: X, SC Hbm usage Y”, oznacza to, że łączne użycie rdzeni TensorCore (TC) i SparseCore (SC) przekracza limit HBM. Porównaj te 2 wartości, aby zidentyfikować wąskie gardło:

  • Wysokie 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ć maksymalne 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 związany z wypełnieniem: SparseCore wyrównuje tabele zagnieżdżania do 32 B (8 liczb zmiennoprzecinkowych). Tabele o małej szerokości funkcji (np. < 8 liczb zmiennoprzecinkowych) powodują znaczny narzut dopełnienia, co marnuje pamięć 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ł: 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
  • 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 zaleceniami we wszystkich 3 sekcjach.

Sytuacja 2. Brak pamięci z powodu nieoczekiwanie dużych przydziałów

Jeśli widzisz komunikat o błędzie „Ran out of memory in memory space HBM” i w dziennikach występuje co najmniej jedna nieoczekiwanie duża alokacja (ponad 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 szczytowego wykorzystania pamięci. Usuń pozostałe znaki jax.debug.print().
  • Naprawianie nieefektywnych kształtów siatki lub fragmentacji
    • Nieprawidłowe kształty siatki lub brak adnotacji dotyczących fragmentacji mogą spowodować, że kompilator domyślnie zastosuje replikację danych, co zmusi go do próby umieszczenia bardzo dużych Tensorów na jednym układzie scalonym.
    • Sprawdź kształty dużych alokacji i upewnij się, że XLA prawidłowo określa i propaguje dzielenie na fragmenty.

Sytuacja 3. Brak pamięci z powodu łącznej alokacji

Jeśli zobaczysz komunikat o błędzie „Ran out of memory in memory space HBM” i w dziennikach nie ma nieoczekiwanie dużych tensorów, program wyczerpuje zasoby z powodu łącznej sumy przydziałów przekraczającej limit HBM. W takim przypadku często przydatne jest wizualizowanie profilu 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.

Scenariusz 3.A. Dostosowywanie konfiguracji

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

  • Zmniejsz wielkość wsadu: pamięć potrzebna na aktywacje pośrednie i gradienty jest wprost proporcjonalna do wielkości wsadu. Zmniejszenie rozmiaru partii może często pomóc w ograniczeniu wykorzystania 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 mieszanej precyzji (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ą. Pamiętaj, że ta zmiana może wpłynąć na działanie modelu i należy ją dokładnie rozważyć.

Scenariusz 3.B. 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 układ. Jeśli to możliwe, przejdź na nowsze generacje TPU.
  • Uruchomienie na większej topologii układu: 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żyj przenoszenia obliczeń na hosta w JAX: techniki przenoszenia obliczeń na hosta pozwalają użytkownikowi przenosić duże tensory do pamięci procesora hosta (np. przenoszenie aktywacji i przenoszenie stanu optymalizatora).

Scenariusz 3.C. 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. Wypełnienie wpływa zarówno na tablice danych wejściowych, jak i na tensory pośrednie (tymczasowe HLO), co może znacznie rozszerzać wykorzystanie 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 przeglądarce pamięci Xprof wyświetli się karta szczegółów 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 (wielkość wsadu, wymiar osadzania, rozmiar ukryty) są wielokrotnościami liczby 128. Pamiętaj, że ta zmiana może wpłynąć na działanie modelu i należy ją dokładnie rozważyć.

Scenariusz 3.D. Dostrajanie kluczowych flag pamięci wpływających na XLA

Kluczowe flagi pamięci można dostosować, aby zwiększyć wydajność kosztem mniejszego wykorzystania pamięci. Ta strategia powinna być jednak stosowana w ostateczności, ponieważ może negatywnie wpłynąć na skuteczność.

Scenariusz 3.E. Tune XLA rematerialization pass/manual checkpointing

Jeśli model jest bliski zmieszczenia się w pamięci, możesz użyć dekoratora jax.checkpoint z parametrem jax.grad, aby ręcznie kontrolować, które wartości pośrednie są zapisywane w przypadku przejścia do przodu, a które są ponownie obliczane w przypadku przejścia do tyłu, co pozwala zaoszczędzić cykle obliczeniowe kosztem pamięci HBM.

Możesz też wymusić przekazanie XLA::Rematerialization, aby nadać priorytet oszczędzaniu pamięci, potencjalnie kosztem wolniejszej kompilacji:

Flaga Opis Wpływ / kompromis
--xla_tpu_max_hbm_size_mib Ręczne ustawianie limitu rozmiaru HBM używanego przez etap ponownego materiału. Wymusza na kompilatorze większy wysiłek, aby dopasować program do limitu mniejszego niż rzeczywista fizyczna pamięć HBM.
--xla_tpu_rematerialization_algo=PEAK_PRIORITY Koncentruje działania w miejscach szczytowego wykorzystania pamięci. Może być bardziej wydajny w przypadku agresywnego zmniejszania ilości pamięci niż algorytm domyślny.
--xla_tpu_rematerialization_max_block_size_limit=32 Określa maksymalną liczbę instrukcji w bloku, które można jednocześnie ponownie zmaterializować. Zwiększenie tej wartości pozwala zaoszczędzić pamięć, ale znacznie wydłuża czas kompilacji.
--xla_tpu_rematerialization_block_effort_factor=10.0 Określa nakład pracy (czas kompilacji) poświęcony na wyszukiwanie bloków do ponownego utworzenia. Wyższe wartości umożliwiają dokładniejsze wyszukiwanie oszczędności pamięci kosztem dłuższego czasu kompilacji.
--xla_tpu_pre_fusion_remat=true Umożliwia dodatkowe przekształcenie przed przekształceniem fuzji. Może znaleźć więcej oszczędności pamięci, ale wydłuża czas kompilacji i może potencjalnie wpłynąć na stabilność numeryczną.

Pamiętaj, że wprowadzanie zmian we flagach XLA powinno być ostatecznością, ponieważ może negatywnie wpłynąć na wydajność.

Scenariusz 3.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 konfigurowaniu profilowania znajdziesz w artykułach Pierwsze kroki z Xprofem i Profilowanie w TensorBoardzie.