07 · Entrenament: optimitzadors, normalització i regularització

Es calcula el gradient dels 4 microlots de 8 exemples sumant-los en un búfer, i després el del lot sencer de 32 exemples d'una tirada. La diferència màxima entre les dues còpies del gradient, paràmetre a paràmetre, és 2,78·10⁻¹⁷.

Coincideixen fins a l'últim bit que la coma flotant permet. Aquesta és tota la justificació de l'acumulació: un batch efectiu de milions de tokens no requereix una GPU que hi càpiga, requereix paciència. El que sí que canvia és el clipping, que s'ha d'aplicar una sola vegada sobre el gradient ja acumulat: si el fessis microlot a microlot, ja no seria el mateix gradient.

El desglossament d'AdamW amb aquesta configuració és 16,0 GB de pesos + 32,0 de còpia mestra + 16,0 de gradients + 64,0 d'estats = 128,0 GB, que són 16,0 bytes per paràmetre.

Posa 13 B de paràmetres amb bf16 i còpia mestra i surten exactament 208 GB: és la mateixa xifra de l'exercici del fitxer 02, i ara es veu d'on surt cada terme. Canviar AdamW per Adam de 8 bits en treu 48,0 GB sense tocar res més.

Res d'això no inclou les activacions de debò, que depenen del batch, de la longitud de seqüència i de si fas gradient checkpointing; el control de marge n'és una aproximació grollera. Tampoc no inclou la fragmentació ni els búfers de comunicació, que en entrenament distribuït no són menyspreables.

1. RMSNorm contra LayerNorm: la diferència, mesurada

Sobre aquest mateix lot, la diferència màxima component a component entre les dues és 1,0057. Si primer es centren les files a mà, passa a 2,22·10⁻¹⁶: són la mateixa operació tret del centrat. Com que les activacions d'un Transformer ja tenen la mitjana prou a prop de zero, treure el centrat gairebé no canvia el resultat i estalvia una passada sobre el vector, el paràmetre $\beta$ i un parell d'operacions per element. D'aquí el 10–15 % de velocitat.

2. Dependència del lot: per què BatchNorm va quedar fora

Es deixa el primer token igual i es canvia la resta del lot. La sortida d'aquell token hauria de ser la mateixa, i amb BatchNorm no ho és:

Zero exacte, no "petit". Aquesta és la propietat que necessites en una seqüència: la sortida d'un token no pot dependre de quins altres tokens hi hagi al lot, perquè en inferència autoregressiva sovint n'hi ha un de sol, perquè les longituds són variables i perquè el mateix text ha de donar el mateix resultat vagi amb qui vagi. BatchNorm ho incompleix per construcció, i per això cal mantenir-hi estadístiques mòbils per a la inferència, amb tot el que això arrossega.

Amb pre-norma, $\mathbf{x}_{l+1} = \mathbf{x}_l + F(\text{Norm}(\mathbf{x}_l))$: el flux residual no es toca mai, i cada bloc hi suma una contribució de mida controlada i gairebé independent de les anteriors. Sumar N vectors independents de mida semblant dona una norma que creix com $\sqrt{N}$, és a dir, exponent 0,5.

Amb post-norma, $\mathbf{x}_{l+1} = \text{Norm}(\mathbf{x}_l + F(\mathbf{x}_l))$: la norma del flux es torna a fixar a cada bloc, de manera que no creix gens i l'exponent ha de ser zero.

La diferència pràctica no és estètica. Amb post-norma, el gradient que torna cap a les primeres capes ha de travessar una normalització a cada bloc, i cadascuna el reescala: en profunditat, això és inestable i obliga a un warmup llarg. Amb pre-norma, el camí residual queda net d'extrem a extrem —és exactament la identitat del widget de propagació de gradients del fitxer 06— i l'entrenament aguanta desenes de blocs sense escalfament especial. El preu és aquest creixement de 6,51×, que es compensa escalant la inicialització de les capes de sortida dels blocs per $1/\sqrt{2N}$ i posant una norma final abans de la capa de sortida.

Puja el dropout per sobre de zero i aquestes quatre xifres passaran a dir alguna cosa: ara mateix les dues primeres files són la mateixa corrida.

El weight decay sí que actua en tots dos règims: amb λ = 0,000, la norma dels pesos queda en 1,000× la de la mateixa corrida sense decay (24,73 contra 24,73). Al pas d'AdamW, el decay entra com θ ← θ − ηλθ per separat de l'actualització adaptativa, i la màscara el manté fora dels biaixos i dels paràmetres de norma.

Reconstruint el valor des dels camps —(−1)^s · significand · 2^exponent— surt 0,09997559, i el valor representat és 0,09997559: difereixen en 0.

I per a fp32, on el maquinari sí que sap fer-ho, la descomposició es contrasta contra la que dona DataView.setFloat32: signe 0, exponent 123, mantissa 5033165 — idèntics.

Prova el botó de 10⁵: fp16 desborda i bf16 no se n'assabenta. Tots dos ocupen setze bits; el que canvia és com es reparteixen entre exponent i mantissa, i en entrenament el rang val més que la precisió.

La mediana de ‖Δθ‖/‖θ‖ al llarg dels dos últims terços d'aquesta corrida és 0,003, contra el 10⁻³ que s'acostuma a donar per bo: dins de l'ordre de magnitud sa. La regla no surt de cap teorema —és una heurística que es compleix en una franja molt ampla de models i escales— però té la virtut que no depèn ni de la mida del model ni de la pèrdua, de manera que serveix per comparar corrides que no es podrien comparar per la pèrdua. El màxim al llarg de la corrida és 0,017, i la norma dels pesos ha crescut 1,18×.

Diagnòstic d'aquest cas: Si passa una vegada, el clipping ja ho ha absorbit. Si es repeteix sempre al mateix punt, el problema són les dades, no l’optimització: mira aquell tros del corpus.

Mira els sis panells i després torna al selector. Hi ha casos on la corba de pèrdua sola no els distingeix: "les dues corbes planes" pot ser un learning rate massa baix, un bug al forward o dades impossibles, i el que els separa és la ràtio update/param (el model es mou o no es mou) i la norma del gradient (arriba senyal o no n'arriba). Registrar només la pèrdua és el que fa que un entrenament de tres setmanes s'hagi de repetir.

El que μP vol que quedi invariant amb l'amplada no és el learning rate: és quant es mouen les sortides del model en un pas. Amb Adam, l'actualització de cada pes té mida ≈ η passi el que passi, de manera que el canvi d'un logit és una suma de $n$ termes de mida η i creix com $n$: per mantenir-lo constant, el learning rate de les capes que reben l'amplada s'ha de dividir per $n$. Això és exactament el que fa μP, i es pot mesurar en un sol pas.

La conseqüència pràctica és que, amb la parametrització estàndard, l'LR òptim es desplaça cap avall a mesura que el model creix, i amb μP es queda quiet. Executa l'escombrat per veure-ho: són 64 entrenaments complets, i aquí costen unes dècimes de segon. En un laboratori de debò, cada punt d'aquell gràfic és una corrida, i és per això que poder buscar hiperparàmetres en un model de 40M i fer-los servir en un de 40B val el que val.