un — Träna en språkmodell: Backprop i anpassad CUDA [DESCRIPTION /]

un

gäst

1 / ?

Lokala gradienter multipliceras

Forward & Backward Kernels

Framåtgåendet

ANDREA-120M:s framåtgående leder inmatningen genom en sekvens av operationer:

x = embed(token_ids)         # tokeninbäddningar
for layer in 12_layers:
x = x + attn(LN(x))      # attention-delblock
x = x + mlp(LN(x))       # MLP-delblock
logits = LN(x) @ embed.T     # bunden utdata-projektion
loss   = cross_entropy(logits, targets)

Varje operation läser in tensorerna & producerar ut tensorerna. Forward-passen avslutas i en enda skalär: cross-entropy-förlusten för denna batch.

Bakåtpasset

Träningen uppdaterar vikterna i riktningen som minskar förlusten. För att få uppdateringsriktningar behöver motorn:

dL/dW för varje lärbar W i modellen

Kedjeregeln ger detta. För en kedja loss = f(g(h(x))):

dL/dx = (dL/df) * (df/dg) * (dg/dh) * (dh/dx)

Varje faktor är en lokal gradient: hur utdata från en operation ändras när dess indata ändras med en liten mängd. Multiplikation av lokala gradienter bakåt genom grafen sprider förlustsignalen till varje vikt.

Reverse-Mode Differentiering

Backprop beräknar gradienter i omvänd ordning: börjar från dL/dlogits = 1, sedan går bakåt genom cross-entropy, sedan utdataprojektion, sedan layer norm, sedan tolv transformerblock, sedan embeddings. Vid varje steg multiplicerar man den inkommande gradienten med den lokala Jacobimatrixen.

Reverse-mode är effektivt när utdata är en enda skalär (förlusten) & det finns många indata (vikterna). En bakåtpass producerar gradienter för varje vikt i modellen. Forward-mode skulle behöva en pass per vikt; för ANDREA-120M med ~120M vikter är forward-mode oöverkomligt.

Varför Reverse-Mode

ANDREA-120M har ~120M vikter & producerar en enda skalärloss per träningssteg. Jämför reverse-mode automatisk differentiering mot forward-mode. Ange (1) vilken mode som producerar alla viktgradienter i en enda bakåtpassning; (2) hur många forward-mode-passningar som skulle behövas för att beräkna alla 120M viktgradienter; (3) vilken mode ANDREA använder & varför.

Varje Forward-op Får En Backward-tvilling

Parningsdisciplinen

microgpt_cuda.cu levererar två CUDA-kärnor för varje operation: en som beräknar framåtutdata, en som beräknar indatagradienter givet utdatagradienter. Parningen är en-till-en:

Framåt-kärna	Bakåt-kärna	Operation
`k_embed_fwd`	`k_embed_bwd`	Token-inbäddningsuppslag
`k_layernorm_fwd`	`k_layernorm_bwd`	Lagernormalisering
`k_attn_qkv_fwd`	`k_attn_qkv_bwd`	Q, K, V-projektioner
`k_attn_fwd`	`k_attn_bwd`	Skalad prickprodukt-attention
`k_attn_out_fwd`	`k_attn_out_bwd`	Utprojektion W_O
`k_mlp_fwd`	`k_mlp_bwd`	MLP (med GELU)
`k_residual_add`	`k_residual_add_bwd`	Residualanslutning
`k_loss_fwd`	`k_loss_bwd`	Cross-entropy-förlust

Åtta operationspar täcker hela transformern. Plus några hjälpkärnor: k_grad_norm_partial, k_grad_norm_final, k_grad_scale för gradientklippning (se aktivitet 75).

En baklängeskärnas uppgift

Givet gradienten som strömmar in från senare lager (grad_output), beräknar en baklängeskärna:

1. grad_input: gradienten med avseende på operationens indatatensor. Denna skickas vidare bakåt.

2. grad_weight: gradienten med avseende på lärbara parametrar i operationen. Denna går in i optimizer-tillståndet.

Båda beräknas i en enda kernel-start. CUDA-trådar samarbetar på tiles av gradienttensorn parallellt.

Sparade tensorer

Bakåträkningen behöver ofta värden från framåtgåendet. Till exempel behöver k_layernorm_bwd medelvärdet & variansen som beräknats under framåt; k_mlp_bwd behöver GELU-preaktiveringen. Träningsmotorn lagrar dessa i dedikerade buffertar under framåt, och läser dem under bakåt.

