Rendre le GPU rapide

Onzième article de la série Smallem. Précédent : Des sous-mots et un entraînement à l'échelle.

Les kernels de l'article 5 étaient justes et lents. Cet article raconte comment le temps d'un pas d'entraînement du plus gros modèle est passé de 652 à 157 millisecondes. C'est aussi l'article où l'on comprend vraiment comment fonctionne une carte graphique, parce que chaque gain vient d'une propriété précise du matériel.

Le travail s'est fait en deux campagnes. La première a eu lieu juste après l'assemblage du Transformer, avant le premier vrai entraînement. La seconde après l'entraînement à l'échelle de l'article précédent.

La règle : mesurer d'abord

L'intuition est mauvaise conseillère en matière de performance. Trois précautions ont guidé tout le travail.

• Un banc d'essai, toujours le même, qui donne un nombre : des tokens par seconde.
• Synchroniser avant de lire l'horloge. On l'a vu à l'article 5 : un kernel lancé n'est pas un kernel terminé.
• Écarter les premiers pas, qui paient des coûts uniques : chargement des kernels, premières allocations.

Et une quatrième, apprise en route : une seule mesure peut être fausse de 30 %. Les nombres finaux ont été répétés trois fois.

Première campagne

La réserve de tampons

Le premier gain n'a rien à voir avec le calcul. Chaque opération crée un tenseur, donc alloue de la mémoire sur la carte, et la libère quand le tenseur disparaît. Un pas d'entraînement en fait des centaines.

Allouer coûte cher. Mais libérer coûte bien plus, pour une raison subtile : un kernel encore en file pourrait être en train d'utiliser cette mémoire. Avant chaque libération, il fallait donc attendre que la carte ait tout fini. Le programme passait son temps à attendre, et l'asynchronisme vanté à l'article 5 était réduit à néant.

La solution : ne plus libérer. Un tampon abandonné retourne dans une réserve, rangée par taille. Un entraînement redemande exactement les mêmes tailles à chaque pas, donc presque toutes les demandes sont servies depuis la réserve.

impl<T> Drop for Pooled<T> {
    fn drop(&mut self) {
        if let Some(buffer) = self.buffer.take() {
            self.pool
                .free
                .borrow_mut()
                .entry(buffer.len())
                .or_default()
                .push(buffer);
        }
    }
}

Réutiliser un tampon qu'un kernel en file lit encore est-il dangereux ? Non, et la raison est la même qu'à l'article 5. Tout passe par un seul stream : le kernel qui écrira dans le tampon réutilisé est mis en file après celui qui le lit, donc exécuté après lui.

Résultat : la grande configuration passe de 606 à 336 ms par pas. Presque la moitié du temps partait dans la gestion de la mémoire.

La matmul par tuiles

La matmul naïve lançait un thread par élément de sortie. Chaque thread relisait en mémoire globale toute une ligne de a et toute une colonne de b. Deux threads voisins relisaient la même ligne, chacun de son côté.

Un GPU possède plusieurs sortes de mémoire. La mémoire globale est grande et lente. Chaque bloc de threads dispose aussi d'une petite mémoire partagée, bien plus rapide, que ses threads sont seuls à voir.

D'où l'idée des tuiles. Un bloc de threads se charge d'une tuile carrée de la sortie, et ses threads coopèrent :

1. ensemble, ils copient en mémoire partagée la portion de a et la portion de b dont ils ont besoin ;
2. ils attendent que tout le monde ait fini de copier ;
3. chacun calcule en lisant la mémoire partagée ;
4. ils attendent que tout le monde ait fini de lire, puis passent à la portion suivante.

Chaque valeur n'est lue qu'une fois en mémoire globale, au lieu d'une fois par thread qui en a besoin.

Les étapes 2 et 4 sont des barrières : thread::sync_threads() bloque chaque thread jusqu'à ce que tous ceux du bloc y soient arrivés. C'est la première fois que des threads s'attendent. Sans barrière, un thread rapide lirait une case que son voisin n'a pas encore remplie.

