Un tenseur sur CPU

Quatrième article de la série Smallem. Précédent : Des données et un premier modèle.

Un réseau de neurones est une longue suite d'opérations sur des tableaux de nombres. Avant d'en écrire un, il faut décider précisément ce que sont ces tableaux et ce que fait chaque opération. Cet article le fait sur CPU, le plus simplement possible.

Ce code ne servira jamais à entraîner un vrai modèle. Il est la définition de chaque opération. Quand on écrira les kernels GPU, on saura qu'ils sont justes parce qu'ils donneront les mêmes nombres que lui.

Ce qu'est un tenseur

Un tenseur est un tableau de nombres à plusieurs dimensions. Un vecteur a une dimension, une matrice en a deux, et un lot de phrases dont chaque mot est un vecteur en a trois.

La structure tient en deux champs :

pub struct CpuTensor<F> {
    data: Vec<F>,
    shape: Vec<usize>,
}

• data contient tous les nombres, à plat, dans un seul vecteur.
• shape, la forme, dit comment lire ce vecteur : combien de dimensions, et quelle taille pour chacune.

Les nombres seuls ne suffisent pas. Les six nombres [1, 2, 3, 4, 5, 6] peuvent être un vecteur de 6 éléments, une matrice de 2 lignes et 3 colonnes, ou une matrice de 3 lignes et 2 colonnes. C'est la forme qui tranche :

data  = [1, 2, 3, 4, 5, 6]

shape = [2, 3]          shape = [3, 2]
  1 2 3                   1 2
  4 5 6                   3 4
                          5 6

Le rangement suit une convention, dite ligne par ligne : on écrit la première ligne en entier, puis la deuxième, et ainsi de suite. Autrement dit, le dernier indice est celui qui varie le plus vite.

Le type F est générique : c'est f32 ou f64. On aura besoin des deux. Le f32 permet de comparer avec le GPU, qui calcule en simple précision. Le f64 servira à vérifier les gradients à l'article 6.

Les pas

Pour trouver l'élément (i, j) d'une matrice [2, 3] dans le vecteur à plat, on calcule i × 3 + j. Le 3 est le pas de la première dimension : de combien on avance dans le vecteur quand i augmente de 1.

La règle se généralise. Pour la forme [2, 3, 4], les pas sont [12, 4, 1] :

• avancer de 1 sur le dernier indice, c'est avancer de 1 élément ;
• avancer de 1 sur l'indice du milieu, c'est sauter une rangée de 4 ;
• avancer de 1 sur le premier indice, c'est sauter un bloc de 3 × 4 = 12.

pub fn strides(shape: &[usize]) -> Vec<usize> {
    let mut strides = vec![1; shape.len()];
    for dim in (1..shape.len()).rev() {
        strides[dim - 1] = strides[dim] * shape[dim];
    }
    strides
}

L'élément de coordonnées (i, j, k) est donc à l'indice i × 12 + j × 4 + k. Les pas vont résoudre élégamment les deux problèmes suivants.

Transposer, c'est échanger deux pas

Transposer une matrice, c'est échanger ses lignes et ses colonnes. L'élément (i, j) de la sortie est l'élément (j, i) de l'entrée.

Avec les pas, cela s'écrit sans effort. Pour remplir la sortie, on parcourt ses éléments dans l'ordre, et pour chacun on va chercher le bon élément dans la source. Il suffit de lire la source avec ses pas échangés :

pub fn transpose(&self, axis_a: usize, axis_b: usize) -> Self {
    let mut shape = self.shape.clone();
    shape.swap(axis_a, axis_b);
    let mut strides = strides(&self.shape);
    strides.swap(axis_a, axis_b);

    let data = (0..self.data.len())
        .map(|flat| self.data[source_index(flat, &shape, &strides)])
        .collect();
    Self { data, shape }
}

La fonction source_index fait le travail : elle transforme un indice à plat de la sortie en coordonnées, puis ces coordonnées en indice dans la source, à l'aide des pas qu'on lui donne.

pub fn source_index(mut flat: usize, shape: &[usize], strides: &[usize]) -> usize {
    let mut index = 0;
    for dim in (0..shape.len()).rev() {
        index += (flat % shape[dim]) * strides[dim];
        flat /= shape[dim];
    }
    index
}

