Fehlercode: E2001

Kategorie:Kompilierungszeit: Nicht unterstützter RHS-Datentyp auf Hardware

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

Beispiele 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 Matrix Multiply Unit (MXU) der TPU unterstützt Float32-Vorgänge nativ auf allen Hardwaregenerationen.

BFloat16 und andere quantisierte Datentypen werden jedoch nativ unterstützt (z.B. Int4, Int8 oder Float8) variiert je nach Hardwaregeneration. Dieser Fehler wird ausgelöst, wenn Ihr Kernel versucht, eine Matrixmultiplikation mit einem Datentyp auf die MXU abzubilden, für den Ihre spezifische TPU-Generation nicht die physischen Schaltkreise 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 Software-Emulation), 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 die Datentypen an die Funktionen Ihrer Hardware anpassen. Sie haben folgende Möglichkeiten:

1. In native Typen umwandeln

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

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

2. Kompatibilitätsmodus prüfen

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

Dies dient als „Polyfill“ und fügt Software-Emulationssequenzen für Vorgänge ein, die Ihre Hardware nicht nativ unterstützt.

Das ermöglicht der Kompatibilitätsmodus:

  • MatMuls mit gemischter Präzision:Ermöglicht das Mischen von Ganzzahloperanden mit Gleitkomma-Akkumulatoren durch automatisches Einfügen von Cast-Operationen (z.B. Erweitern von Ganzzahlen auf Float32 vor dem Matmul).
  • Emulation mit niedriger Präzision: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.

In diesem Modus wird die Kompatibilität gegenüber der Spitzenleistung priorisiert, da für die Emulation zusätzliche Anweisungen zum Konvertieren von Datenformaten erforderlich sind, bevor die MXU damit arbeiten kann.

3. Hardware aktualisieren oder Support anfordern

Wenn Ihr Algorithmus für Typen wie Int4 oder Float8 unbedingt native Leistung ohne den Overhead von Casting 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 diesen Vorgang unterstützt, oder wenn dem Compiler auch im Kompatibilitätsmodus ein gültiger Emulationspfad fehlt, reichen Sie bitte eine Funktionsanfrage ein. Wir garantieren in der Regel, dass Vorgänge 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 gibt jedoch keine Garantie dafür, dass die Emulation für ältere Generationen verfügbar ist, da die Kosten für einige davon sehr hoch wären.