Résultat : de 336 à 151 ms par pas.

Les kernels fusionnés

Restaient des dizaines de petites opérations. Une LayerNorm écrite avec les primitives, c'est une moyenne, une soustraction, un carré, une autre moyenne, une racine, une division : une dizaine de lancements, et autant de tenseurs intermédiaires qui traversent la mémoire.

Trois opérations ont reçu leur propre kernel : GELU, la normalisation, et le softmax causal de l'attention. Ce dernier ne construit même plus de masque : le kernel calcule, à partir de la position du thread, s'il regarde le futur.

let position = width - queries + row % queries;
*out_elem = if col > position {
    0.0
} else {
    ((a[flat] * scale - row_max[row]) - row_log_sum[row]).exp()
};

position est la place de la requête dans la séquence. À l'entraînement, il y a autant de requêtes que de positions, et la requête i est à la position i. Le kernel accepte aussi moins de requêtes que de positions : ce sont alors les dernières de la séquence. Ce détail servira à la génération.

C'est ici que servent les opérations de l'article 4. Chacune vit dans son fichier, avec sa référence écrite avec les primitives. Lui donner un kernel, c'est dire, dans le paquet GPU, ce que la carte fait à la place :

/// Deux kernels : les réductions par ligne, puis un thread par élément.
impl Run<Normalize> for Cuda {
    fn run(&self, op: &Normalize, inputs: &[&Self::Tensor]) -> Self::Tensor {
        let [x] = operands(Normalize::NAME, inputs);
        self.normalize(x, op.epsilon as f32)
    }
    // … et sa rétropropagation, elle aussi avec ses propres kernels
}

Le backend CPU ne change pas : il continue d'exécuter la référence. Le kernel est testé contre elle, pour le calcul comme pour le gradient. La version simple reste la définition, la version rapide n'est qu'une façon d'aller plus vite.

Résultat : de 151 à 123 ms par pas. En tout, la grande configuration est devenue 4,9 fois plus rapide, la petite 6,6 fois. C'est avec ces kernels qu'ont été entraînés les modèles des articles 9 et 10.

Seconde campagne

Un profileur

Avec le modèle de 30 millions de paramètres, un pas durait 652 ms. On souhaite identifier là où le temps est passé.

Un chronomètre autour de chaque opération ne suffit pas, puisque les lancements sont asynchrones. On a donc écrit un backend qui enveloppe le backend GPU et attend la carte après chaque opération :

fn timed<T>(&self, name: &str, operation: impl FnOnce(&Cuda) -> T) -> T {
    let start = Instant::now();
    let result = operation(&self.inner);
    self.inner.synchronize();
    // … ajoute start.elapsed() au total de l'opération `name`
    result
}

Le total est un peu plus lent qu'en temps normal, mais la répartition est juste. Et comme ce profileur est un Backend comme un autre, le modèle tourne dessus sans changer une ligne. C'est encore le contrat de l'article 4 qui paie.

Il paie même deux fois. Les primitives sont déléguées une par une, mais les opérations n'ont pas à être énumérées : une seule implémentation les couvre toutes, y compris celles qu'on ajoutera plus tard.

impl<O: Operation> Run<O> for Profiled
where
    Cuda: Run<O>,
{
    fn run(&self, op: &O, inputs: &[&T]) -> T {
        self.timed(O::NAME, |gpu| gpu.run(op, inputs))
    }
    // … et de même pour le gradient, compté à part
}

opération                      appels     ms/pas    part
matmul                            195     556.41   84.5 %
causal_softmax (gradient)           8      20.84    3.2 %
sum                               214      16.33    2.5 %
transpose                         212      14.17    2.2 %

La matmul pèse 84,5 % du temps. Tout le reste est du détail tant qu'elle n'a pas été traitée.

Plusieurs sorties par thread

Dans la matmul par tuiles, chaque thread calculait un élément de sortie. Pour chaque produit, il lisait une valeur de a et une valeur de b en mémoire partagée : deux lectures par multiplication.