Elle marche comme la décomposition d'un nombre en chiffres. flat % shape[dim] extrait la coordonnée de la dernière dimension, flat /= shape[dim] passe à la précédente.

À l'inverse, reshape ne déplace rien : il garde les mêmes données et change seulement la forme. C'est une différence qu'il faut garder en tête, car elle est source de bugs : transformer une matrice [2, 3] en [3, 2] avec reshape ne la transpose pas.

Le broadcasting

On veut souvent combiner des tenseurs de formes différentes. Ajouter un vecteur de biais à chaque ligne d'une matrice, par exemple, ou multiplier tout un tenseur par un seul nombre.

Le broadcasting est la règle qui dit quand c'est permis et ce que cela signifie. Elle vient de NumPy :

1. on aligne les deux formes par la droite
2. deux dimensions sont compatibles si elles sont égales, ou si l'une vaut 1
3. une dimension de taille 1, ou absente, s'étire pour atteindre l'autre

  [4, 3]        [4, 1]        [2, 3]
+    [3]      + [1, 3]      +     []
= [4, 3]      = [4, 3]      = [2, 3]

Dans le premier cas, le vecteur de 3 éléments est ajouté à chacune des 4 lignes. Dans le dernier, la forme [] est celle d'un scalaire, un tenseur à un seul élément : il s'étire à n'importe quelle forme. C'est la façon de multiplier un tenseur par une constante.

shape[ndim - 1 - from_right] = match (dim_a, dim_b) {
    (x, y) if x == y => x,
    (1, y) => y,
    (x, 1) => x,
    _ => panic!("formes incompatibles pour le broadcasting : {a:?} et {b:?}"),
};

Reste à l'exécuter. La solution naïve serait de recopier le petit tenseur autant de fois que nécessaire, puis d'additionner. Les pas permettent d'éviter cette copie.

L'idée : une dimension étirée reçoit un pas de zéro. Avec un pas nul, on a beau avancer le long de cette dimension, on retombe toujours sur le même élément. C'est exactement ce que veut dire « étirer ».

pub fn broadcast_strides(shape: &[usize], target: &[usize]) -> Vec<usize> {
    let own = strides(shape);
    let offset = target.len() - shape.len();
    let mut result = vec![0; target.len()];
    for dim in 0..shape.len() {
        if shape[dim] != 1 {
            result[offset + dim] = own[dim];
        }
    }
    result
}

Prenons le vecteur de biais [3] ajouté à une matrice [4, 3]. Ses pas, vus à travers la forme [4, 3], sont [0, 1]. Quand on parcourt la matrice, changer de ligne ne déplace pas la lecture du biais (pas 0), changer de colonne l'avance d'un élément (pas 1). Chaque ligne relit donc les trois mêmes valeurs.

Toutes les opérations élément par élément passent alors par une seule fonction :

fn zip_with(&self, other: &Self, f: impl Fn(F, F) -> F) -> Self {
    let shape = broadcast_shapes(&self.shape, &other.shape);
    let strides_a = broadcast_strides(&self.shape, &shape);
    let strides_b = broadcast_strides(&other.shape, &shape);
    let len: usize = shape.iter().product();
    let data = (0..len)
        .map(|flat| {
            let a = self.data[source_index(flat, &shape, &strides_a)];
            let b = other.data[source_index(flat, &shape, &strides_b)];
            f(a, b)
        })
        .collect();
    Self { data, shape }
}

L'addition est zip_with(other, |a, b| a + b), la multiplication zip_with(other, |a, b| a * b). Le broadcasting est écrit une fois pour toutes. On retrouvera la même astuce sur GPU, presque mot pour mot.

Les réductions

Une réduction fait disparaître une dimension en combinant ses éléments. La somme le long d'un axe en est le prototype : sommer une matrice [4, 3] le long de l'axe 0 donne un vecteur de 3 éléments, les totaux de chaque colonne.

L'option keepdim garde l'axe réduit, avec la taille 1. Le résultat a alors la forme [1, 3] au lieu de [3]. Ce détail a une utilité précise : le résultat peut se recombiner avec l'entrée par broadcasting. C'est ce qui permet d'écrire « soustraire à chaque ligne sa moyenne » en une expression.

