smallbard

Articles · Smallem : un LLM en Rust et CUDA, à partir de zéro

Les mêmes opérations sur GPU

Article 5 sur 18 ·

Cinquième article de la série Smallem. Précédent : Un tenseur sur CPU.

L'article précédent a fixé une trentaine de primitives et quelques opérations sur des tenseurs. Celui-ci les réécrit pour la carte graphique, dans le second paquet du projet, smallem-cuda. Le but est de pouvoir enchaîner des milliers d'opérations sans que les données quittent jamais la mémoire de la carte.

Pour le moment, on se concentre sur la justesse des kernels en les écrivant le plus simplement possible. On apportera des optimisations dans l'article 11, et on verra qu'elles demandent de revenir sur plusieurs des choix faits ici.

Le modèle d'exécution

Sur un GPU, un kernel est une fonction que des milliers de threads exécutent en même temps, chacun une seule fois, chacun avec son propre numéro.

Presque tous les kernels de cet article suivent alors le même schéma :

un thread par élément de la sortie et chaque thread calcule son élément et l'écrit.

Comme deux threads n'écrivent jamais au même endroit, ils n'ont pas besoin de se coordonner : ni verrou, ni attente. Toute la difficulté consiste à formuler chaque opération sous cette forme.

Les threads sont regroupés en blocs de 256. Pour une sortie de len éléments, on lance assez de blocs pour avoir au moins un thread par élément :

fn launch_config(len: usize) -> LaunchConfig1D {
    let blocks = (len as u32).div_ceil(BLOCK_SIZE).max(1);
    LaunchConfig1D::new(blocks, BLOCK_SIZE, 0)
}

Le dernier bloc contient en général des threads en trop, qui ne font rien.

Le tenseur GPU

pub struct CudaTensor {
    storage: Rc<Storage>,
    shape: Vec<usize>,
}

Les données sont dans un tampon en mémoire GPU, derrière storage. La forme, elle, reste sur l'hôte : ce sont quelques entiers dont le programme a besoin en permanence pour vérifier les opérations et dimensionner les lancements.

Le tampon est partagé par comptage de références (Rc). Ainsi reshape ne coûte rien : le nouveau tenseur pointe sur les mêmes données avec une autre forme, sans rien copier. C'est sans danger pour une raison simple : aucune opération ne modifie un tenseur existant, chacune en crée un nouveau.

Un seul kernel pour toutes les opérations à deux opérandes

L'addition, la soustraction, la multiplication et la division ne diffèrent que par une ligne. Plutôt que quatre kernels, on en écrit un qui reçoit un code d'opération :

#[derive(Clone, Copy, Debug)]
#[repr(u32)]
pub enum BinaryOp {
    Add,
    Sub,
    Mul,
    Div,
    MaskedFill,
    GeluBackward,
}

Le #[repr(u32)] fixe la représentation de l'énumération en mémoire, ce qui permet de la passer telle quelle à la carte. Un kernel peut recevoir des structures et des énumérations par valeur, et faire un match dessus comme n'importe quel code Rust.

Le broadcasting sur GPU

Sur CPU, le broadcasting reposait sur les pas : un pas nul le long d'une dimension étirée. On reprend la même idée. Mais un kernel ne peut pas recevoir un Vec, car il ne peut pas allouer de mémoire. On range donc la forme et les pas dans des tableaux de taille fixe :

pub const MAX_RANK: usize = 4;

#[derive(Clone, Copy, Debug)]
pub struct Layout {
    pub shape: [u32; MAX_RANK],
    pub strides: [u32; MAX_RANK],
}

Un tenseur a donc au plus quatre dimensions sur le GPU. C'est ce dont on aura besoin pour un Transformer : lot, têtes, positions, composantes.

Pour les tenseurs qui en ont moins, on complètera à gauche avec des dimensions de taille 1 et de pas 0. Une dimension de taille 1 n'a qu'une coordonnée possible, 0, et 0 × 0 = 0 : elle ne change rien au calcul. Le kernel peut alors toujours boucler sur quatre dimensions, sans cas particulier.

La conversion d'indice est la même que sur CPU, réécrite pour des tableaux fixes :

impl Layout {
    pub fn source_index(&self, mut flat: u32) -> usize {
        let mut index = 0;
        let mut dim = MAX_RANK;
        while dim > 0 {
            dim -= 1;
            index += (flat % self.shape[dim]) * self.strides[dim];
            flat /= self.shape[dim];
        }
        index as usize
    }
}

Cette fonction est une fonction Rust ordinaire. Elle est compilée à la fois pour le processeur et pour le GPU : un kernel peut appeler du code Rust normal, tant qu'il n'alloue pas de mémoire.

Le kernel complet tient alors en quelques lignes :

