Code d'erreur : E2001

Catégorie : Temps de compilation : type de données RHS non compatible sur le matériel

Cette erreur se produit lorsque le type de données utilisé pour l'opérande Right-Hand Side (RHS) (partie droite) 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 prise en charge native de BFloat16 et d'autres types de données quantifiées (par exemple, Int4, Int8 ou Float8) varie selon la génération du matériel. Cette erreur se produit lorsque votre noyau tente de mapper une multiplication matricielle à l'unité MXU à l'aide d'un type de données que votre génération de TPU spécifique ne possède pas le circuit physique pour exécuter.

Cette erreur indique généralement que la passe de canonisation du compilateur, qui tente de convertir automatiquement les types non compatibles en types compatibles (par exemple, par émulation logicielle), n'a pas pu trouver de règle de conversion valide ou n'a pas pu le faire, car le mode 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 avez le choix entre les options suivantes :

1. Caster vers des types natifs

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

• Pourquoi : Float32 est le type de données universel compatible en mode natif avec l'unité MXU sur toutes les générations de TPU.
• Compromis : cela entraîne un coût de VPU (Vector Processing Unit), c'est-à-dire 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 dans le mode 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".

Il sert de "polyfill" en injectant des séquences d'émulation logicielle pour les opérations que votre matériel ne prend pas en charge de manière native.

Ce que permet le mode de compatibilité :

• MatMuls de précision mixte : permet de mélanger des opérandes entiers avec des accumulateurs flottants en insérant automatiquement des opérations de conversion (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 les types à virgule flottante 4-bit (4E2M1FN) ou 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 l'unité MXU puisse les traiter.

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

Si votre algorithme nécessite des performances natives pour des types tels que Int4 ou Float8 sans la surcharge de la conversion ou de l'émulation, vous devrez l'exécuter sur une génération de TPU plus récente avec prise en charge 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 Compatibilité, veuillez envoyer une demande de fonctionnalité. Nous garantissons généralement que les opérations sont compatibles avec les versions ultérieures. Ainsi, 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. Toutefois, 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).