Articles · Smallem : un LLM en Rust et CUDA, à partir de zéro
Bilan
Dernier article de la série Smallem. Précédent : L'attention fusionnée.
Le projet partait d'une question : que se passe-t-il entre le texte qu'on tape et le texte qui revient ? Dix-sept articles plus tard, il existe un programme qui y répond en entier. Cet article fait le point sur ce qui a été construit, ce qui manque, et ce qu'on en retire.
Ce qui existe
Un dépôt d'environ 25 000 lignes de Rust, tests compris, avec six dépendances, en deux paquets : l'un sans GPU, en Rust stable, l'autre pour la carte. Il contient :
- un tokeniseur BPE, de son apprentissage à son format de fichier
- deux implémentations de tenseurs, sur CPU et sur GPU, derrière un même contrat : une trentaine de primitives, et des opérations qui tiennent chacune dans un fichier
- une trentaine de kernels écrits à la main, dont une multiplication de matrices et une attention fusionnée sur Tensor Cores
- un autograd à tape
- un GPT complet, qui charge les poids de GPT-2 et retrouve ses sorties
- les pièces des modèles récents, chacune en option : RMSNorm, SwiGLU, RoPE, attention latente, mélange d'experts, prédiction multi-token, couches traversées deux fois, clés partagées entre couches
- deux optimiseurs, AdamW et Muon, ce dernier avec trois variantes
- une boucle d'entraînement avec planification du taux, écrêtage, accumulation et reprise
- un générateur de texte avec cache, qui affiche le texte à mesure qu'il s'écrit.
Et 490 tests, qui sont la raison pour laquelle tout cela tient debout. Les trois quarts tournent sans carte graphique, en deux secondes.
La progression, en deux tableaux
La qualité, d'abord. La perte est en nats par octet de texte : c'est la seule unité qui se compare d'un modèle à l'autre.
| Modèle | Paramètres | Perte | Hésite entre |
|---|---|---|---|
| Hasard uniforme | 0 | 5,55 | 257 octets |
| Bigramme | 66 049 | 2,53 | 12,6 octets |
| MLP à contexte fixe | 402 721 | 1,90 | 6,7 octets |
| GPT, octets | 834 432 | 1,36 | 3,9 octets |
| GPT, octets | 10 844 544 | 1,20 | 3,3 octets |
| GPT, sous-mots | 29 676 544 | 0,94 | 2,6 octets |
La dernière colonne est la perplexité : le nombre de choix équiprobables entre lesquels le modèle hésite, en moyenne, pour chaque octet.
Ce tableau s'arrête à l'article 10. Les articles 12 à 15 ont ensuite fait évoluer l'architecture, sur un budget plus court et fixe (49 millions de tokens, 16 minutes), pour pouvoir comparer :
| Variante | Perte | Débit, tokens/s |
|---|---|---|
| GPT-2, le modèle de l'article 10 | 1,000 | 50 700 |
| + RMSNorm, SwiGLU, RoPE | 0,959 | 52 300 |
| + attention latente (cache 6,4 fois plus petit) | 0,984 | 46 600 |
| + mélange d'experts équilibré | 0,940 à 0,944 | 33 000 |
| + prédiction multi-token | 0,955 | 42 200 |
| + Muon à la place d'AdamW, article 15 | 0,925 | 22 500 |
| + Muon, matrices empilées, quatre itérations | 0,926 | 43 900 |
| + deux traversées des couches | 0,952 | 27 600 |
| + clés partagées par 4 couches sur 8 (cache 2 fois plus petit) | 0,964 | 55 000 |
Chaque ligne à partir de la troisième change une seule chose par rapport à la deuxième. Le bruit entre deux graines vaut 0,002.
Ce tableau compare à nombre de pas égal, et les débits disent à quel point c'est trompeur. L'article 16 a donc tout repris à durée égale, contre le modèle de la deuxième ligne entraîné aussi longtemps :
| Variante | Sa perte | Sa durée | La référence, même durée |
|---|---|---|---|
| Mélange d'experts | 0,940 | 25 min | 0,925 |
| Prédiction multi-token | 0,955 | 19 min | 0,942 |
| Deux traversées | 0,952 | 30 min | 0,913 |
| Clés partagées, 4 sur 8 | 0,964 | 15 min | 0,963 |
| Muon de l'article 15 | 0,925 | 36 min | 0,900 |
| Muon empilé, quatre itérations | 0,926 | 19 min | 0,946 |
| Le même, sur 36 minutes | 0,882 | 36 min | 0,900 |
Une seule idée bat la référence à la minute : Muon, une fois son orthogonalisation réécrite. La dernière ligne est le meilleur modèle de la série.
La vitesse, ensuite, pour le plus gros modèle :
| Moment | Temps par pas | Tokens par seconde |
|---|---|---|
| Avant la seconde campagne d'optimisation | 652 ms | 12 300 |
Après, en f32 | 224 ms | 36 500 |
| Après, sur Tensor Cores | 157 ms | 52 000 |
| Avec l'attention fusionnée | 127 ms | 64 400 |
Ce que le projet ne fait pas
Il serait malhonnête de présenter Smallem comme un modèle abouti.
Le modèle est petit et peu entraîné. Trente millions de paramètres et deux heures d'entraînement, là où les modèles qu'on utilise au quotidien se comptent en dizaines de milliards de paramètres et en mois de calcul sur des milliers de cartes. Son texte est grammatical et tient un sujet, mais il tourne facilement en boucle et ses faits sont inventés. Il n'a vu que 11 % d'un seul des quatorze fichiers du corpus.
Il est plus lent que les outils industriels. De 1,1 à 1,7 fois à opérations égales en simple précision. En précision mixte, 2,8 fois plus lent que PyTorch avec son attention fusionnée, alors que nous avons maintenant la nôtre : l'écart tient au reste, et surtout aux activations que PyTorch garde en 16 bits.
Chaque variante n'a été entraînée qu'avec une seule graine, et avec le taux d'apprentissage réglé pour la référence. Les écarts de quelques millièmes des articles 12 à 16 sont donc fragiles. Les comparaisons à durée égale reposent en plus sur une courbe de référence interpolée entre quatre points.
Trois résultats publiés n'ont pas été retrouvés. L'équilibrage par biais de DeepSeek-V3, NorMuon et SAMuon-lite ne font pas mieux ici que ce qu'ils remplacent. Nos lots sont soixante fois plus petits que ceux de ces articles : cela ne les contredit pas.
Deux points envisagés sont restés ouverts. La comparaison avec cutile-rs, l'autre modèle de programmation de CUDA Rust, demande une version plus récente du toolkit CUDA que celle de la machine. Et le chargement asynchrone des lots n'a pas été implémenté, parce que la mesure a montré qu'il n'y avait rien à y gagner.
Une question est restée sans réponse. Deux entraînements identiques sur GPU ne donnent pas exactement les mêmes poids. L'addition atomique des gradients d'embedding est suspectée, mais cela n'a pas été vérifié.
Ce qui a rendu le projet faisable
Avec le recul, quatre habitudes ont compté.
Toujours avoir une référence
Chaque morceau rapide ou subtil du projet est adossé à une version lente et évidente.
- Chaque kernel GPU est comparé à l'opération CPU.
- Chaque règle de gradient est comparée aux différences finies.
- Le BPE incrémental est comparé au BPE naïf.
- La génération avec cache est comparée à la génération sans cache.
- Le mélange d'experts clairsemé est comparé à sa version dense.
- La matmul sur Tensor Cores est comparée au CPU, au bit près sur des entiers.
- Le modèle entier est comparé à GPT-2.
Dans ce domaine, un bug ne fait presque jamais planter le programme. Il produit des nombres, plausibles et faux. Sans référence, on ne voit rien.
Avancer par étapes vérifiables
Le projet n'a jamais eu deux choses nouvelles à déboguer en même temps. Quand le Transformer a été assemblé, tous les kernels et toute la rétropropagation étaient déjà testés : les seules erreurs possibles étaient dans l'assemblage.
L'ordre a parfois semblé lent. Écrire un bigramme, puis un tenseur CPU qu'on n'utilisera jamais pour entraîner, avant de toucher au GPU, donne l'impression de reculer. C'est ce qui a permis d'aller vite ensuite.
Trois contrôles avant chaque entraînement
La perte initiale doit valoir ln(vocabulaire). Le modèle doit pouvoir apprendre un lot
par coeur. Le CPU et le GPU doivent suivre la même courbe. Ces trois vérifications prennent
quelques secondes et épargnent des heures.
Mesurer avant de croire
La campagne de performance a contredit l'intuition à plusieurs reprises. Des tuiles plus grosses étaient plus lentes. Lire à l'avance ne servait à rien. Le chargement des données, cible classique d'optimisation, pesait un dix-millième du temps. Et près de la moitié du temps de départ partait dans la libération de mémoire, ce que rien ne laissait soupçonner.
Ce qu'on a appris sur les LLM
Il n'y a pas de magie. Un GPT est une trentaine de primitives sur des tableaux, répétées.
L'attention tient en dix lignes. La rétropropagation est une boucle à l'envers avec un
match. Ce qui impressionne dans un grand modèle ne vient pas de la complexité du
mécanisme, mais de l'échelle.
Presque tout est de la multiplication de matrices. Elle représentait 84,5 % du temps d'un pas. C'est pour cela que les cartes graphiques, puis leurs circuits spécialisés, ont compté autant dans l'histoire récente du domaine.
La qualité se paie en calcul. À aucun moment un modèle de cette série n'a manqué de données : pertes d'entraînement et de validation sont restées confondues du début à la fin. Ce qui bornait le résultat, c'était le nombre de tokens qu'on avait le temps de voir. D'où l'importance de tout le travail de performance.
Le choix des tokens compte autant que le modèle. Passer des octets aux sous-mots a presque multiplié par quatre la quantité de texte couverte par chaque token du contexte, sans toucher à l'architecture.
Une idée se mesure contre le bruit, et contre le temps. Deux entraînements identiques à la graine près diffèrent de 0,002 : tout écart plus petit ne dit rien. Et un gain à nombre de pas égal peut s'inverser à durée égale : c'est arrivé au mélange d'experts, à la prédiction multi-token et à la double traversée. Les articles 12 à 16 ont produit plus de résultats mitigés que de succès.
Une idée trop chère est parfois une idée mal écrite. Muon perdait à durée égale à l'article 15, et l'explication semblait toute trouvée : il lui faut des lots énormes. Le profileur en a donné une autre, 1 680 petits produits matriciels par pas. Une fois ceux-ci regroupés, Muon gagne. Avant de conclure qu'une technique n'est pas faite pour soi, il faut regarder où part son temps.
Une technique répond à un problème d'échelle. L'attention latente économise un cache qui ne nous gêne pas. Les experts promettent un coût de calcul inchangé que nos kernels ne tiennent pas. Une seconde traversée économise des paramètres, quand c'est le temps qui nous manque. Ces idées ne sont pas mauvaises : elles répondent à d'autres contraintes que les nôtres.
Ce qu'on a appris sur les GPU
Un GPU ne boucle pas. Il exécute la même fonction des milliers de fois en parallèle. Écrire un kernel, c'est reformuler un calcul pour que chaque thread sache quoi faire avec son seul numéro.
Tout est une question de mémoire. Les calculs sont rarement le facteur limitant. Ce qui coûte, c'est d'amener les données jusqu'aux calculs. Presque chaque optimisation de l'article 11 a consisté à rapprocher des données : de la mémoire globale vers la mémoire partagée, puis vers les registres. L'attention fusionnée de l'article 17 l'a confirmé d'une autre façon : elle ne fait pas une multiplication de moins que l'attention en trois opérations, elle évite seulement d'écrire ses résultats intermédiaires en mémoire. Cela a suffi à faire tomber un mur.
L'asynchronisme est à la fois un cadeau et un piège. Il permet à l'hôte de préparer la suite pendant que la carte calcule. Il rend aussi toute mesure de temps naïve complètement fausse.
Ce qu'on a appris sur Rust
Rust n'est pas le langage habituel de ce domaine, et le projet permet d'en juger sur pièces.
Le système de types a rendu de vrais services. Le contrat Backend a permis d'écrire le
modèle une fois et de le faire tourner sur CPU, sur GPU et sur un profileur. Le
DisjointSlice a rendu impossibles, par construction, les race conditions sur les
écritures. Et pouvoir appeler une fonction Rust ordinaire depuis un kernel garantit que
l'hôte et la carte appliquent la même formule.
L'écosystème, en revanche, est jeune. CUDA Rust est un projet de recherche, et cela s'est vu : une chaîne d'outils épinglée à une nightly précise, des versions de crates en conflit, un fichier de kernels que deux compilations simultanées corrompent, un outil d'inspection qui regardait le mauvais fichier. Aucun de ces obstacles n'était grave, mais chacun a coûté du temps. Tenir le code GPU dans son propre paquet a limité les dégâts : le reste du projet ne dépend que de Rust stable.
Pour aller plus loin
Le projet peut se prolonger dans plusieurs directions, par ordre de difficulté croissante.
- Entraîner plus longtemps, avec Muon. C'est le gain le plus sûr : il suffit de
lancer un entraînement avec plus de pas et
--optimizer muon:short, et la reprise le protège des interruptions. Avec l'attention fusionnée, un pas de Muon coûte environ 1,2 fois celui d'AdamW, soit à peu près 49 000 tokens par seconde : un fichier entier du corpus se parcourt en moins de cinq heures. - Supprimer les transpositions en apprenant à la matmul à lire une matrice transposée.
- Fusionner le biais dans la multiplication, et son gradient en un seul kernel.
- Stocker les activations sur 16 bits, comme le fait PyTorch en précision mixte.
- Étendre l'attention fusionnée aux têtes de 96 composantes de l'attention latente, qui passe encore par trois opérations.
- Refaire les comparaisons des articles 12 à 16 avec plusieurs graines, chercher le taux d'apprentissage de Muon, puis le combiner aux pièces qui survivent.
- Écrire l'attention clairsemée, sur un contexte de 2 048 positions.
Le mot de la fin
Le but n'était pas d'obtenir un bon modèle. Il existe des outils bien meilleurs pour cela, et ils tiennent en trois lignes de Python.
Le but était de pouvoir lire ces trois lignes en sachant ce qu'elles déclenchent : les octets qui deviennent des tokens, les tokens des vecteurs, les requêtes qui cherchent leurs clés, le gradient qui remonte la tape, les threads qui se partagent une tuile. Rien de tout cela n'est plus une boîte noire.