Codice di errore: E2001

Categoria: Compile Time: Unsupported RHS DataType on Hardware

Questo errore si verifica quando il tipo di dati utilizzato per l'operando lato destro (RHS) in una moltiplicazione di matrici (ad es. jax.lax.dot_general, jax.lax.conv, jax.numpy.matmul o l'operatore @) non è supportato in modo nativo dalla specifica generazione di TPU in uso.

Messaggi di errore di esempio:

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

Backend XLA: TPU

Panoramica

L'unità di moltiplicazione a matrice (MXU) della TPU supporta in modo nativo le operazioni Float32 su tutte le generazioni di hardware.

Tuttavia, il supporto nativo per BFloat16 e altri tipi di dati quantizzati (ad es. Int4, Int8 o Float8) varia in base alla generazione dell'hardware. Questo errore viene attivato quando il kernel tenta di mappare una moltiplicazione di matrici all'MXU utilizzando un tipo di dati che la tua generazione di TPU specifica non ha il circuito fisico per l'esecuzione.

Questo errore in genere indica che il passaggio di canonicalizzazione del compilatore, che tenta di convertire automaticamente i tipi non supportati in tipi supportati (ad es. tramite emulazione software), non è riuscito a trovare una regola di conversione valida o non è riuscito a farlo perché la modalità di compatibilità era disattivata.

Debug

Per risolvere questo errore, devi allineare i tipi di dati alle funzionalità dell'hardware. Hai le seguenti opzioni:

1. Trasmetti ai tipi nativi

La correzione più affidabile consiste nel convertire manualmente gli operandi in un tipo di dati supportato dall'hardware (come Float32 o BFloat16 su TPU v4+) all'interno del kernel prima dell'operazione matmul.

• Motivo:Float32 è il tipo di dati universale supportato in modo nativo dall'MXU su tutte le generazioni di TPU.
• Compromesso:questo comporta un costo per la VPU (Vector Processing Unit), ovvero i cicli necessari per eseguire il cast, ma garantisce che il kernel verrà eseguito sull'hardware attuale.

2. Controllare la modalità compatibilità

In genere, il compilatore può gestire automaticamente questi problemi di mancata corrispondenza dei tipi nella modalità di compatibilità, che è abilitata per impostazione predefinita. Controlla di nuovo le configurazioni XLA per assicurarti che --xla_mosaic_compat_mode non sia impostato su false.

Funziona come un "polyfill", inserendo sequenze di emulazione software per operazioni che l'hardware non supporta in modo nativo.

Cosa consente la modalità compatibilità:

• MatMul a precisione mista:consente di combinare operandi interi con accumulatori float inserendo automaticamente operazioni di cast (ad es. estendendo gli interi a Float32 prima di MatMul).
• Emulazione a bassa precisione:su alcune generazioni di hardware, emula tipi non supportati come la rappresentazione in virgola mobile 4-bit (4E2M1FN) o la rappresentazione in virgola mobile 8-bit (8E4M3FN) estendendoli a tipi supportati come BFloat16 o Float32 prima dell'esecuzione.

Tieni presente che questa modalità dà la priorità alla compatibilità rispetto alle prestazioni di picco, poiché l'emulazione richiede istruzioni aggiuntive per convertire i formati dei dati prima che l'MXU possa operare su di essi.

3. Eseguire l'upgrade dell'hardware o richiedere assistenza

Se il tuo algoritmo richiede rigorosamente prestazioni native per tipi come Int4 o Float8 senza l'overhead di casting o emulazione, dovrai eseguirlo su una generazione di TPU più recente con supporto nativo.

Richiesta di funzionalità: se ritieni che il tuo hardware supporti questa operazione o se il compilatore non trova un percorso di emulazione valido anche in modalità di compatibilità, invia una richiesta di funzionalità. In genere garantiamo la compatibilità delle operazioni. Pertanto, se il kernel viene eseguito su una generazione di TPU, dovrebbe essere eseguito su tutte le generazioni future. Tuttavia, non è garantito che sia disponibile l'emulazione per le generazioni precedenti (per alcune delle quali i cast sarebbero molto costosi).