Het trainen van een 3.8B LLM tot 0.384 CORE voor $998
Er ligt een groot, onderbelicht gebied tussen de 'nanoGPT-speeltjes' en de noodzaak voor een volledig onderzoekslaboratorium, waar één persoon met een paar duizend dollar een betekenisvol model kan trainen.
Ik wilde zelf ervaren hoe taal en begrip ontstaan uit willekeurige gewichten, en de aspecten leren die je alleen kunt leren door vanaf nul te beginnen. Dit project is geschreven in de avonduren, gedebugd op een RTX 5090 en voltooid op gehuurde B200's. Het is sterk geïnspireerd door nanochat van Andrej Karpathy.
Het resultaat is een model met 3,8 miljard parameters dat een score van 0,384 behaalt op CORE. Het is getraind op 65 miljard tokens in 43 uur voor een bedrag van $998.
Wat volgt is een overzicht van wat werkte, wat niet werkte, en wat ik nog steeds niet weet.
| Model | Parameters | Tokens | Hardware | Tijd | Kosten | CORE |
|---|---|---|---|---|---|---|
| GPT-2 (OpenAI) | 1.5B | — | — | — | — | 0.2565 |
| nanochat d26 | ~561M | 11.2B | 8× H100 | ~3u | — | ~0.258 |
| nanochat d32 | ~1B | — | 8× H100 | ~33u | ~$1000 | 0.310 |
| little-lm (1024 ctx) | 3.848B | 57.3B | 8× B200 | 35.9u | $820 | 0.338 |
| little-lm (2048 ctx) | 3.848B | 65.3B | 8× B200 | 43u | $998 | 0.384 |
Mijn model is groter dan nanochat d32 en kostte een vergelijkbare hoeveelheid tijd. B200's boden een betere waarde per eenheid werk dan H100's. Voor ongeveer hetzelfde bedrag als de $1.000-configuratie van nanochat, scoort dit model aanzienlijk hoger. Dit is een bemoedigend datapunt over wat bereikbaar is buiten een laboratorium of een megabedrijf met miljoenen aan compute-budget. Naarmate de grens verschuift, brengt $1.000 je steeds verder.
Setup
Ik heb little-lm gebouwd als een config-driven framework voor het trainen van kleine decoder-only LLM's. Elke run wordt volledig gespecificeerd door een YAML-bestand: model, dataset, optimizer, schedule en callbacks. Componenten registreren zichzelf in een globale registry en worden op naam opgelost, waardoor het wisselen van een optimizer of dataset een wijziging van één regel in de config is.
Goede infrastructuur betaalt zich bijna onmiddellijk terug. Gewone software engineering discipline (zoals separation of concerns, schone interfaces en uitwisselbare componenten) is zeer belangrijk in AI-werk. Het kostte me in het begin wat tijd, en later nog een paar keer om slechte contracten of suboptimaliteiten te herstellen. Maar deze investering in tijd betaalt zich terug bij het eerste convergentieprobleem dat je tegenkomt. Ik ontdekte dat een goede infra een infra is die bijna nooit vereist dat je handmatig code bewerkt. Als je de config kunt lezen en precies begrijpt wat er gebeurt, zonder verborgen mechanismen, dan heb je goed werk geleverd. Het volgende rapport is het resultaat van het kunnen uitdrukken van experimenten als een YAML-diff van drie regels in plaats van een aparte branch.
Het uiteindelijke model is Llama-stijl: RMSNorm, RoPE, GQA (24 query heads, 8 KV heads), relu² MLPs, QK-norm, logit softcap, per laag leerbare residual scalars, en ResFormer-stijl value embeddings.
Componenten en Parameters
| Component | Parameters |
|---|---|
| Token embeddings | 154.5M |
| LM head (untied) | 154.5M |
| 28 decoder layers | 2,818.7M |
| Value embeddings (14 tabellen) | 721.2M |
| Totaal | 3.848B |
Het is vermeldenswaardig dat de value embeddings 19% van het totale aantal parameters beslaan. Dit betreft 14 tabellen van vocab × kv_dim, één per twee lagen.
Resultaten
Vroege experimenten
Voor de goede runs waren er veel slechte.
Ik trainde een 858M Llama op FineWeb-Edu voor 16,4 miljard tokens, gedurende 5,8 dagen op een enkele A100. De instellingen waren: AdamW op 2.5e-4, cosine decay naar nul, 5% warmup, batch 256 via gradient accumulation, 2048 context.
Het resultaat: PIQA 60,45%. GPT-2 124M scoort ongeveer 63%. Ik had zes dagen aan compute besteed om iets te bouwen dat slechter was dan een model uit 2019 dat zeven keer kleiner was. De generaties waren repetitief en bijna onzinnig.
De verliescurve (loss curve) vertelde het verhaal:
- Cosine decay naar nul: De curve werd volledig vlak na ongeveer 70% van de stappen. De laatste 30% van het compute-budget leverde in essentie niets op, omdat de learning rate mogelijk te laag was. Een lineaire cooldown houdt de rate veel langer nuttig.
- Peak LR te conservatief: 2.5e-4 is laag voor 858M parameters. Je kunt bij zulke kleine modellen behoorlijk agressief zijn.
- AdamW op alles: Muon zou bij deze schaal aanzienlijk beter moeten zijn per token voor de matrixparameters. Dit werd snel aangetoond in ablation runs.
- De data: FineWeb-Edu is redelijk, maar niet het beste beschikbare.
Uit deze post-mortem kwamen vijf wijzigingen voort. Samen vormen zij het verschil tussen de bovengenoemde run en een model dat GPT-2 met een ruime marge verslaat:
- Trapeziumvormige LR-planning: Warmup voor 5%. Houd het daarna vlak en eindig met een lineaire cooldown over de laatste 50% naar 5% van de piek. Het doel is dat het model tot het einde blijft leren in plaats van uit te rollen in de staart. In de 3.8B run daalde de eval loss nog steeds bij de laatste stap, wat precies het gedrag is dat de 858M run niet vertoonde.
- Muon voor matrixparameters, AdamW voor de rest: Muon is langzamer per stap (Newton-Schulz orthogonalisatie is niet gratis, ongeveer 25% in een shallow-accumulation benchmark), maar die kosten worden slechts één keer per optimizer-stap betaald. Bij 7 gradient-accumulation stappen verwatert dit tot ~4%. Gemeten tegen de totale looptijd is de convergentie over het geheel genomen veel sneller.
- ClimbMix in plaats van FineWeb-Edu: Dit zorgde voor een enorme sprong in convergentiesnelheid, precies zoals Karpathy ook ontdekte.
- FP8 + vocab padding: FP8-training via
torch.scaledmmmet dynamische tensorwise scaling op alle drie de GEMMs, en het padden van de vocab van 50.257 naar 50.304 (een veelvoud van 64) zodat de tensor cores efficiënt werken. Samen leverde dit een throughput-stijging van +33% op, hoofdzakelijk door FP8. - 1024 context in plaats van 2048: Het halveren van de context verdubbelt ruwweg de batch size bij een vast geheugen. De throughput per token verandert nauwelijks. We worden nog steeds gedomineerd door de MLPs, wat een goed teken is dat we de hardware effectief gebruiken.
Verloop van de volledige run
| Stap | Tokens | Eval loss | CORE |
|---|---|---|---|
| 2,500 | 5.7B | 2.3278 | 0.2389 |
| 5,000 | 11.5B | 2.2072 | 0.2752 |
| 7,500 | 17.2B | 2.1571 | 0.2934 |
| 10,000 | 22.9B | 2.1269 | 0.3104 |
| 12,500 | 28.7B | 2.1075 | 0.3147 |
| 15,000 | 34.4B | 2.0710 | 0.3224 |
| 17,500 | 40.1B | 2.0395 | 0.3294 |
| 20,000 | 45.9B | 2.0160 | 0.3267 |
| 22,500 | 51.6B | 1.9963 | 0.3345 |
| 25,000 | 57.3B | 1.9868 | 0.3384 |
In steady state werd er ongeveer 480.000 tokens/sec verwerkt, wat 57,3 miljard tokens in 33 uur plaatst. De werkelijke tijd was 35,9 uur. Het verschil wordt veroorzaakt door de CORE-evaluaties, die elk ongeveer 15 minuten duurden (tien stuks over de hele run) en 7% van het totaal in beslag namen.
Het opnieuw uitvoeren van exact dit recept met een context van 2048 tokens resulteerde in een score van 0,3840. Bijna al dat verschil bleek te komen door taken die zeer contextafhankelijk zijn.
Wat betreft de GPU's: 92% SM-activiteit en 40% SM-bezetting (occupancy). Hoge activiteit betekent dat de SM's bijna nooit stilstonden. Er was geen sprake van dataloader-starvation of netwerkwachten, wat het resultaat is van het lokaal downloaden van de shards in plaats van streamen (wat ons kwetsbaar zou maken voor kleine netwerkhaperingen van Hugging Face). De lage bezetting is kenmerkend voor opeenvolgende grote GEMMs: matmul-kernels ruilen bezetting in voor register-tile grootte. Het systeem was compute-bound en goed gevoed, wat een sterk signaal is dat de hardware goed wordt benut.
Dit komt neer op ongeveer 1.047 TFLOP/s sustained per B200, of ~25% MFU tegenover de dense FP8 piek van Blackwell. (Tegenover de bf16 piek leest dit als 50%, wat belangrijker is omdat niet alle lineaire lagen in FP8 draaien). De distributed strategy was simpel DistributedDataParallel. Bij 3.8B op een enkele node was gradiëntcommunicatie nooit de beperkende factor, en sharded-optimizer machinery bleek onnodig.
Doorvoer (Throughput) verhogen
Het huren van GPU's is niet goedkoop. Terwijl je op het werk vaak eerst denkt aan de kwaliteit van het model, telt de doorvoer veel zwaarder wanneer het je eigen geld is.
Ik heb hier echt werk in gestoken op een enkele RTX 5090 voordat ik een node huurde. Baseline 858M model, bf16, compiled: 26.144 tok/s. Eindresultaat: 37.621 tok/s.
- FP8 (+25%): Alle drie de GEMMs (1 forward en 2 backwards) in FP8 met dynamische tensorwise scaling. Dit vereist SM90+, maar levert een mooie sprong in doorvoer op.
- Vocab padding (+33% cumulatief): Het padden van 50.257 → 50.304 kost 47 ongebruikte embedding-rijen en ontgrendelt het snelle tensor-core pad. Dit is bijna gratis.
- Fused linear cross-entropy (+44% cumulatief): Liger’s
FusedLinearCrossEntropyLossfuseert delm_headmatmul in de loss en hakt deze intern in stukken, zodat de volledige (B*T, vocab) logits tensor nooit volledig in het geheugen wordt geplaatst. Bij dezelfde batch size is het 6% langzamer:
| Config | Throughput | VRAM |
|---|---|---|
| Baseline CE, batch 6 | 34,724 tok/s | 27,852 MiB |
| Fused CE, batch 6 | 32,952 tok/s | 19,630 MiB |
| Fused CE, batch 8 | 35,979 tok/s | 24,028 MiB |
| Fused CE, batch 10 | 37,621 tok/s | 28,872 MiB |
Hoewel het per stap langzamer is, wint het een aanzienlijke hoeveelheid VRAM terug (8 GB op mijn 5090), waardoor de toename in micro-batch size het verlies van 6% meer dan compenseert.
- Niet-gated MLPs: Het weglaten van de gate projection (SwiGLU → relu², twee matmuls in plaats van drie) bij het kleine model verhoogde de snelheid van 183.035 → 214.173 tok/s en bespaarde 6 GB VRAM. Eén kanttekening uit de ablations: een SwiGLU intermediate ratio van 2,75 werkt niet bij relu². Het model leert merkbaar slechter. Gebruik 4× voor niet-gated.
- bf16 master weights: Het houden van de optimizer master weights in bf16 in plaats van fp32 verminderde het VRAM-gebruik met 27% en verhoogde de throughput van 640K naar 1,4M tok/s op de 1.5B config. Dat was een enorme versnelling (2,2×). Het kwaliteitsverlies is reëel maar klein: CORE 0,22 vs 0,23 na 4.000 stappen. Wanneer je optimaliseert voor capaciteit per dollar, is zorgvuldige dtype-afhandeling een van de krachtigste knoppen waar je aan kunt draaien.
- Hardware: Dezelfde code, 150M model, FP8: RTX 5090 op 184.662 tok/s, B200 op 477.440 tok/s. Een factor 2,59 door hardware alleen, nog zonder rekening te houden met de extra VRAM die een grotere batch size mogelijk maakt.
Wat niet werkte
- Document-boundary masking met flex attention: Het inpakken van documenten in één sequentie laat tokens over grenzen heen kijken. Ik heb dit correct geprobeerd op te lossen met per-token document ID's en een mask, maar ik heb het uiteindelijk verwijderd. Andrej Karpathy ontdekte ook dat cross-document leakage niet veel slechter is bij BOS-aligned packing. Best-fit packing verving dit in ongeveer 10 regels code.
- Liger RMSNorm en RoPE: RoPE was 2,2× sneller in een microbenchmark, maar produceerde geen meetbare verandering in de end-to-end throughput. RMSNorm was ronduit langzamer dan de ingebouwde
F.rms_normvan PyTorch 2.9 (0,41ms vs 0,25ms). Beiden zijn teruggedraaid. - Nanochat-stijl initialisatie: Embeddings op N(0, 0.8), lineaire gewichten uniform, output projecties nul-geïnitialiseerd zodat de residual stream begint als pure identity, en LM head op N(0, 0.001). Theoretisch veel mooier dan de N(0, 0.02) van GPT-2. De loss curve startte marginaal lager en overlapte na ~1.500 stappen. Geen meetbaar kwaliteitsverschil. Ik heb het behouden voor de esthetiek, niet op basis van bewijs.
- Streaming datasets: Geweldig om mee te beginnen, verkeerd voor een echte run. Zelfs als het netwerk gezond lijkt, gaven lokale shards 2-3% meer throughput, en incidentele netwerkdips kosten veel meer dan dat. Voor runs langer dan een paar uur is het de moeite waard om de data één keer vooraf te downloaden.
Ablation op value-embedding
Value embeddings besloegen 721M parameters voor een 3.8B model. Ik trainde hetzelfde model met dezelfde config maar met value_embeddings: false en vergeleek dit met de originele run tot 12.500 stappen (29B tokens).
| Parameters | Loss @12.5K | CORE @12.5K | Throughput |
|---|---|---|---|
| Value embeddings aan (3.848B) | 2.1075 | 0.3147 | 479,445 tok/s |
| Value embeddings uit (3.128B) | 2.1171 | 0.3047 | 477,908 tok/s |
Dit resulteerde in 0,46% betere loss en 3,2% betere CORE, voor 19% meer parameters. De throughput is identiek, omdat value embeddings lookups zijn. Ze kosten geheugen en optimizer-state, maar essentieel geen FLOPs.
Twee interessante bevindingen:
- Value embeddings kochten het equivalent van ongeveer 1.200 trainingsstappen. 19% meer parameters is dus ongeveer 5% meer training waard.
- CORE bewoog ongeveer zeven keer meer dan de loss (3,2% vs 0,46%), en het gat kromp gestaag tijdens de training. Dit is belangrijk als je CORE gebruikt om beslissingen te nemen: het is een accuratesse-metriek, dus items nabij de beslissingsgrens klappen om bij minuscule logit-wijzigingen.
Value embeddings zijn nuttig voor een klein model en kosten bijna niets qua throughput. Een beetje VRAM hieraan besteden geeft het model een vorm van bias naar bepaalde concepten die nuttig kunnen zijn voor CORE.
Discussie
Misleidende micro-benchmarks
Men zou geneigd zijn te geloven dat 1024 tokens context voldoende is voor een hoge CORE-score. Maar bij het analyseren van de per-taak logs blijkt dit onjuist voor taken die zeer contextgevoelig zijn.
Drie van de 22 CORE-taken hebben prompts die vrijwel nooit in 1024 tokens passen:
| Taak | Prompts afgekapt | Stap 2.5K | Stap 25K |
|---|---|---|---|
| squad | 10570 / 10570 (100%) | 0.1478 | 0.0000 |
| boolq | 3265 / 3270 (99.8%) | 0.5798 | 0.5131 |
| bigbenchlanguageid | 9965 / 10000 (99.7%) | 0.2454 | 0.2538 |
SQuAD is het meest opvallend. Het stagneert niet, maar daalt monotoon naar precies nul. Het model wordt gestaag slechter in deze taak naarmate het langer traint.
Twee details verklaren dit. SQuAD is een 10-shot taak in de DCLM bundle, dus elke prompt bevat tien uitgewerkte voorbeelden gevolgd door de eigenlijke vraag. De mediaan is 1.998 tokens. Geen enkele past in 1024. Wanneer een prompt te lang is, houdt mijn harness de laatste maxseqlen tokens vast.
De testpassage staat aan het einde, dus die overleefde altijd. Wat werd afgekapt, waren de tien demonstraties. Het model las de passage en de vraag, maar zag bijna nooit de voorbeelden die het de verwachte output-formattering leren. Omdat SQuAD wordt gescoord op exact-token match, scoort vloeiende proza elke keer nul.
Dit verklaart ook de daling. Een vroeg, high-entropy model produceert incidenteel iets korts en generieks dat toevallig matcht. Naarmate het model scherper wordt, committeert het zich aan goed gevormde voortzettingen en verdwijnen de toevallige hits. Ironisch genoeg zorgde het beter worden in taal voor een slechtere score bij het gokken.
Boolq vertoont een mildere versie van hetzelfde patroon. Taalidentificatie beweegt zich helemaal niet weg van de kansberekening. Kortom, de score van 0,338 werd gemeten terwijl drie van de 22 taken bijna nul scoorden om redenen die niets te maken hadden met modelkwaliteit, maar enkel met de contextlengte.
Het effect van grotere context
Als we de hoogst mogelijke CORE-score willen, hebben we een grotere context nodig. Dit heeft echter gevolgen voor de training throughput.
Door de contextlengte te verdubbelen en de micro-batch te halveren om het VRAM constant te houden, bleven de tokens per optimizer-stap identiek. CORE ging van 0,3384 naar 0,3840.
Bij stap 20.000 hadden de twee runs vrijwel dezelfde eval loss (2,0160 vs 2,0164), maar verschilden ze met 0,034 op CORE. Het was verrassend om zo'n lage correlatie te zien tussen CORE en eval loss op de ClimbMix dataset.
| Taak | 1024 | 2048 | Afgekapt |
|---|---|---|---|
| squad | 0.0000 | 0.3114 | 100% → 47% |
| boolq | 0.5131 | 0.7095 | 99.8% → 3.2% |
| bigbenchlanguageid | 0.2538 | 0.2585 | 99.7% → 14% |
| overige 19 taken | +0.008 gecombineerd |
SQuAD en boolq alleen zijn verantwoordelijk voor 83% van de winst. De overige twintig taken bewogen in totaal +0,008, wat ongeveer is wat 14% meer tokens op zichzelf oplevert. Enkele taken werden slechter: commonsenseqa daalde met 0,072 en csalgorithms met 0,031. Over 22 taken is beweging in beide richtingen verwacht.
Een context van 2048 was een waardevolle beslissing voor de meting, niet per se voor de kwaliteit. Het kostte 9% throughput (480K → 437K tok/s) en leverde buiten de context-gebonden taken bijna niets op. 1024 is prima voor training, terwijl 2048 sommige specifieke taken ontgrendelt.
Toekomstig werk
Beperkingen
Er zijn vier zaken die ik niet heb geablateerd:
- Peak LR: Gebaseerd op nanochat's $\sqrt{768/d_{model}}$. Ik wilde geen geld uitgeven aan het scannen van learning rates.
- Planning: Ik ben overgestapt van cosine naar trapeziumvormig, maar er kunnen efficiëntere schedules bestaan.
- QK-norm: Stond standaard aan en is nooit uitgeschakeld.
- GQA ratio: Gebruikt als een handige knop om geheugen te besparen.
De meeste van deze keuzes zijn overgenomen van nanochat. Dat is een verdedigbare manier om een klein budget te beheren — iemand anders heeft het experiment al betaald — maar het betekent dat ik erop vertrouw dat de resultaten van Karpathy overdraagbaar zijn naar mijn model, data en schaal.
Open vragen
Er is veel interessant werk dat ik zou willen doen met meer tijd en middelen:
- Value embeddings versus herallocatie: De vergelijking was VE tegenover niets. De relevante vraag is VE tegenover het besteden van die 721M parameters aan iets anders.
- 1024 versus 2048 bij gelijke wall-clock tijd: De rerun veranderde de context en duurde langer, wat de meetvraag beantwoordt, maar niet de kwaliteitsvraag.
- Waarom commonsense_qa regreseedde met 0,072 bij een langere context, terwijl die taak geen lange prompts bevat.
- Sharding van de optimizer, zoals in nanochat. Ik gebruikte plain DDP met een single-GPU Muon, wat betekent dat elke rank een volledige kopie van de optimizer-state houdt. nanochat gebruikt ZeRO-2 sharding in de optimizer. De winst in geheugen is zeker, en vrijgekomen geheugen vertaalt zich in batch size, wat weer tokens betekent voor hetzelfde budget.
- Additionele data-exploratie: Ik heb weinig tijd gehad voor data-analyse op zowel de CORE benchmark als de ClimbMix dataset. Ik ben ervan overtuigd dat dit zou helpen om nog hogere prestaties te behalen met hetzelfde compute-budget.
Slotbeschouwing
GPT-2 was in 2019 een grensverleggend resultaat, geproduceerd door een goed gefinancierd lab met een groot team. Het 1.5B model scoort 0,2565 op CORE. Zeven jaar later heb ik dat met een ruime marge verslagen in mijn avonduren, voor $998, op hardware die ik per uur huurde.
De grens is verschoven, en alles is meegestegen. Werk waarvoor vroeger een laboratorium nodig was, kan nu door één engineer in de avonduren worden gedaan. Ik vraag me af wat voor krankzinnige machines we over zeven jaar kunnen bouwen!
Appendix: De configuratie
De volledige run, samengevoegd uit de YAML-bestanden in één blok:
model:
hidden_size: 3072
intermediate_size: 12288 # 4x, non-gated
num_hidden_layers: 28
num_attention_heads: 24
num_key_value_heads: 8 # 3:1 GQA
head_dim: 128
hidden_act: relu2
gated_mlp: false
qk_norm: true
logit_softcap: 15.0
layer_scale: true
value_embeddings: true # 14 tables, alternating layers
tie_word_embeddings: false
rope_theta: 10000.0
rms_norm_eps: 1.0e-6
vocab_pad_to: 64 # 50257 -> 50304
max_position_embeddings: 2048
dtype: bf16
engine:
compile: true
fp8: true
precision: bf16
total_batch_size: 2293760 # 20 x 2048 x 7 grad_accum x 8 GPUs
loss: LigerFusedLinearCrossEntropyLoss(softcap=15.0)
optimizer: # composite, one group per parameter class
matrix: Muon lr=0.02 momentum=0.95 wd=0.0
embeddings: AdamW lr=0.1414 betas=(0.8, 0.995) eps=1e-10 wd=0.001
lm_head: AdamW lr=0.002828 betas=(0.8, 0.96) eps=1e-10 wd=0.01
value_embeds: AdamW lr=0.0707 betas=(0.8, 0.995) eps=1e-10 wd=0.01
scalars: AdamW lr=0.005 betas=(0.8, 0.95) eps=1e-10 wd=0.05
scheduler:
trapezoidal:
warmup_ratio: 0.05
warmdown_ratio: 0.50
final_lr_frac: 0.05
data:
dataset: nvidia/Nemotron-ClimbMix (karpathy/climbmix-400b-shuffle shards)
tokenizer: gpt2 (tiktoken)
block_size: 2048
packing: best-fit, BOS-aligned
batch_size: 20 per rank
num_workers: 11
trainer:
max_steps: 32000 # stopped at ~28,000 -> 65.3B tokens
eval_every: 4000 # must divide max_steps or the final CORE is skipped
De AdamW learning rates volgen de $\sqrt{768/d_{model}}$ schalingsregel van nanochat; de Muon LR van 0,02 is daar ook van overgenomen.
Appendix: Voorbeelden van tekstgeneratie
- The capital of France is Paris. It is the largest city in France and the second largest city in Europe
- The french revolution happened in 1789 and 1799, and was a time of great change in france
- *At the center of the milky way there is a supermassive black hole. It is called Sagittarius A\**
- Electrons orbit around the nucleus of an atom in a series of energy levels. The energy levels are numbered
- Newton discovered the laws of motion and gravity. He also discovered the law of universal gravitation. Newton's
Groetjes,