Contexte et objectifs
Hugo Vergnes a cherché à démontrer qu’un modèle de 3,8 milliards de paramètres pouvait être entraîné en dehors d’un laboratoire de recherche, avec un budget inférieur à 1 000 $. Le projet, nommé little‑lm, vise à reproduire les performances de modèles comme nanoGPT tout en restant accessible à un individu disposant de quelques milliers de dollars.
Architecture et optimisations
Le modèle suit le design Llama : RMSNorm, RoPE, GQA (24 têtes de requête, 8 têtes KV), MLPs relu², normalisation QK, logit soft‑cap et scalaires résiduels par couche. Les embeddings de valeur représentent 19 % du nombre total de paramètres, répartis sur 14 tables vocabulaire × kv_dim, insérées à chaque seconde couche. Trois optimisations majeures ont été appliquées : utilisation de l’optimiseur Muon pour les matrices, entraînement en FP8 via
torch._scaled_mm(...) avec mise à l’échelle dynamique tensor‑wise, et remplissage du vocabulaire de 50 257 à 50 304 tokens (multiple de 64) afin d’optimiser l’utilisation des cœurs tensoriels.Stratégie d’entraînement et performances
Le processus a consommé 65 milliards de tokens en 43 heures, soit environ 480 000 tokens/s en régime stable. Le temps total de mur était de 35,9 h, la différence étant due aux évaluations CORE (15 min chacune, 7 % du temps). Le LR a suivi un planning trapézoïdal : 5 % de warm‑up, plateau, puis décroissance linéaire sur les 50 % finaux, évitant la stagnation observée sur un modèle de 858 M paramètres. Le contexte a été limité à 1024 tokens, doublant la taille de lot à mémoire constante et maintenant le débit token‑wise. Sur les GPU B200, l’activité SM a atteint 92 % avec une occupation de 40 %, traduisant une utilisation compute‑bound efficace. Le débit soutenu s’élève à 1 047 TFLOP/s par B200, soit ~25 % du pic FP8 dense de Blackwell et ~50 % du pic bf16, ce qui confirme une exploitation optimale du matériel.
Analyse des coûts et limites
Le coût total, 998 $, comprend la location de deux B200 et le temps de calcul. Comparé à la configuration nanochat (~1 000 $), le modèle atteint 0,384 CORE contre 0,3840 avec un contexte de 2048 tokens, montrant que la réduction de contexte impacte légèrement les scores dépendants du contexte. L’absence de sharded‑optimizer et la simple utilisation de DistributedDataParallel indiquent que la communication n’est pas un goulot d’étranglement à cette échelle. Toutefois, la dépendance à FP8 nécessite du matériel SM90+, limitant la portabilité. Le rapport ne fournit pas de mesures de généralisation au‑delà de CORE, et la robustesse face à des jeux de données plus diversifiés reste à évaluer.