Huit bits, puis quatre

Troisième article de la série Smallem : servir un modèle. Précédent : Générer à plusieurs.

Chaque poids du modèle est un f32 : quatre octets. Générer un token les lit tous, et le premier article a montré que la borne de la génération est la bande passante. Ranger les poids sur un octet, ou un demi, en lirait quatre ou huit fois moins.

C'est la quantification. Elle a une seconde raison d'être, plus terre à terre : la mémoire. Un modèle de 70 milliards de paramètres en 16 bits fait 140 Go, en 4 bits, 35 Go. Pour nos 29 millions de paramètres, ce n'est pas un problème. Mais la question reste la même : que perd le modèle quand on arrondit ses poids ?

Arrondir

Le principe est simple. On prend un groupe de poids, on cherche la plus grande valeur absolue, et on découpe l'intervalle en autant de pas que le format a de valeurs :

échelle = max |w| / 127          (int8 ; 7 en int4)
code    = arrondi(w / échelle)   un entier entre −127 et 127
w'      = code × échelle         le poids que le modèle utilise ensuite

On range le code, un petit entier, et une échelle par groupe. L'erreur d'un poids est au plus d'un demi-pas.

Toute la question est la taille des groupes. Une échelle par colonne de la matrice, un « canal de sortie », suffit en 8 bits : 255 valeurs couvrent bien l'intervalle d'une colonne. En 4 bits, il n'en reste que 15, et un seul poids un peu plus grand que les autres élargit le pas pour toute la colonne. On prend alors des groupes de 32 à 128 poids consécutifs, chacun avec son échelle.

Ranger des entiers dans des f32

Les tenseurs du projet ne connaissent qu'un type : f32 sur la carte, f32 ou f64 sur CPU. Ajouter un type entier aux backends serait un gros chantier pour un seul usage.

L'article précédent avait déjà rangé des numéros de lignes dans des f32, en notant qu'un f32 représente exactement tout entier jusqu'à 2²⁴. On peut donc y mettre 24 bits de codes : trois codes de 8 bits, ou six de 4 bits. Le kernel relit le mot comme un entier, puis en extrait chaque code par décalage :

let word = words[index] as u32;
let stored = (word >> (bits * slot)) & mask;
let weight = (stored as i32 - offset) as f32 * scale;

Le prix est d'un quart : 1,33 octet par poids en int8 au lieu de 1, 0,67 en int4 au lieu de 0,5. En échange, les mêmes tenseurs, les mêmes transferts, et une définition de référence qui tourne sur tous les backends.

Une opération QuantizedMatmul reçoit les codes, les échelles, et les mêmes poids déquantifiés en f32. La carte lit les codes dans les produits de la génération, où l'on attend la mémoire, et la forme déquantifiée dans les gros produits d'une évaluation, où l'on calcule plus qu'on ne lit. Les deux donnent les mêmes nombres, aux arrondis de l'ordre des sommes près, et les tests le vérifient pour une ligne, neuf, quarante et trois cents.

Ce choix a un coût qu'il faut dire : notre modèle quantifié garde aussi sa copie en f32. La mémoire des poids donnée plus bas est celle de la forme compacte, ce qu'un serveur garderait mais la nôtre occupe encore la carte.

Seules les couches linéaires sont quantifiées : les attentions et les feed-forward, 25 millions de poids. La table d'embedding, qui sert aussi à la sortie, reste en f32.

Ce que perd le modèle

La perte de validation se mesure sur les 200 lots de la mesure finale d'un entraînement. Le modèle non quantifié retrouve exactement son chiffre de l'article 16 de la première série : 0,8823 nat par octet.

L'int8 ne coûte rien de mesurable : l'écart est de 0,00004, vingt fois sous le bruit de 0,002 entre deux entraînements. En génération gloutonne, il écrit exactement le même texte que le f32, sur 300 tokens.

L'int4 coûte de trois à cinq fois le bruit. Des groupes plus petits aident, mais chaque échelle coûte quatre octets : avec des groupes de 32 plutôt que 64, le modèle grossit de 9 %. Et le texte change dès le premier token généré : un modèle quantifié est un autre modèle, proche, mais qui ne choisit pas toujours le même token.

Arrondir mieux : GPTQ

Arrondir chaque poids au plus proche minimise l'erreur sur chaque poids. Ce n'est pas ce qu'on veut. On veut que la sortie de la couche change le moins possible, sur des entrées typiques :

minimiser  ‖ X · W − X · W' ‖²      X : des entrées réelles de la couche

Les deux ne sont pas la même chose quand les composantes de x sont corrélées. Si l'erreur d'arrondi d'un poids pousse la sortie vers le haut, on peut la compenser en arrondissant vers le bas un autre poids, dont la composante d'entrée varie avec la première.

C'est l'idée de GPTQ (Frantar et al., 2022). Les lignes de W, une par composante de l'entrée, sont quantifiées l'une après l'autre. L'erreur de chacune est aussitôt reportée sur les lignes qui restent, dans la mesure où leurs entrées sont corrélées à la sienne. Tout ce qu'il faut savoir des entrées tient dans une matrice [K, K], H = Σ xᵀ · x, accumulée sur quelques lots. Et le facteur de report se lit dans le facteur de Cholesky de son inverse, H⁻¹ = Uᵀ · U :

pour chaque ligne k :
    q_k     = arrondi(w_k)
    erreur  = (w_k − q_k) / U[k, k]
    w_j    -= erreur × U[k, j]        pour chaque ligne j > k