On peut faire mieux en donnant à chaque thread un petit pavé de 4 × 4 éléments. Pour un indice k donné, il lit 4 valeurs de a et 4 valeurs de b, soit 8 lectures, et en tire 16 produits :

let a0 = TILE_A[(ty * PATCH) * DEPTH + k];
// … a1, a2, a3
let b0 = TILE_B[k * TILE + tx * PATCH];
// … b1, b2, b3
c00 += a0 * b0;
c01 += a0 * b1;
// … seize produits en tout

Une lecture pour deux produits, au lieu de deux lectures par produit : quatre fois moins. Un bloc de 16 × 16 threads calcule maintenant une tuile de 64 × 64.

Le débit de la matmul passe de 2,8 à environ 11 TFLOP/s (mille milliards d'opérations flottantes par seconde), et le pas de 652 à 258 ms.

Un warp par ligne

Plusieurs kernels réduisaient chaque ligne d'un tenseur : le maximum et la somme d'une ligne pour le softmax, la moyenne et la variance pour la normalisation. Ils lançaient un thread par ligne, qui la parcourait seul.

Cela heurte une propriété du matériel. Sur un GPU, les threads avancent par groupes de 32, les warps. Quand les 32 threads d'un warp lisent 32 cases voisines, la mémoire les sert en une seule transaction. Quand chacun lit dans son coin, il faut 32 transactions. Avec un thread par ligne, c'était le pire cas : les 32 threads d'un warp lisaient 32 lignes différentes, éloignées les unes des autres.

On retourne donc le problème : un warp par ligne. Le thread de rang l parcourt les éléments l, l + 32, l + 64… À chaque tour, les 32 threads lisent 32 éléments voisins.

Restent 32 résultats partiels à réunir. Les threads d'un warp peuvent s'échanger des valeurs directement, de registre à registre, sans passer par la mémoire :

pub fn warp_sum(mut value: f32) -> f32 {
    let mut distance = WARP as u32 / 2;
    while distance > 0 {
        value += warp::shuffle_xor_f32(value, distance);
        distance /= 2;
    }
    value
}

À chaque tour, un thread ajoute à sa valeur celle du thread dont le rang diffère par un bit : 16, puis 8, 4, 2, 1. Après cinq tours, chaque valeur a parcouru tout le warp, et chaque thread tient le total. Il n'y a ni mémoire ni barrière.

Les gains sont nets : la rétropropagation du softmax passe de 20,9 à 6,3 ms, le softmax causal de 9,5 à 3,2 ms. Le pas descend à 224 ms.

Les Tensor Cores

Les cartes récentes contiennent des circuits spécialisés dans la multiplication de matrices, les Tensor Cores. Une seule instruction, mma.sync, multiplie une tuile de 16 × 16 par une tuile de 16 × 8. Elle est exécutée par un warp entier : chacun des 32 threads fournit dans ses registres un fragment des deux tuiles, et reçoit un fragment du résultat.

Il y a une contrepartie. Ces circuits travaillent sur des nombres de 16 bits, pas de 32. Deux formats existent.

Le f16 oblige à multiplier la perte par un facteur pour que les gradients, souvent minuscules, ne tombent pas à zéro. Le bf16 garde les exposants du f32 : il n'a pas ce problème. C'est lui qu'on a retenu.

C'est ce qu'on appelle la précision mixte. Seuls les opérandes de la multiplication sont arrondis sur 16 bits ; les sommes restent en f32, et tous les tenseurs aussi.

L'arrondi doit être fait proprement. Un bf16 est la moitié haute d'un f32. Couper simplement les 16 bits bas tirerait toutes les valeurs vers zéro : un biais systématique, qui s'accumulerait pendant l'entraînement. On ajoute donc un demi-pas avant de couper, ce qui arrondit au plus proche :

pub fn bf16_pair(lo: f32, hi: f32) -> u32 {
    let lo = lo.to_bits();
    let hi = hi.to_bits();
    let lo = lo.wrapping_add(0x7FFF + ((lo >> 16) & 1)) >> 16;
    let hi = hi.wrapping_add(0x7FFF + ((hi >> 16) & 1)) & 0xFFFF_0000;
    lo | hi
}

La difficulté du kernel est ailleurs : l'instruction impose quel thread tient quel élément de chaque tuile, selon une disposition fixe et peu intuitive. Si on se trompe, rien ne plante : on calcule un mauvais produit.

Le test qui protège de cette erreur est simple : des entiers entre -8 et 8 sont représentés exactement en bf16, et leurs sommes restent exactes en f32. Sur de telles entrées, le résultat doit donc être celui du CPU au bit près. La moindre erreur de disposition donnerait d'autres nombres.

#[test]
fn mixed_matmul_is_exact_on_small_integers() {
    // …
    let expected = eval(&cpu, &inputs, |b, t| b.matmul(&t[0], &t[1]));
    let actual = eval(&gpu, &inputs, |b, t| b.matmul(&t[0], &t[1]));
    assert_eq!(expected, actual, "formes {shape_a:?} × {shape_b:?}");
}

Le débit atteint 14 à 22 TFLOP/s sur les couches linéaires, et le pas descend à 157 ms.

Et la qualité ? Sur la grande configuration, 1 500 pas donnent une perte de validation de 1,4287 en f32 et de 1,4286 en précision mixte. L'arrondi ne coûte rien.

Le cache de génération

Jusqu'ici, pour générer un token, le modèle recalculait tout, sur toute la séquence. Pour le 400ᵉ token, il refaisait le travail des 399 précédents.

Or le passé ne change pas. À cause du masque causal, ce que calcule une position ne dépend que d'elle-même et des précédentes. Ajouter un token à la fin ne modifie rien de ce qui a déjà été calculé.

Pour traiter le nouveau token, l'attention a besoin des clés et des valeurs de toutes les positions précédentes. On les garde donc en mémoire d'un tour à l'autre : c'est le cache KV. Seul le nouveau token traverse le modèle.

Il n'y a pas pour autant deux façons de calculer. C'est ici que sert le paramètre cache de l'attention, laissé de côté à l'article 8. Quand il est présent, les clés et valeurs des nouvelles positions sont ajoutées à celles du passé, et le calcul continue tel quel :

if let Some(cache) = cache {
    k = cache.keys.extend(tape, k);
    v = cache.values.extend(tape, v);
}
// [B, H, N, D] × [B, H, D, T] → [B, H, N, T]
let scores = tape.matmul(q, tape.transpose(k, 2, 3));

La séquence est simplement plus longue du côté des clés (T positions connues) que du côté des requêtes (N nouvelles). Le softmax causal sait traiter ce cas, on l'a vu plus haut : les N requêtes sont les dernières des T positions.

L'entraînement et la génération passent donc par la même fonction. L'entraînement est le cas sans passé, où N = T. La génération est le cas où N vaut 1. Et l'amorce est un cas intermédiaire : elle traverse le modèle en une seule passe, quelle que soit sa longueur.

Le cache appartient au Generator de l'article 9, qui garde déjà la séquence. Son calcul des scores se réduit à ceci :

match &mut self.cache {
    Some(cache) => {
        let new = &self.tokens[self.pending.max(window)..];
        self.pending = self.tokens.len();
        last_logits(self.backend, self.model, new, Some(cache))
    }
    None => last_logits(self.backend, self.model, &self.tokens[window..], None),
}

pending est le rang du premier token qui n'a pas encore traversé le modèle. Au premier tour, toute l'amorce est en attente ; ensuite, un seul token. window est le début de ce que le modèle peut voir : zéro tant que la séquence tient dans le contexte.

Le cas sans cache montre une limite. Quand la séquence dépasse le contexte, le modèle ne garde que les derniers tokens. Toutes les positions se décalent alors d'un cran, et avec elles leurs embeddings de position. Le cache ne vaut plus rien : le générateur l'abandonne et recalcule tout, en appelant la même fonction sans cache.

Un test vérifie que découper une séquence en morceaux ne change rien : traitée d'un coup ou en trois fois à travers le cache, l'attention rend les mêmes nombres. Un autre vérifie que la génération avec et sans cache produit exactement les mêmes tokens. Sur 400 tokens, elle passe de 85 à 222 tokens par seconde.

Lire le code produit par le compilateur

À plusieurs reprises, la seule façon de comprendre un kernel lent a été de lire le PTX, l'assembleur que le compilateur produit pour la carte.

Les tableaux vont en mémoire. Un tableau local comme [f32; 16] est rangé dans une mémoire locale, lente, dès que son adresse sert. Dans une boucle, il coûte une écriture mémoire par tour. Seize variables séparées restent dans des registres. Voilà pourquoi les accumulateurs de la matmul s'appellent c00, c01... et ne forment pas un tableau.

Chaque accès est vérifié. Chaque a[i] sur une tranche devient un test et un saut vers un piège en cas de dépassement. C'est la sécurité de Rust, et elle a un prix. La matmul sur Tensor Cores a donc un chemin rapide, sans aucun test, pour les tuiles entièrement à l'intérieur des matrices.

Ce qui n'a pas marché

Les essais ratés font partie du résultat.

• Des tuiles plus épaisses pour la matmul f32 : deux fois plus lent.
• Des tuiles de 256 × 256 avec 1 024 threads par bloc pour les Tensor Cores : quatre fois plus lent, probablement parce que de si gros blocs laissent peu de blocs actifs en même temps sur la carte.
• Lire la tranche suivante à l'avance, pour masquer le délai de la mémoire : aucun gain.
• Le gradient des poids en un seul grand produit, plutôt qu'un produit par séquence suivi d'une somme : plus lent, 176 ms de matmul au lieu de 157.
• Supprimer les vérifications de bornes à la compilation : 5 à 7 % de gain, trop peu pour renoncer à la sécurité par défaut.
• Préparer les lots en parallèle du calcul : préparer un lot prend 14 microsecondes, soit 0,01 % d'un pas. Il n'y avait rien à gagner, et cela n'a pas été implémenté.
• Envoyer les données à la carte sans attendre : essayé, puis retiré. Le débit n'a pas bougé, pour du code unsafe en plus.

Aucun de ces résultats n'était évident avant la mesure. Plusieurs sont l'inverse de ce qu'on aurait parié.

Le bilan chiffré

Le plus gros modèle traite maintenant 52 000 tokens par seconde, contre 12 300 lors de son entraînement. Le fichier que l'article précédent aurait mis 19 heures à parcourir se parcourt en 4 heures et demie.

Et par rapport à PyTorch ?

Il fallait une référence extérieure. Le même modèle a été écrit avec PyTorch et mesuré sur la même carte.

À opérations égales, Smallem est 1,1 à 1,7 fois plus lent que PyTorch en f32, et deux fois en précision mixte.

PyTorch n'a pas été profilé, mais nos propres mesures désignent les causes probables. La principale est la multiplication de matrices, qui pèse encore 157 ms sur 224 : PyTorch utilise cuBLAS, une bibliothèque réglée carte par carte depuis des années. Viennent ensuite des passes mémoire que PyTorch évite : il ne recopie rien pour transposer, il fusionne l'ajout du biais, et en précision mixte il stocke les activations sur 16 bits, là où nous convertissons à chaque lecture.

Ce qu'il faut retenir

• On mesure avant d'optimiser, et on synchronise avant de mesurer.
• La mémoire globale est lente, la mémoire partagée rapide, les registres plus encore. Presque chaque gain consiste à faire descendre des données d'un cran.
• Les threads d'un warp doivent lire des cases voisines.
• Les Tensor Cores demandent des nombres de 16 bits ; le bf16 garde la plage du f32.
• Une optimisation n'est acquise qu'avec un test qui la compare à la version simple.

Le modèle est maintenant assez rapide pour servir de banc d'essai. La suite remplace ses pièces une à une par celles des modèles récents : Trois pièces modernes, et comment mesurer une idée.