Kategoria: Compile Time: Scoped Vmem OOM
Ten błąd oznacza, że program wymaga więcej pamięci wektorowej o ograniczonym zakresie (Vmem), niż zostało przydzielone.
Przykładowe komunikaty o błędach:
RESOURCE_EXHAUSTED: Ran out of memory in memory space vmem while allocating on stack for %my-custom-kernel = bf16[2048,4096]{1,0:T(8,128)(2,1)} custom-call(...) ...
Backendy XLA: TPU
Przegląd
TPU mają pamięć wektorową (VMEM), która jest lokalną pamięcią podręczną używaną wyłącznie przez TensorCore (TC). Kompilator zarządza pamięcią wirtualną w przypadku różnych typów alokacji:
- Alokacje w zakresie instrukcji: tymczasowe miejsce na dane w pamięci Vmem podczas wykonywania pojedynczej instrukcji HLO. Obejmuje to bufor zakresu operandów (np. w przypadku podwójnego buforowania) i przelewanie rejestrów.
- Przydziały w zakresie programu: przydziały, które wykraczają poza zakres pojedynczej instrukcji HLO. Są to zwykle tymczasowe i pośrednie wyniki HLO, które są danymi wejściowymi lub wyjściowymi instrukcji HLO.
Błąd braku pamięci wirtualnej w zakresie czasu kompilacji występuje, gdy alokacje w zakresie instrukcji przekraczają limit alokacji dla tej instrukcji. Ten limit jest kontrolowany
- globalnie dla całego programu za pomocą flagi --xla_tpu_scoped_vmem_limit_kib.
i
- na niestandardowe jądro za pomocą parametru vmem_limit_bytes.
Te błędy są zwykle spowodowane wewnętrznym błędem kompilatora lub przekroczeniem limitu przydziału przez niestandardowe jądro.
Debugowanie
Dokładnie przeanalizuj komunikat o błędzie, aby sprawdzić, czy błąd wynika z niestandardowego jądra, czy ze standardowego HLO. Błąd spowodowany przez jądro niestandardowe powinien mieć następujący podpis:
Ran out of memory in memory space vmem while allocating on stack for %my-custom-call = <output-shape> custom-call(<params>), custom_call_target="tpu_custom_call" ...
- Niestandardowy jądro w zakresie pamięci wirtualnej OOM: jeśli błąd wskazuje na niestandardowe jądro, przejdź do sekcji Ponowne dostrajanie jądra.
- Problemy z pamięcią Vmem poza jądrem: jeśli błąd Vmem OOM występuje z powodu operacji poza jądrem niestandardowym, prawdopodobnie jest to wewnętrzny błąd kompilatora. Wyślij raport o błędzie w XLA z zrzutem HLO.
Ponowne dostrojenie jądra
Jeśli błąd pochodzi z niestandardowego jądra, użyj tych technik, aby zmniejszyć wymagania jądra dotyczące pamięci:
- Dostosuj rozmiary bloków: zmniejsz rozmiary bloków (rozmiary kafelków) w konfiguracji jądra, aby ograniczyć użycie pamięci wirtualnej w zakresie.
- Ustawianie limitów pamięci wirtualnej w zakresie poszczególnych jąder: wyraźnie zażądaj wymaganej ilości pamięci dla danego jądra za pomocą parametru vmem_limit_bytes.
- Modyfikowanie kolorowania pamięci: jawne kolorowanie/ograniczanie danych wejściowych/wyjściowych jądra do Vmem za pomocą funkcji pallas.tpu.with_memory_space_constraint. Uważaj, aby nie przypisywać zbyt wielu danych wejściowych do pamięci Vmem, ponieważ może to spowodować ogólny błąd braku pamięci Vmem.
- Dostosuj limit pamięci wirtualnej: jeśli dostosowanie do konkretnego jądra jest trudne lub problem dotyczy wielu jąder, możesz dostosować globalny limit pamięci wirtualnej za pomocą flagi --xla_tpu_scoped_vmem_limit_kib.