La multiplication de matrices

C'est l'opération centrale. Dans un Transformer, elle représente la quasi-totalité du temps de calcul, et l'article 11 lui sera presque entièrement consacré.

Le produit d'une matrice A de forme [rows, inner] par une matrice B de forme [inner, cols] est une matrice [rows, cols]. Son élément (i, j) est le produit scalaire de la ligne i de A par la colonne j de B :

for i in 0..rows {
    for j in 0..cols {
        let mut sum = F::ZERO;
        for k in 0..inner {
            sum += a[i * inner + k] * b[k * cols + j];
        }
        out[i * cols + j] = sum;
    }
}

Trois boucles imbriquées. C'est la version la plus lente possible, et c'est bien ainsi : elle est impossible à rater.

Un Transformer ne multiplie jamais deux matrices isolées. Il multiplie des lots de matrices. La convention est que le produit porte sur les deux dernières dimensions, et que toutes celles qui précèdent sont des dimensions de lot, soumises au broadcasting :

[B, T, C] × [C, D]       → [B, T, D]      une couche linéaire
[B, H, T, D] × [B, H, D, T] → [B, H, T, T]   les scores d'attention

Dans le premier cas, la matrice de poids [C, D] n'a pas de dimension de lot : elle est réutilisée pour chacune des B séquences. C'est le broadcasting, appliqué cette fois à des matrices entières plutôt qu'à des nombres.

Le softmax, et un piège numérique

Un réseau produit des scores quelconques, positifs ou négatifs, qu'on appelle des logits. Un modèle de langage doit rendre des probabilités. Le softmax fait la conversion :

softmax(x)ᵢ = exp(xᵢ) / Σⱼ exp(xⱼ)

L'exponentielle rend tout positif, et la division fait que la somme vaut 1.

Écrite telle quelle, cette formule ne marche pas. Avec un logit de 100, exp(100) vaut environ 10⁴³ : un f32, qui plafonne vers 10³⁸, répond « infini », et le résultat devient infini / infini, c'est-à-dire rien du tout.

La parade repose sur une propriété du softmax : ajouter la même constante à tous les logits ne change pas le résultat. On soustrait donc le plus grand logit à tous les autres avant de prendre l'exponentielle. Le plus grand devient 0, son exponentielle vaut 1, et les autres sont plus petites : rien ne peut plus déborder.

pub fn softmax(&self) -> Self {
    self.map_rows(|row, max| {
        let mut total = F::ZERO;
        for x in row.iter_mut() {
            *x = (*x - max).exp();
            total += *x;
        }
        for x in row.iter_mut() {
            *x = *x / total;
        }
    })
}

On reverra ce genre de précaution plusieurs fois. En calcul numérique, une formule juste sur le papier n'est pas toujours un programme juste.

Le contrat : primitives et opérations

Une fois tout cela écrit, il faut en faire un contrat : la liste de ce que le reste du projet a le droit de demander à des tenseurs. Ce contrat est en deux étages.

Les primitives

Le premier étage est un trait, qui réunit les briques élémentaires :

pub trait Primitives {
    type Scalar: Float;
    type Tensor: Clone;

    fn from_vec(&self, data: Vec<Self::Scalar>, shape: &[usize]) -> Self::Tensor;
    fn to_vec(&self, tensor: &Self::Tensor) -> Vec<Self::Scalar>;
    fn shape<'a>(&self, tensor: &'a Self::Tensor) -> &'a [usize];

    fn add(&self, a: &Self::Tensor, b: &Self::Tensor) -> Self::Tensor;
    fn mul(&self, a: &Self::Tensor, b: &Self::Tensor) -> Self::Tensor;
    fn sum(&self, a: &Self::Tensor, axis: usize, keepdim: bool) -> Self::Tensor;
    fn reshape(&self, a: &Self::Tensor, shape: &[usize]) -> Self::Tensor;
    fn transpose(&self, a: &Self::Tensor, axis_a: usize, axis_b: usize) -> Self::Tensor;
    fn matmul(&self, a: &Self::Tensor, b: &Self::Tensor) -> Self::Tensor;
    fn softmax(&self, a: &Self::Tensor) -> Self::Tensor;
    // … une trentaine de primitives en tout
}

