Kod błędu: 1001

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

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" ...
  • Custom Kernel Scoped Vmem OOM: jeśli błąd wskazuje na niestandardowe jądro → przejdź do sekcji Ponowne dostrojenie jądra.
  • Problemy z pamięcią wirtualną poza jądrem: jeśli błąd braku pamięci wirtualnej występuje z powodu operacji poza jądrem niestandardowym, prawdopodobnie jest to błąd wewnętrzny kompilatora. Zgłoś błąd w XLA, dołączając zrzut 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 o ograniczonym 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 jednak, aby nie przypisywać zbyt wielu danych wejściowych i wyjściowych do pamięci Vmem, ponieważ może to spowodować ogólny błąd braku pamięci VMEM.
  • Jeśli dostosowanie do konkretnego jądra jest trudne lub problem dotyczy wielu jąder, możesz dostosować globalny limit pamięci Vmem za pomocą flagi --xla_tpu_scoped_vmem_limit_kib.