Minnes kostnad: ungefär samma form som framåtutgången för varje sparad tensor. För ANDREA-120M med batch=8, seq=1024, d_model=768, är en sparad tensor 8 × 1024 × 768 × 4 bytes = 25 MB. Över 12 lager & flera sparade tensorer per lager dominerar aktiveringar VRAM under träning (~5-10 GB på ett 24 GB kort).

Spårning av ett bakåtsteg

ANDREA-120M slutför en framåtgång genom en transformerblock. Spåra vad som händer under bakåtgången genom samma block (i pre-norm-struktur: `x = x + Attention(LN(x))` sedan `x = x + MLP(LN(x))`). Namnge bakåtkärnorna i den ordning de aktiveras, & ange vilken framåtkärna varje en para med. Täck minst 4 kärnor.

Var gradienter lever i minnet

En gradienttensor per viktensor

Varje lärbar viktensor i ANDREA-120M har en matchande gradienttensor med identisk form. För varje block:

W_Q       [768, 768]     ↔   grad_W_Q       [768, 768]
W_K       [768, 768]     ↔   grad_W_K       [768, 768]
W_V       [768, 768]     ↔   grad_W_V       [768, 768]
W_O       [768, 768]     ↔   grad_W_O       [768, 768]
W_1       [768, 3072]    ↔   grad_W_1       [768, 3072]
W_2       [3072, 768]    ↔   grad_W_2       [3072, 768]
LN1.gamma [768]          ↔   grad_LN1.gamma [768]
LN1.beta  [768]          ↔   grad_LN1.beta  [768]
LN2.gamma [768]          ↔   grad_LN2.gamma [768]
LN2.beta  [768]          ↔   grad_LN2.beta  [768]

Plus token embeddings, position embeddings, & a final layer norm. The gradient buffer total memory matches the weight memory: ~120M floats, ~480 MB at FP32, ~240 MB at FP16.

Ackumulation över mikrobatcher

ANDREA's batch_size = 8 fits in VRAM at FP16. Larger effective batches require gradient accumulation: run multiple forward+backward passes on small batches, summing gradients into the same buffer, then take one optimizer step.

for microbatch in range(n_microbatches):
forward(mikrobatch)
backward()           # LÄGGER TILL i grad-buffrarna, skriver inte över
scale_grads(1.0 / n_mikrobatcher)  # genomsnitt över mikrobatcher
optimizer_step()
zero_grads()             # nollställ för nästa träningssteg

Backward-kärnor använder +=-semantik, inte =. Varje anrop lägger till gradientbidrag till den befintliga buffern; buffern håller den löpande summan tills zero_grads() rensar den.

Optimiser-tillståndet

AdamW (aktivitet 73) håller två ytterligare buffertar per vikt: första momentet m & andra momentet v. Totalt minnesanvändning under träning:

vikter:    1× viktantal
gradienter:  1× viktantal
Adam m:     1× viktantal
Adam v:     1× viktantal
sparade akt: ~2-4× beroende på lager & batch
──────────────────────────────────────────
totalt:      ~6-8× viktantal

ANDREA-120M vid FP16: ~240 MB × 4 buffertar (vikt, grad, m, v) + ~5-10 GB aktiveringar = ~10-12 GB totalt. Bekvämt under RTX 4090:s 24 GB tak. ANDREA-12M tränad i 1,4 GB; 10× parametertskalning ger ~10× minne.

Dimensionering av gradientbuffertar

ANDREA-120M har ~120 000 000 vikter & använder gradientackumulation över 4 mikrobatcher per träningssteg. Beräkna: (a) gradientbuffertstorlek i MB vid FP16; (b) total minne för vikter + gradienter + Adam m + Adam v vid FP16; (c) hur många separata `forward()` + `backward()`-anrop som utförs per träningssteg. Visa din aritmetik.

Full kontroll över minne & precision

Vad generiska ramverk kostar

PyTorch & JAX gör autograd bekvämt: skriv Python-kod, få gradienter automatiskt. Kostnaden: ett generiskt dispatch-lager mellan din kod & CUDA. Varje operation går genom Python-tolkarens overhead, ramverksbokföring, & dynamisk kärnval. För träning av en liten språkmodell på en GPU spelar den overheaden roll.

Konkreta kostnader som ANDREA undviker:

