smallbard

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èleParamètresPerteHésite entre
Hasard uniforme05,55257 octets
Bigramme66 0492,5312,6 octets
MLP à contexte fixe402 7211,906,7 octets
GPT, octets834 4321,363,9 octets
GPT, octets10 844 5441,203,3 octets
GPT, sous-mots29 676 5440,942,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 :

VariantePerteDébit, tokens/s
GPT-2, le modèle de l'article 101,00050 700
+ RMSNorm, SwiGLU, RoPE0,95952 300
+ attention latente (cache 6,4 fois plus petit)0,98446 600
+ mélange d'experts équilibré0,940 à 0,94433 000
+ prédiction multi-token0,95542 200
+ Muon à la place d'AdamW, article 150,92522 500
+ Muon, matrices empilées, quatre itérations0,92643 900
+ deux traversées des couches0,95227 600
+ clés partagées par 4 couches sur 8 (cache 2 fois plus petit)0,96455 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 :

VarianteSa perteSa duréeLa référence, même durée
Mélange d'experts0,94025 min0,925
Prédiction multi-token0,95519 min0,942
Deux traversées0,95230 min0,913
Clés partagées, 4 sur 80,96415 min0,963
Muon de l'article 150,92536 min0,900
Muon empilé, quatre itérations0,92619 min0,946
Le même, sur 36 minutes0,88236 min0,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 :

MomentTemps par pasTokens par seconde
Avant la seconde campagne d'optimisation652 ms12 300
Après, en f32224 ms36 500
Après, sur Tensor Cores157 ms52 000
Avec l'attention fusionnée127 ms64 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.

  1. 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.
  2. Supprimer les transpositions en apprenant à la matmul à lire une matrice transposée.
  3. Fusionner le biais dans la multiplication, et son gradient en un seul kernel.
  4. Stocker les activations sur 16 bits, comme le fait PyTorch en précision mixte.
  5. Étendre l'attention fusionnée aux têtes de 96 composantes de l'attention latente, qui passe encore par trois opérations.
  6. 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.
  7. É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.