Mise à l'échelle de FlashAttention, entraînement multi-GPU et quantification des poids

L'infrastructure d'IA moderne s'appuie sur des noyaux d'attention adaptés au matériel, le parallélisme distribué multi-GPU et la quantification à faible nombre de bits pour entraîner et servir efficacement des modèles massifs. Ces trois techniques comblent l'écart entre l'échelle théorique des modèles et les limites physiques des GPU.

1. FlashAttention-2 et FlashAttention-3 sur les GPU Hopper (H100)

Si le premier FlashAttention a résolu le goulet d'étranglement de la bande passante mémoire, il laissait encore une part importante de la puissance de calcul brute des GPU modernes inutilisée. FlashAttention-2 (2023) et FlashAttention-3 (2024) ont repensé la coordination des threads GPU et des moteurs matériels.

• FlashAttention-2 (meilleure répartition du travail).

• Réduction des opérations autres que les multiplications matricielles. Les Tensor Cores des GPU exécutent les multiplications matricielles (matmul) environ 16× plus vite que les cœurs CUDA standard n'exécutent les opérations scalaires comme la remise à l'échelle Softmax. FlashAttention-2 a réécrit les calculs du Softmax en ligne pour reporter la remise à l'échelle à la toute fin de la boucle interne, réduisant les instructions lentes autres que matmul.
• Parallélisation sur la longueur de séquence. La version initiale parallélisait surtout le travail sur la taille des lots et les têtes d'attention, laissant des cœurs GPU inactifs lors d'un entraînement avec de longues séquences et de petits lots. FlashAttention-2 ajoute du parallélisme sur la dimension de longueur de séquence et inverse l'ordre des boucles (boucle externe sur les blocs Q, boucle interne sur K et V), ce qui minimise les écritures lentes en mémoire partagée.

• FlashAttention-3 (exploitation du matériel Hopper H100).

• Déplacement asynchrone des données (TMA). L'architecture Hopper de NVIDIA a introduit le Tensor Memory Accelerator (TMA), une unité matérielle dédiée qui copie les données entre la HBM globale et la SRAM sur puce sans monopoliser les threads de calcul.
• Spécialisation des warps (modèle producteur-consommateur). FlashAttention-3 attribue des rôles dédiés aux warps GPU (groupes de 32 threads). Les warps producteurs utilisent le TMA pour précharger les futures tuiles de K et V en SRAM, tandis que les warps consommateurs alimentent les Tensor Cores avec les tuiles courantes (instructions WGMMA). Les transferts mémoire se chevauchent ainsi avec les calculs matriciels et les calculs Softmax.
• Prise en charge de la faible précision FP8. FlashAttention-3 exploite les Tensor Cores FP8 de Hopper, doublant le débit brut par rapport à FP16. Pour éviter la perte de précision due aux activations aberrantes en FP8, il utilise la quantification par blocs (mise à l'échelle indépendante de petites tuiles) et le traitement incohérent (multiplication de Q et K par des matrices orthogonales aléatoires pour atténuer les valeurs extrêmes avant la quantification).

2. Stratégies de mise à l'échelle multi-GPU (parallélisme 3D)

Lorsqu'un modèle possède des dizaines ou des centaines de milliards de paramètres, ses poids, gradients et états d'optimiseur dépassent la VRAM d'un seul GPU. L'entraînement à grande échelle combine trois stratégies distinctes de parallélisation, souvent appelées parallélisme 3D.

• Parallélisme de données & ZeRO / FSDP. Dans le parallélisme de données standard, chaque GPU conserve une copie complète du modèle, gaspillant de la mémoire. Les frameworks modernes utilisent ZeRO (Zero Redundancy Optimizer) ou PyTorch FSDP (Fully Sharded Data Parallel), qui répartit les paramètres, les gradients et les états d'optimiseur entre tous les GPU, en rassemblant les poids de chaque couche juste à temps pour le calcul.

• Parallélisme de tenseurs (intra-couche). Au lieu de dupliquer une multiplication matricielle massive de réseau à propagation avant Y = XW, le parallélisme de tenseurs divise W par colonnes entre les GPU (W = [W_1, W_2]). Chaque GPU calcule simultanément sa moitié de la sortie et échange des résultats partiels via NVLink à haut débit (900 GB/s sur les H100).

• Parallélisme de pipeline (inter-couches). Les couches 1–8 sont sur le GPU 1, les couches 9–16 sur le GPU 2, et ainsi de suite. Pour éviter que les GPU en aval restent inactifs en attendant que les couches précédentes terminent (la « bulle de pipeline »), les frameworks découpent chaque lot d'entraînement en micro-lots afin de maintenir tous les GPU occupés comme sur une chaîne de montage.

3. Quantification (GPTQ, AWQ) et noyaux de déquantification à la volée

Pendant l'inférence (génération des tokens un par un), les grands modèles de langage sont limités par la bande passante mémoire, et non par le calcul. Charger 70 milliards de paramètres de 16 bits (140 GB) depuis la VRAM pour générer un seul token laisse les Tensor Cores en attente des transferts mémoire. La quantification compresse les poids de nombres flottants de 16 bits en entiers de 4 bits (INT4), réduisant la taille mémoire et le temps de transfert d'un facteur proche de 4×.

• GPTQ (quantification après entraînement par calibration hessienne).

• Arrondir naïvement les poids de 16 bits en entiers de 4 bits détruit la précision du modèle. GPTQ fait passer un petit jeu de données de calibration dans le modèle et utilise des calculs du second ordre (la matrice hessienne inverse) pour mesurer la sensibilité de chaque poids.
• Il quantifie les poids colonne par colonne ; chaque fois que l'arrondi d'un poids introduit une erreur, GPTQ ajuste mathématiquement les poids restants non quantifiés de cette ligne pour annuler l'erreur.

• AWQ (quantification des poids tenant compte des activations).

• Les chercheurs ont découvert qu'environ 0,1 % à 1 % des canaux de poids sont bien plus importants que les autres, car ils multiplient de grandes valeurs d'activation.
• Au lieu de conserver ces poids essentiels en 16 bits (ce qui crée des organisations mémoire en précision mixte complexes et lentes), AWQ augmente l'échelle de ces canaux de poids saillants avant la quantification et réduit celle des activations d'entrée correspondantes à l'exécution, préservant la précision tout en gardant 100 % des poids dans un format uniforme de 4 bits.

• Comment les noyaux CUDA personnalisés déquantifient à la volée (W4A16).

• Les GPU ne peuvent pas multiplier directement une activation de 16 bits (FP16) par un poids entier de 4 bits (INT4). Des noyaux d'inférence spécialisés (comme Marlin ou ExLlama) réalisent donc une exécution W4A16 (poids de 4 bits, activations de 16 bits) :
• 1. Lecture HBM compressée : le noyau charge les poids de 4 bits empaquetés (huit valeurs INT4 dans un seul entier de 32 bits) depuis la VRAM globale lente, divisant par 4× l'utilisation de la bande passante mémoire.
• 2. Dépaquetage dans les registres : dans les registres ultrarapides du GPU, le noyau CUDA utilise des décalages binaires rapides (>>) et des masques binaires (&) pour dépaqueter les entiers INT4, puis applique les facteurs d'échelle/de point zéro stockés pour les reconvertir en valeurs FP16 à la volée.
• 3. Calcul immédiat sur Tensor Cores : les poids FP16 fraîchement dépaquetés alimentent directement les Tensor Cores avec les activations FP16 sans jamais réécrire les poids décompressés de 16 bits en VRAM.