Ce trait décrit « un endroit où vivent des tenseurs, avec ce qu'on peut leur faire ». Le backend CPU, Cpu<F>, en est une implémentation. Le backend GPU en sera une autre.

La liste est courte, et elle a vocation à ne plus bouger. Un Transformer entier, entraînement compris, se ramène à une trentaine de primitives.

Les opérations

Tout ne mérite pas d'être une primitive. Une normalisation, par exemple, s'écrit avec des moyennes, des soustractions et une racine : elle se construit au-dessus des primitives. On appelle ces constructions des opérations.

Chaque opération est une petite structure, qui porte ses réglages, et qui vit dans son propre fichier :

pub struct Normalize {
    pub epsilon: f64,
}

Un backend sait exécuter une opération s'il implémente un trait pour elle :

pub trait Run<O: Operation>: Primitives {
    fn run(&self, op: &O, inputs: &[&Self::Tensor]) -> Self::Tensor;
    fn run_backward(&self, op: &O, inputs: &[&Self::Tensor],
                    output: &Self::Tensor, grad: &Self::Tensor) -> Vec<Option<Self::Tensor>>;
}

La seconde méthode sert à l'apprentissage ; on y revient à l'article 6.

Quand une opération s'écrit avec les primitives, cette écriture est sa définition. Elle s'appelle sa référence, et elle est générique : elle fonctionne sur n'importe quel backend, puisqu'elle ne lui demande que des primitives.

impl Reference for Normalize {
    fn reference<B: Primitives>(&self, backend: &B, inputs: &[&B::Tensor]) -> B::Tensor {
        let [x] = operands(Self::NAME, inputs);
        let (centered, inverse_std) = self.centered_and_inverse_std(backend, x);
        backend.mul(&centered, &inverse_std)
    }
}

Le backend CPU exécute cette référence telle quelle, sans une ligne de plus. Le backend GPU pourra faire mieux, avec un kernel dédié, et on saura que ce kernel est juste parce qu'il donne les mêmes nombres que la référence.

Avec ce découpage, ajouter une opération au projet c'est ajouter un fichier : la structure, sa définition et, on le verra, sa règle d'apprentissage y sont réunies. Rien d'autre dans le projet n'a à la connaître.

Le backend

Un Backend, enfin, est simplement un type qui fournit les primitives et sait exécuter toutes les opérations dont un modèle a besoin :

pub trait Backend:
    Primitives + Run<Gelu> + Run<Normalize> + Run<CausalSoftmax> + Run<CrossEntropy> // + …
{
}

Tout le reste du projet, l'autograd, les couches, le modèle, l'optimiseur, sera écrit en fonction de ce trait, sans savoir où tournent les calculs. Le même modèle pourra donc tourner sur CPU en f64 pour être vérifié, et sur GPU en f32 pour être entraîné, sans qu'une ligne ne change.

Deux paquets

Ce contrat trace aussi la frontière du projet, qui est fait de deux paquets Rust.

Le second dépend du premier, jamais l'inverse. Le premier ne sait pas qu'une carte graphique existe : il n'a besoin ni du toolkit, ni de la nightly. Ses tests, qui vérifient toutes les mathématiques du projet, tournent en deux secondes sur n'importe quelle machine.

Rien de particulier n'a été fait pour rendre cette séparation possible. Elle découle du contrat : puisque le modèle ne connaît que Backend, un backend peut être fourni de l'extérieur.

Ce qu'il faut retenir

• Un tenseur, c'est des données à plat et une forme qui dit comment les lire.
• Les pas font le lien entre coordonnées et indice à plat. Échanger deux pas transpose, un pas nul étire.
• Le produit matriciel porte sur les deux dernières dimensions, les autres sont des lots.
• Le softmax soustrait le maximum avant l'exponentielle, sans quoi il déborde.
• Le contrat a deux étages : une trentaine de primitives, et au-dessus des opérations, chacune dans son fichier avec sa définition de référence.
• Le trait Backend réunit les deux : c'est le contrat entre le modèle et l'endroit où il calcule.

Chaque primitive et chaque opération est testée sur de petits exemples calculés à la main. La référence est en place. On peut maintenant écrire la même chose sur la carte graphique : Les mêmes opérations sur GPU.