Kod błędu: E2001

Kategoria: Compile Time: Unsupported RHS DataType on Hardware

Ten błąd występuje, gdy typ danych używany w przypadku prawostronnego operandu w mnożeniu macierzy (np. jax.lax.dot_general, jax.lax.conv, jax.numpy.matmul lub operator @) nie jest natywnie obsługiwany przez używaną generację TPU.

Przykładowe komunikaty o błędach:

INTERNAL: Mosaic failed to compile TPU kernel: Unsupported matmul RHS type on target: 'vector<256x256xi8>'
...

The MLIR operation involved:
%13440 = "tpu.matmul"(%13435, %13437, %13439) <dimension_numbers = #tpu.dot_dimension_numbers<...>

Backendy XLA: TPU

Przegląd

Jednostka mnożenia macierzy (MXU) TPU natywnie obsługuje operacje Float32 na wszystkich generacjach sprzętu.

Jednak natywna obsługa typu BFloat16 i innych typów danych skwantyzowanych (np. Int4, Int8 lub Float8) różni się w zależności od generacji sprzętu. Ten błąd jest wywoływany, gdy jądro próbuje zmapować mnożenie macierzy na MXU przy użyciu typu danych, którego dana generacja TPU nie może wykonać ze względu na brak odpowiednich obwodów fizycznych.

Ten błąd zwykle oznacza, że etap kanonizacji kompilatora, który próbuje automatycznie przekonwertować nieobsługiwane typy na obsługiwane (np. za pomocą emulacji oprogramowania), nie mógł znaleźć prawidłowej reguły konwersji lub nie mógł tego zrobić, ponieważ tryb zgodności był wyłączony.

Debugowanie

Aby rozwiązać ten błąd, musisz dopasować typy danych do możliwości sprzętu. Masz do wyboru te opcje:

1. Rzutowanie na typy natywne

Najbardziej niezawodne rozwiązanie to ręczne przekształcenie operandów na typ danych obsługiwany przez sprzęt (np. Float32 lub BFloat16 na TPU v4+) w jądrze przed operacją matmul.

  • Uzasadnienie: Float32 to uniwersalny typ danych obsługiwany natywnie przez MXU we wszystkich generacjach TPU.
  • Kompromis: wiąże się to z kosztem jednostki VPU (Vector Processing Unit) – cyklami wymaganymi do wykonania rzutowania, ale gwarantuje, że jądro będzie działać na bieżącym sprzęcie.

2. Sprawdzanie trybu zgodności

Zwykle kompilator może automatycznie rozwiązywać problemy z niezgodnością typów w trybie zgodności, który jest domyślnie włączony. Sprawdź konfiguracje XLA, aby upewnić się, że wartość --xla_mosaic_compat_mode nie jest ustawiona na false.

Działa to jak „polyfill”, czyli wstrzykuje sekwencje emulacji oprogramowania dla operacji, których sprzęt nie obsługuje natywnie.

Co umożliwia tryb zgodności:

  • Mnożenie macierzy z użyciem mieszanej precyzji: umożliwia mieszanie operandów całkowitych z akumulatorami zmiennoprzecinkowymi przez automatyczne wstawianie operacji rzutowania (np. rozszerzanie liczb całkowitych do Float32 przed mnożeniem macierzy).
  • Emulacja o niskiej precyzji: w przypadku niektórych generacji sprzętu emuluje nieobsługiwane typy, takie jak 4-bit zmiennoprzecinkowy (4E2M1FN) lub 8-bit zmiennoprzecinkowy (8E4M3FN), rozszerzając je przed wykonaniem do obsługiwanych typów, takich jak BFloat16 lub Float32.

Pamiętaj, że w tym trybie priorytetem jest zgodność, a nie maksymalna wydajność, ponieważ emulacja wymaga dodatkowych instrukcji do przekształcania formatów danych, zanim MXU będzie mogła na nich działać.

3. Ulepszanie sprzętu lub proszenie o pomoc

Jeśli Twój algorytm wymaga natywnej wydajności w przypadku typów takich jak Int4 lub Float8 bez narzutu związanego z rzutowaniem lub emulacją, musisz uruchomić go na nowszej generacji TPU z natywną obsługą.

Prośba o dodanie funkcji: jeśli uważasz, że Twój sprzęt obsługuje tę operację lub jeśli kompilatorowi brakuje prawidłowej ścieżki emulacji nawet w trybie zgodności, prześlij prośbę o dodanie funkcji. Zazwyczaj gwarantujemy, że operacje są zgodne z wersjami przyszłymi. Jeśli więc Twój kernel działa na TPU określonej generacji, powinien działać na wszystkich przyszłych generacjach, ale nie ma gwarancji, że będzie emulowany na starszych generacjach (w przypadku niektórych z nich rzutowania byłyby bardzo kosztowne).