Accélérer le deep learning sur anciens GPU avec un logiciel FP8
Face à la croissance exponentielle des modèles d’intelligence artificielle et des volumes de données, la bande passante mémoire des GPU est devenue un goulot d’étranglement majeur. Bien que les nouvelles architectures matérielles comme les GPU Ada et Blackwell intègrent nativement le support du format FP8 — permettant des gains significatifs en vitesse d’entraînement et d’inférence — la majorité des utilisateurs, notamment dans les milieux académiques ou les petites équipes, continuent d’utiliser des GPU plus anciens, tels que l’NVIDIA RTX 3050 6Go, qui ne supportent pas cette précision matérielle. Cet écart technologique a motivé la création de Feather, une bibliothèque open source conçue pour offrir des performances proches du FP8 sur du matériel existant, via une approche logicielle innovante. L’idée centrale repose sur le regroupement (packing) de plusieurs valeurs de faible précision (FP8 ou FP16) dans un seul conteneur FP32, en exploitant des opérations bit à bit. Grâce à cette technique, on peut stocker jusqu’à quatre valeurs FP8 ou deux FP16 dans une seule variable FP32, réduisant ainsi de 75 % à 50 % la charge mémoire. Bien que cette méthode implique des opérations de pack/unpack supplémentaires, elle s’avère très efficace dans les cas où le calcul est limité par la mémoire — ce qui est fréquent dans les opérations de deep learning, comme les produits matrice-vecteur (GEMV) ou les attention mécanismes. Le modèle hiérarchique de mémoire GPU illustre ce problème : les registres et la SRAM sont extrêmement rapides, mais très limités en capacité (environ 20 Mo), tandis que la mémoire HBM (VRAM) est bien plus grande mais 7 fois plus lente. Les unités de calcul sont souvent en attente, bloquées par le transfert de données. En réduisant la taille des données à transférer via le packing, Feather diminue la charge sur la bande passante mémoire, ce qui améliore les performances globales. Le défi technique réside dans l’implémentation de ces opérations sur des formats non natifs. Contrairement au FP16, le FP8 n’est pas directement supporté par NumPy, et les formats E5M2 et E4M3 ont des structures d’exposant différentes. Pour le E5M2, le regroupement est simple car il partage le même nombre de bits d’exposant que le FP16. Pour E4M3, une solution plus complexe est nécessaire, et la bibliothèque ml_dtypes est utilisée pour gérer les conversions avec précision. Les calculs sont ensuite réalisés via des noyaux Triton, un langage spécialisé pour les GPU qui permet d’écrire des kernels en Python, sans avoir à manipuler le CUDA directement. Le noyau décompresse les données FP32 en FP8, effectue les opérations en FP32 (pour éviter les débordements), puis stocke le résultat. Ce processus, bien que nécessitant un peu de surcharge, permet d’obtenir des gains substantiels. Des benchmarks montrent des améliorations spectaculaires : sur une opération GEMV (16384×16384), Feather atteint 3,3x de vitesse avec E5M2 (contre 2,13x avec E4M3), proche du gain théorique de 4x. Pour FlashAttention (séquence de 8192, dimension 512), les gains sont également significatifs, avec une précision numérique maintenue dans les tolérances acceptables pour le deep learning. Cependant, Feather reste à un stade expérimental. Les surcharges liées au pack/unpack, au lancement des noyaux et à la gestion des formats complexes (notamment E4M3) limitent encore ses performances optimales. Il n’est pas adapté aux cas exigeant une précision absolue ou des calculs très courts. En résumé, Feather ouvre la voie à une accélération logicielle du FP8 sur du matériel ancien, rendant les techniques avancées du deep learning accessibles à un public bien plus large. Son approche innovante, combinant packing, Triton et des bibliothèques spécialisées, montre que l’innovation logicielle peut compenser les limites matérielles. L’outil est open source, et les contributions sont fortement encouragées pour améliorer sa robustesse, sa performance et sa couverture de formats.