1. Python-tolkens latens. Varje PyTorch-operation korsar gränsen mellan Python/C++. För ~100 kärnstart per träningssteg vid ~9 steg/min innebär det ~900 gränskorsningar per minut. Dispatch på C-nivå eliminerar detta.

2. Ramsallokatorns oförutsägbarhet. PyTorch:s cachande allokator ger god genomströmning i genomsnitt men oförutsägbar toppminne. ANDREA:s träningsmotor förallokerar varje buffert vid uppstart; ingen omallokering under träning, ingen fragmentering, inga överraskande OOM vid steg 100K.

3. Generisk kärnval. PyTorch väljer kärnor vid körning via heuristiker. ANDREA väljer kärnor vid kompileringstid, optimerade för RTX 4090 tensor core-plattor.

4. Blandad precision-rörledning. ANDREA-120M:s FP16 cuBLAS-sökväg & ANDREA:s FP8 E4M3 tensor core-experiment kräver exakt kontroll över vilka tensorer som lever vid vilken precision. Generiska ramverk exponerar denna kontroll genom lagerbaserade API:er; anpassad CUDA skriver det direkt.

Kompromissen

Anpassad CUDA kostar: mer kod att skriva, fler buggar att hitta, ingen community-ekosystem. ANDREA:s microgpt_cuda.cu är ~6000 rader handskriven CUDA som tog månader att felsöka. Varje ny operation kräver att skriva en forward-kärna, en backward-kärna & tester.

Vad ANDREA vinner:

- Fullständig reproducerbarhet. Tränings-pipelinen är en C-binär plus en Python-proxy. Ingen versionsdrift över PyTorch-släpp, inga CUDA-versionsmismatch med ramverks-wheels.

- Bit-exakta repriser. SIGTERM utlöser en checkpoint-skrivning som fångar varje tensor exakt som GPU:n ser den. Repris tar upp samma förlustbana som körningen var på.

- Förutsägbar minne. ANDREA-120M tränad i 200K steg utan OOMs. Minnesanvändning redovisades vid motorns start.

- Direkt hårdvaruåtkomst. Tensor core-tile-storlekar, FP8 E4M3-inställningar, asynkrona minneskopier: allt direkt adresserbart i CUDA, ogenomskinligt i generiska ramverk.

Reproducerbarhet som uppdrag

ANDREA-whitepaper-sektion 9 listar den fullständiga reproducerbarhetsstacken:

Träningsmotor: microgpt/microgpt_cuda.cu
Träningsproxy: microgpt/training_proxy.py
Experimentkonfigurationer: experiments/ANDREA-*-TRAIN.json
Datapipeline: scripts/pull-hermes3.py, scripts/prep-megachat.py
Dashboard: scripts/live-loss-dashboard.html
Banditspecifikation: docs/FIREHOSE-BANDIT.md
Modellokumentation: docs/ANDREA.md

Hårdvarukrav: en NVIDIA GPU med ≥8 GB VRAM (RTX 3060 eller bättre). Vem som helst kan reproducera ANDREA-12M från dessa artefakter. Den anpassade CUDA-sökvägen är en del av varför: inga versionsfrysningar av ramverk, inga beroendeöverraskningar om fem år.

Signaler & Kontrollpunkter

CUDA-träningsloopen svarar på två POSIX-signaler:

- SIGTERM: skriv en omedelbar kontrollpunkt, sedan avsluta. Används vid rent stopp av träning.

- SIGUSR1: skriv en omedelbar kontrollpunkt, fortsätt träningen. Används under polish-pivoten i v3 för att fånga tillstånd utan att avbryta körningen.

Kontrollpunktformat: [int32 step][int32 n_params][n_params × float32 weights][n_params × float32 m][n_params × float32 v]. Stegräknare, viktantal, sedan vikter följt av Adam-momenter. Återupptar bit-exakt. Proxyn arkiverar .samples.json & .state.json separat vid polish; .loss.json arkiveras aldrig (den ackumulerar hela träningshistoriken).

Varför inte PyTorch

ANDREA kunde ha använt PyTorchs autograd istället för att skriva `microgpt_cuda.cu` för hand. Ge två distinkta ingenjörsmässiga skäl till varför ANDREA valde anpassad CUDA. Ett skäl ska referera till minne eller precisionkontroll; det andra ska referera till reproducerbarhet, ramverksberoenden eller långsiktig underhåll.