Code d'erreur : E2001

Catégorie : Compile Time: Unsupported RHS DataType on Hardware

Cette erreur se produit lorsque le type de données utilisé pour l'opérande Right-Hand Side (RHS) (côté droit) dans une multiplication matricielle (par exemple, jax.lax.dot_general, jax.lax.conv, jax.numpy.matmul ou l'opérateur @) n'est pas compatible en mode natif avec la génération de TPU spécifique utilisée.

Exemples de messages d'erreur :

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

Backends XLA : TPU

Présentation

L'unité de multiplication matricielle (MXU) du TPU est compatible en mode natif avec les opérations Float32 sur toutes les générations de matériel.

Toutefois, la compatibilité native avec BFloat16 et d'autres types de données quantifiés (par exemple, Int4, Int8 ou Float8) varie en fonction de la génération de matériel. Cette erreur se produit lorsque votre noyau tente de mapper une multiplication matricielle à la MXU à l'aide d'un type de données que votre génération de TPU spécifique ne peut pas exécuter en raison de son circuit physique.

Cette erreur indique généralement que la passe de canonicalisation du compilateur, qui tente de convertir automatiquement les types non compatibles en types compatibles (par exemple, via l'émulation logicielle), n'a pas pu trouver de règle de conversion valide ou a été empêchée de le faire, car le mode de compatibilité était désactivé.

Débogage

Pour résoudre cette erreur, vous devez aligner vos types de données sur les capacités de votre matériel. Vous disposez des options suivantes :

1. Caster vers des types natifs

La solution la plus fiable consiste à caster manuellement vos opérandes vers un type de données compatible avec le matériel (comme Float32 ou BFloat16 sur TPU v4+) dans votre noyau avant l'opération matmul.

• Pourquoi : Float32 est le type de données universel compatible en mode natif avec la MXU sur toutes les générations de TPU.
• Compromis : cela entraîne un coût de VPU (unité de traitement vectoriel) : les cycles nécessaires pour effectuer le cast, mais cela garantit que votre noyau s'exécutera sur le matériel actuel.

2. Vérifier le mode de compatibilité

En règle générale, le compilateur peut gérer automatiquement ces problèmes d'incompatibilité de type en mode de compatibilité , qui est activé par défaut. Vérifiez les configurations XLA pour vous assurer que --xla_mosaic_compat_mode n'est pas défini sur "false".

Cela agit comme un "polyfill", en injectant des séquences d'émulation logicielle pour les opérations que votre matériel ne prend pas en charge en mode natif.

Ce que le mode de compatibilité permet :

• MatMuls en précision mixte : permet de mélanger des opérandes entiers avec des accumulateurs flottants en insérant automatiquement des opérations de cast (par exemple, en étendant les entiers à Float32 avant le matmul).
• Émulation de faible précision : sur certaines générations de matériel, émule les types non compatibles tels que le point flottant 4-bit (4E2M1FN) ou le point flottant 8-bit (8E4M3FN) en les étendant à des types compatibles tels que BFloat16 ou Float32 avant l'exécution.

Notez que ce mode privilégie la compatibilité par rapport aux performances maximales, car l'émulation nécessite des instructions supplémentaires pour convertir les formats de données avant que la MXU puisse les utiliser.

3. Mettre à niveau le matériel ou demander de l'aide

Si votre algorithme nécessite strictement des performances natives pour des types tels que Int4 ou Float8 sans la surcharge du cast ou de l'émulation, vous devrez l'exécuter sur une génération de TPU plus récente avec une compatibilité native.

Demande de fonctionnalité : si vous pensez que votre matériel est compatible avec cette opération ou si le compilateur ne trouve pas de chemin d'émulation valide, même en mode de compatibilité, veuillez envoyer une demande de fonctionnalité. Nous garantissons généralement que les opérations sont compatibles avec les versions ultérieures. Par conséquent, si votre noyau s'exécute sur une génération de TPU, il devrait s'exécuter sur toutes les générations futures, mais l'émulation pour les générations plus anciennes n'est pas garantie (pour certaines d'entre elles, les casts seraient très coûteux).