Fehlercode: E2001

Kategorie:Kompilierzeit: Nicht unterstützter RHS-Datentyp auf der Hardware

Dieser Fehler tritt auf, wenn der Datentyp, der für den Operanden auf der rechten Seite (RHS) in einer Matrixmultiplikation verwendet wird (z.B. jax.lax.dot_general, jax.lax.conv, jax.numpy.matmul oder der Operator @), von der verwendeten TPU-Generation nicht nativ unterstützt wird.

Beispiel für Fehlermeldungen :

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<...>

XLA-Back-Ends:TPU

Übersicht

Die Matrixmultiplikationseinheit (Matrix Multiply Unit, MXU) der TPU unterstützt nativ Float32-Operationen auf allen Hardwaregenerationen.

Die native Unterstützung für BFloat16 und andere quantisierte Datentypen (z.B. Int4, Int8 oder Float8) variiert jedoch je nach Hardwaregeneration. Dieser Fehler wird ausgelöst, wenn Ihr Kernel versucht, eine Matrixmultiplikation der MXU mit einem Datentyp zuzuordnen, für den Ihre spezifische TPU-Generation keine physischen Schaltkreise zur Ausführung hat.

Dieser Fehler weist in der Regel darauf hin, dass der Canonicalization- Pass des Compilers, der versucht, nicht unterstützte Typen automatisch in unterstützte Typen zu konvertieren (z.B. über die Softwareemulation), keine gültige Konvertierungsregel finden konnte oder daran gehindert wurde, weil der Kompatibilitätsmodus deaktiviert war.

Debugging

Um diesen Fehler zu beheben, müssen Sie Ihre Datentypen an die Funktionen Ihrer Hardware anpassen. Sie haben folgende Möglichkeiten:

1. In native Typen umwandeln

Die zuverlässigste Lösung besteht darin, Ihre Operanden manuell in einen von der Hardware unterstützten Datentyp (z. B. Float32 oder BFloat16 auf der TPU v4+) in Ihrem Kernel vor der Matmul-Operation umzuwandeln.

  • Grund:Float32 ist der universelle Datentyp, der von der MXU auf allen TPU-Generationen nativ unterstützt wird.
  • Kompromiss:Dies ist mit Kosten für die VPU (Vector Processing Unit) verbunden, nämlich den Zyklen, die für die Umwandlung erforderlich sind. Es garantiert jedoch, dass Ihr Kernel auf der aktuellen Hardware ausgeführt wird.

2. Kompatibilitätsmodus prüfen

In der Regel kann der Compiler diese Probleme mit Typkonflikten im Kompatibilitätsmodus automatisch beheben, der standardmäßig aktiviert ist. Prüfen Sie die XLA-Konfigurationen noch einmal, um sicherzugehen, dass --xla_mosaic_compat_mode nicht auf „false“ gesetzt ist.

Dies fungiert als „Polyfill“ und fügt Softwareemulationssequenzen für Operationen ein, die von Ihrer Hardware nicht nativ unterstützt werden.

Was der Kompatibilitätsmodus ermöglicht :

  • MatMuls mit gemischter Genauigkeit:Ermöglicht das Mischen von Integer-Operanden mit Float-Akkumulatoren, indem automatisch Umwandlungsoperationen eingefügt werden (z.B. das Erweitern von Ganzzahlen auf Float32 vor der Matmul-Operation).
  • Emulation mit niedriger Genauigkeit:Auf bestimmten Hardwaregenerationen werden nicht unterstützte Typen wie 4-bit-Gleitkommazahlen (4E2M1FN) oder 8-bit-Gleitkommazahlen (8E4M3FN) emuliert, indem sie vor der Ausführung auf unterstützte Typen wie BFloat16 oder Float32 erweitert werden.

Beachten Sie, dass in diesem Modus die Kompatibilität Vorrang vor der Spitzenleistung hat, da für die Emulation zusätzliche Anweisungen erforderlich sind, um Datenformate zu konvertieren, bevor die MXU sie verarbeiten kann.

3. Hardware aktualisieren oder Support anfordern

Wenn Ihr Algorithmus unbedingt die native Leistung für Typen wie Int4 oder Float8 ohne den Overhead von Umwandlung oder Emulation erfordert, müssen Sie ihn auf einer neueren TPU-Generation mit nativer Unterstützung ausführen.

Funktionsanfrage:Wenn Sie der Meinung sind, dass Ihre Hardware diese Operation unterstützt, oder wenn dem Compiler auch im Kompatibilitätsmodus ein gültiger Emulationspfad fehlt, reichen Sie bitte eine Funktionsanfrage ein. In der Regel garantieren wir, dass Operationen vorwärtskompatibel sind. Wenn Ihr Kernel also auf einer TPU- Generation ausgeführt wird, sollte er auch auf allen zukünftigen Generationen ausgeführt werden. Es wird jedoch nicht garantiert, dass die Emulation für ältere Generationen funktioniert (für einige davon wären die Umwandlungen sehr teuer).