Le projet et ses règles

Premier article de la série Smallem.

L'objectif de ce projet est simple: comprendre le fonctionnement d'un LLM et savoir en construire. Pour cela la meilleure façon est d'en écrire un, sans utiliser de framework.

Donc pas de PyTorch, pas de bibliothèque d'algèbre linéaire, pas de bibliothèque de kernels GPU. On écrit le tokeniseur, les tenseurs, la multiplication de matrices, la rétropropagation, l'optimiseur, le Transformer et la boucle d'entraînement. L'objectif étant éducatif, le résultat n'est pas destiné à la production, pour laquelle on utilisera plutôt des frameworks éprouvés.

Le tout est écris en Rust, avec le calcul sur une carte graphique via CUDA Rust.

Ce qu'est un modèle de langage

Un modèle de langage est une fonction qui reçoit un début de texte et renvoie une probabilité pour chaque suite possible.

Le texte est d'abord découpé en tokens, des petits morceaux numérotés. Si le vocabulaire compte 257 tokens, le modèle renvoie 257 nombres positifs dont la somme vaut 1 : la probabilité que le token suivant soit le numéro 0, le numéro 1, et ainsi de suite.

Deux choses en découlent:

• Générer du texte, c'est tirer un token selon ces probabilités, l'ajouter au texte, et recommencer.
• Entraîner le modèle, c'est lui montrer du texte réel et ajuster ses réglages pour qu'il donne une probabilité plus haute au token qui venait vraiment ensuite.

Ce qui change d'un modèle à l'autre, c'est la fonction : une simple table de comptage dans l'article 3, un réseau de 30 millions de paramètres dans l'article 10.

Les règles du jeu

Aucun framework d'apprentissage. Ni candle, ni burn, ni tch. Aucune bibliothèque de kernels toute faite non plus : pas de cuBLAS, pas de cuDNN. Chaque kernel est écrit à la main.

Très peu de dépendances. Les crates de CUDA Rust pour parler au GPU, rand pour le hasard, parquet pour lire le corpus, serde_json pour l'en-tête des fichiers de poids. Rien d'autre. Lire un fichier Parquet n'apprend rien sur les modèles de langage, donc on délègue ; multiplier deux matrices est le cœur du sujet, donc on l'écrit.

Un vrai corpus. Le texte d'entraînement vient de FineWeb-Edu, des pages web éducatives en anglais. Un seul de ses fichiers contient 3,3 Go de texte.

Des octets d'abord. Jusqu'à l'article 9, un token est un octet du texte. C'est le tokeniseur le plus simple qui fonctionne sur n'importe quel texte, et il laisse toute l'attention au modèle. Les sous-mots n'arrivent qu'à l'article 10.

Pourquoi Rust, et pourquoi CUDA Rust

Rust n'est pas le choix évident pour ce genre de projet : presque tout l'écosystème est en Python. C'est justement l'intérêt. En Python, on écrit trois lignes et une bibliothèque fait le reste ; ici il n'y a pas de reste.

Pour le GPU, le projet utilise CUDA Rust, un projet de recherche de NVIDIA. Il permet d'écrire le code qui tourne sur la carte, les kernels, en Rust, dans le même fichier que le reste du programme. CUDA Rust propose deux façons de le faire :

Le projet utilise cuda-oxide, le plus bas niveau des deux. On y gère soi-même les threads du GPU, ce qui est le but quand on veut comprendre.

Deux implémentations de tout

Sur toute la série, on garde ce choix de conception : chaque opération est écrite deux fois, une fois sur CPU, une fois sur GPU.

La version CPU est écrite le plus simplement possible, sans chercher la vitesse. Elle ne sert jamais à entraîner un vrai modèle. Elle sert à deux choses que le GPU fait mal :

• être l'oracle. Chaque kernel GPU est testé en comparant sa sortie à celle de l'opération CPU équivalente, sur des entrées tirées au hasard. Un kernel faux est repéré plus facilement ;
• vérifier les gradients. La méthode de vérification, qu'on verra à l'article 6, exige des nombres en double précision, alors que le GPU travaille en simple précision.

Cette redondance coûte du code, mais c'est elle qui rend le projet faisable. Un kernel GPU qui calcule faux ne plante pas : il renvoie des nombres. Sans référence à laquelle les comparer, on ne s'en aperçoit que des semaines plus tard, quand le modèle refuse d'apprendre et qu'on ne sait plus où chercher.

La route

Le projet avance par couches, chacune posée sur la précédente.

texte ──► tokens ──► tenseurs ──► kernels GPU ──► autograd ──► couches ──► GPT
  3         3           4             5              6            7         8

Les chiffres sont les numéros d'articles. Puis viennent l'entraînement réel (9), les sous-mots et le passage à l'échelle (10), et la performance (11).

L'ordre n'est pas arbitraire. Le modèle le plus simple, un bigramme, arrive avant le premier tenseur : il met en place les données, la perte et la génération sans aucune mathématique. Le tenseur CPU arrive avant le GPU, pour avoir une référence. L'autograd arrive avant le premier réseau. Et le Transformer n'arrive qu'une fois que tout ce dont il a besoin existe et a été testé : le jour où on l'assemble, il n'y a aucun nouveau kernel à écrire.

La machine

Tout a été développé et mesuré sur une seule machine :

• une carte NVIDIA RTX 3090, avec 24 Go de mémoire ;
• Ubuntu 24.04 ;
• le CUDA Toolkit 13.0.

Les temps donnés dans la série sont ceux de cette carte. Sur une autre, les nombres seront différents.

Une remarque sur les nombres

Chaque affirmation chiffrée de cette série a été mesurée : une perte, un débit, un temps par pas. Quand une idée n'a pas marché, c'est dit aussi, avec la mesure qui l'a montré. L'article 11 contient plusieurs optimisations essayées puis abandonnées, et elles en apprennent autant que celles qui ont réussi.

Dans l'article suivant, on installe l'environnement et on fait tourner un premier kernel : Un premier kernel sur le GPU.