#[kernel]
#[launch_bounds(256)]
#[launch_contract(domain = 1, block = (256, 1, 1))]
pub fn binary(
    op: BinaryOp,
    param: f32,
    a: &[f32],
    layout_a: Layout,
    b: &[f32],
    layout_b: Layout,
    mut out: DisjointSlice<f32>,
) {
    let idx = thread::index_1d();
    let flat = idx.get() as u32;
    if let Some(out_elem) = out.get_mut(idx) {
        let x = a[layout_a.source_index(flat)];
        let y = b[layout_b.source_index(flat)];
        *out_elem = match op {
            BinaryOp::Add => x + y,
            BinaryOp::Sub => x - y,
            BinaryOp::Mul => x * y,
            BinaryOp::Div => x / y,
            // …
        };
    }
}

Chaque thread connaît sa position flat dans la sortie. Il la traduit en une position dans a et une position dans b, lit les deux valeurs, calcule et écrit.

Ce même mécanisme sert aussi à la transposition, où l'on échange deux pas, et à expand, qui étire un tenseur : un seul kernel, permute, fait les deux.ch

Les réductions

Pour sommer le long d'un axe, on garde le schéma : un thread par élément de sortie. Chaque thread additionne les n éléments de l'axe réduit qui le concernent.

pub fn sum_axis(a: &[f32], n: u32, inner: u32, mut out: DisjointSlice<f32>) {
    let idx = thread::index_1d();
    let flat = idx.get() as u32;
    if let Some(out_elem) = out.get_mut(idx) {
        let o = flat / inner;
        let i = flat % inner;
        let mut total = 0.0;
        for k in 0..n {
            total += a[((o * n + k) * inner + i) as usize];
        }
        *out_elem = total;
    }
}

La somme de tous les éléments est plus délicate. Elle ne produit qu'un seul nombre : le schéma « un thread par élément de sortie » donnerait un unique thread qui parcourt tout, sans aucun parallélisme.

On procède donc par étages. À chaque passe, le thread t additionne un paquet de 256 éléments et écrit un seul nombre. Un million d'éléments devient 3 907 sommes partielles, puis 16, puis une seule : trois passes.

1 000 000 éléments  ──►  3 907  ──►  16  ──►  1

Cette somme par étages est plus précise que la boucle du CPU. En f32, ajouter un petit nombre à un très grand total perd des chiffres. En additionnant par paquets, on additionne toujours des nombres de taille comparable. Le test qui comparait le GPU au CPU sur 100 000 éléments échouait pour cette raison : c'était le CPU qui avait tort. Il compare maintenant le GPU à une référence en f64.

Le softmax en deux temps

Le softmax illustre bien la contrainte du modèle : Pour calculer un seul élément du résultat, il faut connaître le maximum et la somme de toute sa ligne. Or un thread ne peut écrire que son propre élément : il ne peut pas normaliser la ligne à lui seul.

On le fait donc en deux lancements.

  1. row_stats calcule, pour chaque ligne, les deux réductions nécessaires : son maximum, et ln(Σ exp(x - max)).
  2. softmax_rows lance un thread par élément, qui lit les deux réductions de sa ligne :
let row = flat / width as usize;
let log_probability = (a[flat] - row_max[row]) - row_log_sum[row];
*out_elem = match mode {
    SoftmaxMode::Log => log_probability,
    SoftmaxMode::Probability => log_probability.exp(),
};

La première version additionnait les deux réductions pour ne passer qu'un nombre par ligne, max + ln(Σ …). C'est mathématiquement identique, et pourtant le GPU donnait des résultats dix fois moins précis que le CPU sur une attention saturée. La raison est celle de la somme globale : le second terme, petit, était noyé dans le maximum, qui peut être grand. Garder les deux nombres séparés, et soustraire d'abord le maximum, a ramené l'erreur de 5,6e-3 à 1,1e-3.

La multiplication de matrices, version naïve

La première matmul GPU suivait le schéma à la lettre : un thread par élément (lot, i, j) de la sortie, qui parcourt k et fait le produit scalaire de sa ligne par sa colonne. En substance :

let mut total = 0.0;
for k in 0..inner {
    total += a[start_a + row * inner + k] * b[start_b + k * cols + col];
}
*out_elem = total;

C'est correct, et c'est lent. Chaque thread relit en mémoire globale toute une ligne de a et toute une colonne de b, que ses voisins relisent aussi. Ce kernel a été remplacé à l'article 11 et n'existe plus dans le code.

Les lancements sont asynchrones

Le piège avec le GPU est que lorsque l'hôte lance un kernel, l'appel revient immédiatement. Le kernel n'est pas terminé, il n'a peut-être même pas commencé : il a été ajouté à la file du stream.

C'est une bonne chose car l'hôte peut mettre en file des centaines de kernels pendant que la carte travaille, sans attendre la fin de chacun. Mais cela a deux conséquences :