C'est l'algorithme 1 de l'article, sans le traitement par blocs qui ne sert qu'à occuper une carte : nos matrices font au plus 1368 × 512, et l'hôte les traite en f64 en quelques secondes. Comme dans l'article, un amortissement de 1 % de la moyenne de la diagonale est ajouté au hessien, et les échelles d'un groupe sont fixées quand on atteint sa première ligne, d'après les poids déjà corrigés.

Pour obtenir H, chaque matrice de poids peut maintenant observer ses entrées : pendant un étalonnage, la couche ajoute xᵀ · x à un accumulateur avant chaque produit. Huit lots d'entraînement de 16 séquences suffisent, 65 000 tokens.

Deux tests vérifient l'algorithme. Quand les entrées sont indépendantes, H est diagonal, rien ne peut compenser rien, et GPTQ doit donner exactement l'arrondi au plus proche. Il le donne. Quand elles sont corrélées, l'erreur sur la sortie doit baisser : de plus de 30 % sur un exemple synthétique.

GPTQ divise l'écart par deux, sans rien changer au format ni au kernel : ce sont les mêmes codes, mieux choisis.

AWQ (Lin et al., 2023) prend un autre chemin. Il remarque qu'un pour cent des poids comptent plus que les autres, ceux dont la composante d'entrée est grande, et les protège en les multipliant par une échelle s avant d'arrondir, l'entrée étant divisée d'autant : Q(w · s) · x / s. Leur erreur relative diminue. Mais la division de l'entrée doit être reportée dans la couche précédente : le gain de la normalisation, les valeurs de l'attention, ou la couche up du feed-forward, selon le cas. GPTQ ne touche qu'à la matrice qu'il quantifie. Il était le plus simple à écrire ici.

Des valeurs aberrantes, à 29 millions de paramètres ?

LLM.int8() (Dettmers et al., 2022) avait trouvé pourquoi l'int8 naïf échouait sur les grands modèles. Quelques composantes des entrées des couches prennent des valeurs énormes, toujours les mêmes, d'une couche à l'autre. Les auteurs les définissent ainsi : une composante qui atteint au moins 6, dans au moins un quart des couches, sur au moins 6 % des positions. Elles sont de 3 à 20 fois plus grandes que les autres, et à partir de 6,7 milliards de paramètres, toutes les couches en ont. Arrondies avec le reste, elles écrasent la précision de tout le vecteur.

L'étalonnage sait tenir les mêmes statistiques. Sur deux lots, à l'entrée des attentions et des feed-forward :

Trois composantes reviennent dans presque toutes les couches parmi les plus grandes : 112, 257 et 285. Dans les deux dernières couches, elles dépassent 6 sur 4 à 19 % des tokens. Les critères de LLM.int8() sont donc remplis, de justesse : trois couches sur huit, plus de 6 % des positions. Mais ces composantes ne sont que 2,5 fois plus grandes que la médiane. Ce sont des valeurs aberrantes naissantes, pas encore celles qui ruinent l'int8 : notre int8 n'a rien perdu. On voit comment elles grandissent avec la profondeur.

La vitesse

C'était la promesse de départ : lire moins pour aller plus vite.

Tokens par seconde. Rien, ou presque : l'int8 est même 7 % plus lent.

C'est la conclusion du premier article qui revient. Nous sommes à 13 % de la borne de la bande passante : le temps d'un token n'est pas celui de la lecture des poids, mais celui de 245 kernels courts, dont le coût fixe domine. Un produit d'un vecteur par une matrice 512 × 512 prend 7,2 µs en f32, pour un mégaoctet qui se lit en 1,1. Lire trois fois moins d'octets ne retire qu'une partie de ces 1,1 µs, et décoder les codes ajoute des instructions. L'int8 range trois codes par mot, l'int4 six : pour un même nombre de poids, l'int8 lit deux fois plus de mots que l'int4, et il perd.

Il a fallu trois versions du kernel pour en arriver là. La première lisait un mot à la fois, et l'attendait. La deuxième en lisait quatre d'un coup, mais refaisait une division entière par mot pour trouver son groupe : plus lente encore. La dernière suit le groupe par incréments, décode les six poids d'un mot une seule fois, puis les applique aux lignes. Les tuiles quantifiées des lots, elles, calculaient la position de chaque code séparément ; elles la calculent maintenant une fois pour quatre colonnes voisines.

La quantification paiera quand les kernels attendront vraiment la mémoire : avec des kernels fusionnés, ou un modèle cent fois plus gros, où chaque produit lit des centaines de mégaoctets. Ici, elle ne rapporte que de la place.

Ce qu'il faut retenir

• Quantifier, c'est arrondir des groupes de poids sur une grille de quelques valeurs. En int8 par colonne, notre modèle ne perd rien de mesurable ; en int4, il perd de trois à cinq fois le bruit.
• Arrondir chaque poids au plus proche n'est pas le meilleur arrondi. GPTQ choisit les codes pour préserver la sortie de la couche, et divise l'écart de l'int4 par deux.
• Des valeurs aberrantes apparaissent déjà dans les dernières couches d'un modèle de 29 millions de paramètres, mais encore trop petites pour gêner.
• Lire moins ne va plus vite que si l'on attendait la lecture. Nos kernels attendent surtout leur propre lancement.

Article suivant : Deviner, puis vérifier.