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 del 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 generazione di TPU specifica 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 di hardware. Questo errore viene attivato quando il kernel tenta di mappare una moltiplicazione di matrici all'MXU utilizzando un tipo di dati per cui la generazione di TPU specifica non dispone dei circuiti fisici 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 è stato in grado di farlo perché la modalità compatibilità è stata disattivata.

Debug

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

1. Eseguire il cast ai tipi nativi

La correzione più affidabile consiste nell'eseguire manualmente il cast degli operandi a un tipo di dati supportato dall'hardware (ad es. 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 della 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 in modalità compatibilità , che è attivata per impostazione predefinita. Controlla le configurazioni XLA per assicurarti che --xla_mosaic_compat_mode non sia impostato su false.

Questa opzione funge da "polyfill", inserendo sequenze di emulazione software per le 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 in virgola mobile inserendo automaticamente operazioni di cast (ad es. estendendo gli interi a Float32 prima del matmul).
• Emulazione a bassa precisione: su alcune generazioni di hardware, emula i tipi non supportati come la virgola mobile a 4-bit (4E2M1FN) o la virgola mobile a 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 di dati prima che l'MXU possa operare su di essi.

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

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

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