L’explosion des performances des GPU
(hazyresearch.stanford.edu)- Alors que le coût du calcul IA augmente, Hazy Research résume l’enjeu central de l’optimisation des performances GPU ainsi : maintenir les tensor cores du NVIDIA H100 constamment occupés
- Le H100 atteint 989 TFLOPs en multiplication de matrices en half-precision, mais seulement environ 60 TFLOPs pour les opérations générales ; dès que les tensor cores s’arrêtent, le taux d’utilisation chute fortement
- Pour se rapprocher des performances maximales, il faut traiter ensemble WGMMA, l’agencement de la shared memory, la génération d’adresses et l’occupancy ; sans
wgmma.mma_async, les microbenchmarks restent à environ 63 % du pic - ThunderKittens, un DSL intégré à CUDA et open source, encapsule des complexités comme le swizzling et le register layout via des abstractions de tiles et de vectors, simplifiant l’écriture de kernels de la famille FlashAttention
- Le kernel forward FlashAttention-2 pour H100 s’écrit en environ 100 lignes et est environ 30 % plus rapide que FlashAttention-2 ; le kernel Based linear attention fonctionne à 215 TFLOPs
Les conditions qui déterminent les performances du H100
- L’IA consomme beaucoup de calcul, et Hazy Research travaille depuis quelques années à faire en sorte que l’IA utilise moins de calcul, ou s’exécute plus efficacement avec un budget de calcul donné
- Exemples de réduction du calcul : Based, Monarch Mixer, H3, Hyena, S4
- Exemples d’exécution efficace : FlashAttention, FlashAttention-2, FlashFFTConv
- L’objectif pratique est de synthétiser ce qui a été appris en rendant les GPU plus rapides, et de publier ThunderKittens, un DSL intégré à CUDA pour aider à écrire des kernels rapides
- Plus largement, le sujet est de voir comment la compréhension du matériel a changé la façon d’aborder le calcul IA
Architecture du NVIDIA H100 et goulots d’étranglement
- La discussion prend pour référence un GPU H100 SXM configuré comme suit
- 80 Go de HBM3, bande passante de 3 To/s
- 50 Mo de cache L2, bande passante de 12 To/s, divisés en deux sections de 25 Mo sur l’ensemble du GPU et reliés par un crossbar
-
132 SM
- Chaque SM dispose d’un cache L1 de 256 Ko incluant jusqu’à 227 Ko de shared memory, avec au total environ 33 To/s de bande passante
- Le nouveau matériel de Hopper, le Tensor Memory Accelerator (TMA), prend en charge la génération d’adresses asynchrone et les fetchs mémoire
- Chaque SM est composé de 4 quadrants, chacun avec un warp scheduler, 512 vector registers, un tensor core pour la multiplication de matrices et des instructions intégrées parallèles
- Tout le calcul se fait dans les SM, et l’essentiel est traité dans les registers
- La clé des performances sur H100 est de maintenir les tensor cores continuellement alimentés
- Le H100 fournit 989 TFLOPs en multiplication de matrices half-precision, et environ 60 TFLOPs pour « le reste » des opérations
- Lors des cycles où les tensor cores sont utilisés, il atteint au moins 94 % d’utilisation matérielle
- Lors des cycles où les tensor cores ne sont pas utilisés, il reste à un maximum de 6 % d’utilisation
WGMMA : une instruction nécessaire mais délicate
- Le H100 dispose de
wgmma.mma_async, une instruction warp group matrix multiply accumulate- En PTX :
wgmma.mma_async - En SASS :
HGMMA/IGMMA/QGMMA/BGMMA
- En PTX :
- Sur les GPU précédents,
wmma.mma.syncetmma.syncétaient synchrones : un warp de 32 threads envoyait les données au tensor core puis attendait le résultat wgmma.mma_asyncpermet à 128 threads consécutifs de se synchroniser de façon coopérative sur tous les quadrants du SM et de lancer directement une multiplication de matrices asynchrone depuis la shared memory- Les warps peuvent effectuer d’autres tâches dans les registers pendant que la multiplication de matrices progresse
- Il est possible d’attendre le résultat au moment voulu
- Dans les microbenchmarks, ces instructions se sont révélées nécessaires pour exploiter toute la capacité de calcul du H100
- Sans elles, le GPU a été observé comme restant à environ 63 % du pic d’utilisation
- Cela peut venir du fait que les tensor cores exigent des pipelines matériels profonds même pour les ressources locales
- La plus grande difficulté est la complexité du memory layout
- Un layout de shared memory non swizzled donne un coalescing très mauvais et exige beaucoup de bande passante L2
- Le layout swizzled était mal documenté, ce qui a demandé du temps pour le comprendre
- Le layout swizzled semble ne fonctionner que pour certaines formes de matrices et s’accorde mal avec d’autres fonctionnalités de
wgmma.mma_async - Le matériel peut effectuer une transposition de sous-matrices en route vers le tensor core, mais seulement lorsque le layout n’est pas swizzled
- Dans des kernels comme FlashAttention, le TMA et le cache L2 sont assez rapides pour masquer ce problème dans une certaine mesure
- Pour exploiter pleinement le matériel, il faut coalescer les requêtes mémoire et éviter les bank conflicts ; le contrôle du layout est donc important
Shared memory et bank conflicts
- La latence d’un accès unique à la shared memory semble être d’environ 30 cycles, durée pendant laquelle le tensor core d’un SM peut presque effectuer deux multiplications de matrices carrées 32x32
- Les travaux précédents comme FlashAttention se concentraient surtout sur le goulot HBM-SRAM, et ce goulot était effectivement important auparavant
- Avec l’accélération de la HBM et la progression des tensor cores, plus rapide que d’autres parties de la puce, même la petite latence de la shared memory devient quelque chose qu’il faut éliminer ou masquer
- La shared memory est divisée en 32 banks ; sans précaution, des bank conflicts apparaissent
- Si plusieurs morceaux de mémoire différents sont demandés simultanément dans la même memory bank, les requêtes sont sérialisées
- D’après l’expérience, cela peut ralentir un kernel de manière déséquilibrée
- Les register layouts exigés par les instructions WGMMA et MMA peuvent subir des bank conflicts s’ils sont utilisés naïvement
- La solution consiste à réorganiser la shared memory avec divers motifs de swizzling afin d’éviter les conflits
- Lorsque c’est possible, il vaut mieux éviter les déplacements entre registers et shared memory ; lorsqu’ils sont nécessaires, il est préférable d’utiliser du matériel intégré asynchrone comme WGMMA et TMA pour déplacer les données
- Les transferts synchrones utilisant de vrais warps sont l’approche la plus générale, mais proche du pire fallback
Génération d’adresses et TMA
- Sur H100, les tensor cores comme la mémoire sont si rapides que le travail consistant simplement à générer les adresses mémoire à fetcher consomme une part significative des ressources de la puce
- C’est encore plus visible lorsque des motifs interleaved ou de swizzling complexes s’ajoutent
- Le Tensor Memory Accelerator (TMA) de NVIDIA permet de spécifier un layout de tenseur multidimensionnel en mémoire globale/shared memory, de fetcher de façon asynchrone un subtile de ce tenseur, puis de déclencher une barrière à la fin
- Le TMA réduit le coût de génération des adresses et facilite aussi la construction du pipeline
- Le TMA est considéré, comme
wgmma.mma_async, comme indispensable pour exploiter le potentiel du H100- D’après l’expérience, il pourrait même être plus important que WGMMA
- Il économise des ressources de registers et de dispatch d’instructions
- Il propose aussi une fonctionnalité de reduction asynchrone vers la mémoire globale, utile dans les kernels backward complexes
- Le TMA a lui aussi nécessité un peu de reverse engineering pour comprendre son mode de swizzling, mais cela a été moins pénible que pour WGMMA
L’occupancy masque les coûts
- Dans CUDA, l’occupancy désigne le nombre de threads co-scheduled sur le même matériel d’exécution
- Le warp scheduler d’un quadrant de SM essaie, à chaque cycle, d’émettre une instruction vers un warp prêt à la recevoir
- Le H100 dépend moins de l’occupancy que les générations précédentes sur certains aspects
- Grâce aux fonctionnalités asynchrones, même un seul flux d’instructions peut maintenir simultanément occupés les fetchs mémoire, les multiplications de matrices, les reductions en shared memory et les calculs dans les registers
- Mais l’occupancy reste très utile pour masquer les erreurs et les coûts de synchronisation
- Un pipeline parfaitement conçu peut être rapide sans occupancy supplémentaire
- En pratique, les observations suggèrent que les GPU NVIDIA ont été conçus en tenant compte de l’occupancy
- Comme il existe beaucoup de synchronisation et de risques d’erreur, augmenter l’occupancy améliore souvent l’utilisation effective du matériel
- Sur H100, l’occupancy est utile ; sur A100 et RTX 4090, elle semble respectivement encore plus importante
- Cela pourrait s’expliquer par une dépendance plus forte au dispatch synchrone des instructions que sur H100
ThunderKittens : un petit DSL dans CUDA
- ThunderKittens est un DSL intégré à CUDA conçu pour écrire plus facilement des kernels rapides sur H100
- Il a d’abord été créé pour un usage interne au laboratoire, puis publié
- Le nom vient du fait que les kittens sont mignons et qu’il semblait amusant de taper
kittens::dans le code - ThunderKittens vise la simplicité et fournit quatre types templated
- Register tiles : tenseurs 2D sur le register file
- Register vectors : tenseurs 1D sur le register file
- Shared tiles : tenseurs 2D dans la shared memory
- Shared vectors : tenseurs 1D dans la shared memory
- Les tiles sont paramétrées par height, width et layout
- Les register vectors sont paramétrés par length et layout, et les shared vectors n’utilisent que length
- Les shared vectors ne subissent généralement pas de bank conflicts
- Les opérations fournies manipulent tiles et vectors au niveau du warp ou au niveau d’un warp group coopératif
- initializer : par exemple mettre un shared vector à zéro
- unary op : par exemple
exp - binary op : par exemple
mul - row/column op : par exemple
row_sum
- ThunderKittens étant intégré dans CUDA, son abstraction « échoue avec élégance », contrairement à des bibliothèques comme Triton
- S’il manque une fonctionnalité, on peut l’étendre comme on le souhaite
Exemple FlashAttention et performances
- Un kernel FlashAttention forward simple pour RTX 4090 est présenté comme exemple ThunderKittens
- Il ne gère que headdim=64
ndoit être un multiple de 256- Il est écrit en environ 60 lignes de code CUDA
- Le taux d’utilisation du matériel est de 75 %
- La majeure partie de la complexité vient de l’algorithme lui-même, et non du motif de swizzling ou du register layout
- Le forward pass FlashAttention-2 pour H100 a également été écrit avec ThunderKittens
- ThunderKittens encapsule la complexité du TMA, de WGMMA, des modes de swizzling et des descriptors
- Le kernel compte environ 100 lignes
- Sur H100, il est environ 30 % plus rapide que FlashAttention-2
- ThunderKittens fournit des primitives et encapsule layouts et instructions, comme un « mini-pytorch » utilisable sur GPU
- Based linear attention et des kernels pour d’autres architectures à paraître sont également publiés
- Le kernel Based linear attention fonctionne à 215 TFLOPs
- En tenant compte du recompute propre à l’algorithme, il dépasse 300 TFLOPs
- La linear attention est théoriquement plus efficace, mais son efficacité sur le matériel réel a historiquement été nettement plus faible
- Ce résultat pourrait élargir le champ des applications à haut throughput
Penser en termes de tiles
- ThunderKittens fonctionne bien parce qu’il n’essaie pas de tout faire
- CUDA est beaucoup plus expressif que ThunderKittens
- ThunderKittens est un petit DSL simple
- L’abstraction centrale est le small tile, ce qui correspond à la direction prise par l’IA et le matériel
- ThunderKittens ne prend pas en charge les dimensions inférieures à 16
- Le matériel ne semble pas particulièrement vouloir de dimensions aussi petites non plus
- La question est posée en ces termes : « si une multiplication de matrices est plus petite que 16x16, peut-on vraiment être sûr qu’il s’agit d’IA ? »
- La vision héritée de l’ère CPU, où un word 32-bit est vu comme un register, ne convient pas au matériel IA
- Le vector register 1024-bit de CUDA est considéré comme un pas dans la bonne direction
- Ici, un register contient les données d’un tile 16x16
- L’IA reste centrée sur les multiplications de matrices, les reductions et les reshapes ; l’abstraction en tiles convient donc à la fois à l’IA et au matériel
- À l’avenir, il faudra réorganiser les idées d’IA de façon à bien les mapper sur le matériel
- La taille de l’état récurrent doit être assez grande pour tenir dans un SM
- La densité de calcul ne doit pas être inférieure à ce que le matériel exige
- Adapter ce que le matériel nous apprend à la conception de l’IA sera une direction importante
Projet de support AMD
- La prise en charge du matériel AMD par ThunderKittens devrait arriver prochainement
1 commentaires
Avis sur Hacker News
La question « si la multiplication de matrices est plus petite que 16x16, est-ce vraiment de l’IA ? » est intéressante.
Les exigences du matériel IA deviennent de plus en plus nettes. Les GPU ont été conçus à l’origine pour un usage complètement différent, mais ils ont été utilisés pour l’IA parce que leur matériel de multiplication de matrices était performant ; un « GPU IA » pourrait retirer certaines fonctions présentes dans les vrais GPU.
On observe aussi une tendance vers des représentations numériques plus courtes, comme le flottant 16 bits, le 8 bits, le 2 bits ou le 1 bit, et un point d’équilibre finira sans doute par s’imposer. Cet article montre qu’un matériel favorisant les tuiles 16x16 est assez pertinent. Quelqu’un est probablement déjà en train d’écrire ce genre de chose en VHDL, ou le sera bientôt.
Au final, il semble probable qu’apparaissent des dispositifs plus simples, moins généralistes et moins chers, capables d’exécuter uniquement les opérations « IA » avec le moins possible de surcharge matérielle inutile.
Nvidia travaille probablement aussi dessus, mais d’un point de vue commercial, il peut être plus judicieux de conserver un dispositif qui regroupe jeux/entertainment/cryptomonnaies/IA, autrement dit le format carte vidéo.
[1] https://github.com/hollance/neural-engine/blob/master/docs/a...
Cela rappelle l’époque où Naveen Rao, avec Nervana, créait un pilote Nvidia Maxwell plus rapide que le propre pilote de Nvidia. Toutes les erreurs de documentation d’un produit en forte croissance ne sont pas des contre-mesures concurrentielles, mais vu le temps qu’il a fallu aux chercheurs pour rétroconcevoir wgmma, et compte tenu de la situation politique sino-américaine autour du H100, on a l’impression que Nvidia réutilise d’anciennes méthodes pour protéger son moat.
Donc plutôt que de trop s’attarder sur les particularités du H100, il faut considérer que la question « quel matériel l’IA veut-elle ? » inclut aussi le contexte commercial.
https://www.amd.com/en/products/accelerators/alveo/v80.html
XDNA Architecture
https://www.amd.com/en/technologies/xdna.html
Le passage « Le mensonge de NVIDIA. C’est une représentation extrêmement trompeuse de la véritable disposition 128b swizzled wgmma. Comme ce schéma m’a fait perdre trois semaines de vie irrécupérables, je l’affiche publiquement à la honte » est marquant.
Je me demande combien de gens seraient surpris d’apprendre qu’une énorme part des progrès de l’IA relève de l’ingénierie, comme l’optimisation de la multiplication de matrices, et qu’une bonne partie de cette ingénierie consiste à rétroconcevoir des puces NVIDIA.
Ordonnanceur de warps, quatre quadrants, Tensor Memory Accelerator, disposition wgmma non swizzled…
La frontière entre le jargon GPU et le technoblabla façon Star Trek devient de plus en plus floue.
J’ai parfois eu la même pensée en lisant d’autres articles. Je me demande ce que ressentirait quelqu’un à qui l’on enverrait un lien vers un article d’ici. Il aurait sans doute l’impression d’entrer dans une convention de fans de Trek en train de débattre du cœur de distorsion.
Pour réduire la consommation électrique et augmenter la vitesse de l’inférence IA, le mieux serait probablement de passer à des circuits d’approximation analogiques
On n’a pas besoin d’une multiplication et d’une addition en virgule flottante parfaites ; il suffit d’un dispositif qui prenne deux tensions en entrée et produise une tension de sortie suffisamment proche du résultat de la multiplication
Le gros avantage, c’est qu’au lieu de représenter un float16 par 16 lignes, on représente ce nombre par la tension d’une seule ligne. En théorie, une précision bien supérieure à celle du float32 pourrait aussi être possible. Et comme les valeurs peuvent être connectées directement sans être chargées dans une unité arithmétique et logique, les économies potentielles en surface de die et en énergie pourraient atteindre plusieurs ordres de grandeur
Par exemple, accepter qu’un bit de sortie sur un million soit inversé, en échange d’un meilleur rapport performance/énergie. Ce serait difficile avec du float32, où une seule valeur infinie peut tout casser, mais avec de l’int8, on peut probablement tolérer qu’au lieu de 0, on obtienne parfois 128
[1] Je ne sais pas vraiment si les unités matricielles en virgule flottante du H100 respectent effectivement IEEE 754
Les réseaux neuronaux biologiques ne sont pas presque entièrement connectés comme les réseaux de neurones artificiels classiques ; les coefficients de connexion d’entrée/sortie des neurones sont inférieurs à 10, donc très locaux. En biologie, à notre connaissance, il n’y a pas non plus de rétropropagation, mais plutôt du feedback et des boucles récurrentes
Il pourrait aussi exister des cellules auxiliaires ou des processus, encore inconnus, essentiels au fonctionnement du système nerveux central. Même à haut niveau, il y a probablement une quantité importante de connectivité « codée en dur », dont une partie est déjà connue. Par exemple, les neurones auditifs de l’oreille sont connectés, et quelque chose de similaire à une convolution se produit pour localiser la provenance d’un son. Ce n’est pas un phénomène émergent, mais une fonction possible même sans entraînement
Ce n’est pas surprenant : le vivant a découvert cela au fil de milliards d’années et d’un nombre comparable de générations. En théorie, ce serait aussi possible en logiciel, mais avec plus de mille milliards de neurones dans le cerveau des primates/humains, ce serait extrêmement difficile même sur les machines actuelles à quelques milliers de cœurs. Même le « cloud » ne satisferait probablement pas la connectivité et la latence nécessaires
Ce serait déjà formidable de réussir à modéliser, avec une telle approche, quelque chose du niveau d’un ver ou d’un insecte
Cet article me rappelle le plaisir que j’avais ressenti dans le cours CS 149 de programmation parallèle
Le style de cet article est vraiment impressionnant, et j’ai hâte de voir ça sur AMD MI300x. Si vous voulez utiliser du temps sur ma machine, dites-le-moi
Je me demande à quel point ça fonctionnerait réellement bien, s’il vaudrait mieux économiser un peu plus pour prendre une XTX plutôt qu’une 7900 XT, et dans quelle mesure la baisse de VRAM affecterait l’usage réel
Le lecteur ne devrait pas avoir à aller sur knowyourmeme.com pour comprendre ce que les auteurs essaient de dire. Je ne sais même pas ce que ce titre veut dire, et à mes yeux cela montre à quel point il rate sa cible
Je me demande par où commencer, et quelle feuille de route suivre, pour comprendre complètement ce genre d’article
Ensuite, il est utile d’écrire soi-même un kernel CUDA qui effectue une multiplication vecteur-matrice. Avec pycuda, on peut se concentrer sur le kernel et écrire le reste en Python. Il suffit de dire à ChatGPT que vous voulez créer vous-même une implémentation qui multiplie un vecteur de 4 000 éléments par une matrice 4000x12000, et de lui demander de vous guider tout au long du processus
Pour louer des GPU, Runpod est une bonne option ; actuellement, il y a de tout, des GPU d’entrée de gamme jusqu’aux H100. Au début, un GPU bas de gamme suffit
J’ai passé 2 mois à implémenter et optimiser des kernels de multiplication matricielle avec Spiral
Le graphique du README GitHub (https://github.com/HazyResearch/ThunderKittens/blob/main/att...) est beaucoup trop étourdissant. Ce genre d’histogramme ondulé est-il seulement légal ? :P
[1]: https://matplotlib.org/stable/gallery/showcase/xkcd.html#sph...
Le nom ThunderKittens est excellent. J’aimerais voir ThunderKittens s’attaquer à la rétropropagation de FlashAttention, qui est d’un ordre de grandeur plus difficile que la passe avant.
causale : https://github.com/HazyResearch/ThunderKittens/blob/main/exa...
non causale : https://github.com/HazyResearch/ThunderKittens/blob/main/exa...
Ce genre de recherche n’est-il pas déjà mené aujourd’hui par les équipes qui conçoivent des NPU ? Par exemple, les puces Groq peuvent atteindre les performances actuelles parce qu’elles utilisent une architecture dédiée à l’IA. Côté grand public, Apple Silicon est aussi plutôt performant.
Je ne suis pas du domaine, mais il me semble qu’on atteint vite les limites avec uniquement des processeurs généralistes qui communiquent via des chemins relativement lents. Repenser la conception au niveau matériel, puis à terme faire baisser les prix sur le marché grand public, me paraît être une meilleure stratégie à long terme.
Quand on peut acheter un GPU Nvidia pour quelques centaines de dollars, ou un PC portable gaming avec une 4050 et 6 Go de VRAM pour 900 dollars, il est difficile de qualifier l’IA sur CPU de performante.
Au travail, je n’avais pas non plus de GPU, donc j’ai essayé sur CPU, mais à part utiliser de petits modèles et attendre, ce n’était pas réaliste. J’ai fini par demander un ordinateur avec GPU.
« Techniquement possible » et « agréable à utiliser en pratique » sont deux choses différentes. Nvidia était vraiment agréable à utiliser ; le CPU était pénible et frustrant.