Chronométrer. Une horloge autour d'un appel mesure la mise en file et non pas le calcul. Pour mesurer un temps, il faut d'abord attendre que la file soit vide :

pub fn synchronize(&self) {
    self.stream.synchronize().expect(CUDA_ERROR);
}

Échanger des données. Rapatrier un tenseur sur l'hôte oblige à attendre tous les kernels en file, puis à copier des données. Pendant un entraînement, on ne lit donc qu'un seul nombre de temps en temps : la perte.

L'envoi dans l'autre sens attend lui aussi. La fonction d'envoi de cuda-core synchronise le stream avant de rendre la main, pour que l'hôte puisse disposer de ses données aussitôt. Or un pas d'entraînement envoie les indices des tokens, les cibles et quelques constantes : il attend donc la carte une dizaine de fois.

De plus un tenseur peut être libéré par l'hôte alors qu'un kernel en file ne l'a pas encore lu. Cela reste cependant correct, parce que tout passe par le même stream : les kernels y sont exécutés dans l'ordre où ils ont été mis en file.

Tester : le CPU comme oracle

C'est ici que le tenseur CPU est utile. Pour chaque opération, un test tire des entrées au hasard, exécute l'opération sur les deux backends et compare.

macro_rules! check {
    ($shapes:expr, |$b:ident, $t:ident| $op:expr) => {{
        let inputs = random_inputs($shapes, 0);
        let cpu = eval(&Cpu::<f32>::new(), &inputs, |$b, $t| $op);
        let gpu = eval(&Cuda::new(0).unwrap(), &inputs, |$b, $t| $op);
        assert_close(&cpu, &gpu);
    }};
}

#[test]
fn matmul_batch_times_matrix() {
    // [B, T, C] × [C, D] : une couche linéaire.
    check!(&[&[4, 9, 33], &[33, 21]], |b, t| b.matmul(&t[0], &t[1]));
}

Le corps du test est écrit une fois et compilé deux fois, une par backend. C'est possible parce que les deux remplissent le même trait Backend.

Pour rendre ces tests pertinents et utiles, on peut prendre trois précautions :

  • Des formes biscornues. Les tests utilisent des tailles comme 7 × 300 ou 33 × 21, qui ne sont multiples ni de 256 ni d'aucune taille de bloc. Les bugs se cachent dans les blocs incomplets.
  • Une tolérance. Le GPU ne fait pas les opérations dans le même ordre que le CPU, et il fusionne une multiplication et une addition en une seule instruction. Les résultats ne sont pas identiques au bit près : on tolère un écart relatif de 1e-5.
  • Les cas de broadcasting, dans les deux sens.

Les pièges de l'outillage

Deux difficultés sont propres à cuda-oxide, et valent d'être connues. Elles expliquent la façon dont le paquet smallem-cuda est organisé.

Tout ce qui charge des kernels se compile avec cargo oxide, et non cargo. Avec cargo seul, l'édition de liens échoue sur un symbole au nom étrange, cuda_oxide_artifact_anchor_…. C'est une raison de plus de tenir le code GPU dans son propre paquet : tout ce qui n'en fait pas partie se compile et se teste normalement.

Une bibliothèque de kernels ne doit être compilée qu'une fois. En compilant, cuda-oxide écrit le code des kernels dans un fichier à la racine du paquet. Or cargo compile deux fois une bibliothèque qui a des tests unitaires : une fois pour elle-même, une fois avec ses tests. Les deux compilations tournent en parallèle et écrivent le même fichier ; l'une écrase le travail de l'autre, et des tests échouent au chargement avec un named symbol not found sans rapport apparent avec la cause.

Le paquet GPU n'a donc aucun test unitaire. Tous ses tests sont des tests d'intégration, rangés dans tests/ : ce sont des programmes à part, qui utilisent la bibliothèque sans la recompiler. Cela ne coûte rien, puisque ces tests comparent le GPU au CPU à travers le contrat public, et n'ont pas besoin de voir l'intérieur du backend.

cargo test          # dans smallem : tout sauf le GPU
cargo oxide test    # dans smallem-cuda : chaque kernel contre le CPU

Ce qu'il faut retenir

  • Le schéma de base : un thread par élément de sortie, qui calcule et écrit son élément.
  • Une forme et des pas, dans des tableaux de taille fixe, suffisent à porter le broadcasting sur le GPU.
  • Ce qui ne rentre pas dans le schéma se fait en plusieurs lancements : somme par étages, softmax en deux temps.
  • Un lancement n'est pas une exécution. On synchronise avant de chronométrer, et on lit le moins possible.
  • Chaque kernel est comparé au CPU sur des entrées aléatoires.

À ce stade, 45 tests comparent les deux backends et passent. On sait calculer sur la carte. Il manque de quoi apprendre : les gradients. C'est l'article suivant : La rétropropagation.