17 · Entrenament distribuït i eficiència
Idea central: cap model modern cap en una GPU. L'art de l'entrenament a escala consisteix a repartir pesos, gradients, estats i càlcul entre milers de dispositius amagant la comunicació darrere del càlcul.
1. Per què no hi cap# sense fórmules
Memòria per paràmetre entrenant amb AdamW en precisió mixta:
| Component | Bytes/param |
|---|---|
| Pesos bf16 | 2 |
| Gradients bf16 | 2 |
| Còpia mestra fp32 | 4 |
| Moment 1 (fp32) | 4 |
| Moment 2 (fp32) | 4 |
| Total | ~16 |
Un model de 70B necessita ~1,1 TB només per a estats, sense comptar activacions. Una H100 en té 80 GB.
Per què no hi cap
Quants bytes demana cada paràmetre per entrenar-lo, i quantes GPUs calen només per guardar-los?
Dos moments en fp32: l’estàndard per a LLMs i, sovint, la meitat de tota la memòria.
- bytes per paràmetre
- 16,0
- estats
- 1.120 GB
- GPUs de 80 GB
- 14
- contra AdamW
- 1,00×
| Estat | precisió mixta | aquest projecte, comptat | d'on |
|---|---|---|---|
| Pesos | 2 B · bf16 | 8 B | θ, Float64Array |
| Gradients | 2 B · bf16 | 8 B | grad de batchLossAndGrad |
| Còpia mestra | 4 B · fp32 | 0 B | no cal: θ ja és fp64 |
| Moment 1 | 4 B · fp32 | 8 B | opt.m, Float64Array |
| Moment 2 | 4 B · fp32 | 8 B | opt.v, Float64Array |
| Total | 16 B | 32 B | 922.624 bytes per a 28.832 paràmetres |
Els quatre arrays que un pas d'aquest projecte deixa vius fan exactament 8 bytes per escalar —Float64Array— i la columna surt a 32: el doble de la precisió mixta. La precisió mixta guanya fent els pesos i els gradients en bf16, però paga una còpia mestra en fp32 perquè les actualitzacions petites d'Adam no desapareguin en arrodonir; aquí no cal, perquè θ ja és de precisió doble. I la columna d'aquí no és una fórmula: si l'optimitzador guardés un moment de més, el número de sota ho diria.
I això abans d'una sola activació. Un 70B amb AdamW en precisió mixta són 1.120 GB —els «~1,1 TB» del text— i en una H100 n'hi caben 80. Canviar d'optimitzador ajuda: Adam de 8 bits, 10,0 bytes per paràmetre; Adafactor, 8,1 bytes per paràmetre; Lion, 12,0 bytes per paràmetre. Però cap optimitzador no baixa dels 4 bytes de pesos i gradients, i cap model de frontera no hi cap sencer en una GPU. Per això calen les cinc dimensions del §2: no per anar més de pressa, sinó per poder entrenar.
2. Les cinc dimensions de paral·lelisme# sense fórmules
2.1 Data Parallelism (DP)
Cada GPU té una còpia del model i processa un tros del batch; es sincronitzen els gradients amb un all-reduce. Simple, escala bé, però no resol el problema de memòria.
2.2 ZeRO / FSDP (sharding d'estats)
En comptes de replicar, reparteix:
- ZeRO-1: reparteix els estats de l'optimitzador (4× menys memòria).
- ZeRO-2: + gradients.
- ZeRO-3 / FSDP: + pesos. Cada GPU guarda només un tros; abans de cada capa fa un all-gather del tros que li falta, calcula, i el descarta.
Compromís: memòria per comunicació. FSDP és avui l'opció per defecte en PyTorch per a entrenament i fine-tuning grans.
2.3 Tensor Parallelism (TP)
Reparteix una mateixa matriu entre GPUs: , cada GPU calcula la seva part i es fa un all-reduce per capa. Requereix interconnexió molt ràpida (NVLink) → normalment només dins d'un node (8 GPUs).
2.4 Pipeline Parallelism (PP)
Reparteix capes entre grups de GPUs. El problema és la "bombolla" (GPUs esperant); es mitiga amb micro-batches i planificacions intel·ligents (1F1B, interleaved, zero-bubble).
2.5 Context / Sequence Parallelism (CP/SP) i Expert Parallelism (EP)
- CP/SP: reparteix la seqüència entre GPUs; imprescindible per a context llarg (Ring Attention i variants).
- EP: reparteix els experts d'un MoE; requereix all-to-all per capa.
La combinació ("paral·lelisme 4D/5D")
Els entrenaments reals combinen tot: TP dins del node, PP entre nodes, DP/FSDP per damunt, més CP i EP segons el model. Trobar la configuració òptima és un problema d'optimització en si mateix (i hi ha eines que el resolen automàticament).
Les cinc dimensions de paral·lelisme
Repartir la feina entre milers de GPUs: què baixa la memòria de cada una, què costa en comunicació, i on es perd el temps?
La teva configuració: 512 GPUs = DP 16 × TP 8 × PP 4 × CP 1
11,7 GB per GPU contra 80 GB: hi cap. Comunicació per pas i GPU: DP 8,2 · TP 150 · PP 0,5 · CP 0,0 GB.
El pipeline simulat: 4 etapes, 8 microlots, 1F1B
Temps → (forward = 1, backward = 2). Blau: forward; magenta: backward; gris: la bombolla.
- GPUs
- 512
- memòria per GPU
- 11,7 GB
- comunicació per pas
- 159 GB
- bombolla
- 27,3 %
La taula de sota compara les estratègies amb el mateix model i 16 dispositius. Cap xifra no està escrita a mà: surten de deviceMemory i commVolume, i el test del projecte comprova les dues igualtats de la franja —DP igual que sense, i ZeRO-3 igual a total/n— amb toBe, no amb aproximacions.
| Estratègia | estats / GPU | activacions / GPU | comunicació / pas |
|---|---|---|---|
| Cap paral·lelisme | 1.120 GB | 10,7 GB | — |
| DP × 16 | 1.120 GB | 10,7 GB | 263 GB |
| ZeRO-1 × 16 | 333 GB | 10,7 GB | 263 GB |
| ZeRO-2 × 16 | 201 GB | 10,7 GB | 263 GB |
| ZeRO-3 × 16 | 70,0 GB | 10,7 GB | 394 GB |
| TP × 8 | 140 GB | 1,3 GB | 601 GB |
| PP × 16 | 70,0 GB | 5,4 GB | 4,3 GB |
ZeRO-1 reparteix la còpia mestra i els moments —12 dels 16 bytes— i tendeix a una reducció de 4× quan n creix: aquí, 3,37×. ZeRO-2 hi afegeix els gradients i ZeRO-3, els pesos, i el preu és a l'última columna: ZeRO-3 comunica una vegada i mitja el que DP, perquè abans de cada capa ha de reunir els pesos que no té, al forward i al backward. És el «memòria per comunicació» del text, amb els dos números.
Amb TP, cada capa fa quatre all-reduce de les activacions —dos al forward, dos al backward— de 0,1 GB cadascun. Són 0,94 GB per capa i per GPU, i s'han de moure dins de cada capa, entre un producte de matrius i el següent: no es poden amagar darrere de res que no depengui d'ells. NVLink 4 (dins del node): 450 GB/s; InfiniBand NDR 400 Gb/s (entre nodes): 50 GB/s. Per NVLink la comunicació és 14 % del càlcul de la capa (a un 40 % del pic de la H100 80 GB); per InfiniBand, 1,3 vegades. Per això TP es queda a les 8 GPUs d'un node, i entre nodes s'hi posa PP o DP, que comuniquen molt menys sovint.
El diagrama de sota no és la fórmula dibuixada: és la simulació d'un pas, operació a operació. Cada etapa fa els forwards i els backwards en l'ordre de la planificació, i cada operació espera la que li cal de l'etapa del costat. Les caselles buides són la bombolla, i sumades donen exactament el que diu la fórmula —el test ho comprova per a 256 combinacions—. Puja m i la bombolla baixa; per això els microlots.
1F1B no redueix la bombolla: redueix la memòria. Amb GPipe cada etapa guarda les activacions dels 8 microlots alhora; amb 1F1B, com a molt 4, 3, 2, 1 per etapa —comptat a la simulació—, perquè comença els backwards abans d'acabar tots els forwards.
3. Estalviar memòria# sense fórmules
- Activation checkpointing: no guardis les activacions intermèdies; recalcula-les al backward. ~30% més de temps, fins a 10× menys memòria d'activacions. Es pot fer selectivament (només les capes cares de recalcular es guarden).
- CPU/NVMe offloading: desplaça estats a RAM o disc. Salva execucions impossibles, però és lent.
- Precisió baixa: bf16 → fp8 per als GEMMs, amb acumulació en precisió alta.
- Optimitzadors eficients en memòria: Adafactor, 8-bit Adam, LoRA (que elimina el 99% dels estats d'optimitzador).
Estalviar memòria: activation checkpointing
Quants números desa de debò un forward perquè el backward funcioni? I què costa no desar-los i tornar-los a calcular?
Escalars desats contra la longitud de seqüència
- sense checkpointing
- amb checkpointing: el pic
- pesos d’atenció, H·L² per capa
Model de 53.536 paràmetres. Els pendents de les rectes en log-log són 1 per a la part lineal i 2 per a la quadràtica: on la línia vermella talla la blava, l'atenció ja pesa més que tota la resta.
- sense checkpointing
- 609,4 k escalars
- amb checkpointing
- 179,6 k escalars
- reducció
- 3,4×
- temps de més
- 41 %
- pesos d’atenció
- 43 %
Sense checkpointing, cada bloc desa tot això —comptat sobre els arrays del primer bloc, amb L = 128:
| weights | 65.536 | 44 % |
| gate | 11.264 | 8 % |
| up | 11.264 | 8 % |
| act | 11.264 | 8 % |
| n1 | 4.224 | 3 % |
| n2 | 4.224 | 3 % |
| q | 4.096 | 3 % |
| k | 4.096 | 3 % |
| qr | 4.096 | 3 % |
| kr | 4.096 | 3 % |
| Total del bloc | 148.736 |
Amb checkpointing es guarda només l'entrada de cada bloc —4 × 128 × 32 = 16.384 escalars— i, al backward, cada bloc es torna a calcular des d'aquella entrada just abans de fer-ne el gradient. Al pic hi ha els punts de control, un sol bloc recalculat (148.736) i el cap del model (14.464). El text diu «fins a 10× menys»: aquí en surten 3,4, i el número creix amb les capes, perquè sense checkpointing cada capa hi suma un bloc sencer i amb checkpointing només una entrada. Posa-hi 8 capes i mira'l créixer.
El text diu «~30 % més de temps». La xifra surt de la ràtio entre forward i backward: si el backward en fa el doble, recalcular un forward afegeix 1/3 del total, un 33 %. Aquí la ràtio mesurada és 1,42 i el sobrecost, 41 %. A diferència del prefill i el decode del fitxer 15, aquest número sí que es pot cronometrar al navegador: no depèn de cap paral·lelisme, sinó de quanta aritmètica fa cada fase. La ràtio no és el 2 que se sol dir, i la diferència és d'aquest codi: el forward hi fa feina de més —còpies de Q i K abans i després de RoPE, les contribucions de cada sublayer desades per als widgets del fitxer 11— i el backward en surt 1,42 vegades i no 2. Com més pesa el forward, més costa repetir-lo: d'aquí el 41 % en comptes del 33 %.
Tot el que desa el forward creix linealment amb L excepte una cosa: els pesos d'atenció, H·L·L per capa. El gràfic ho ensenya comptat a cinc longituds, i el test del projecte comprova que el terme quadràtic és exactament capes × caps × L². A L = 256 ja són el 60 % del total, i a context llarg dominen.
Aquí es troba amb el fitxer 14: FlashAttention no guarda mai aquesta matriu —la recalcula per blocs al backward, que és checkpointing aplicat només a l'atenció— i per això la memòria d'activacions d'un entrenament de context llarg passa de quadràtica a lineal. El checkpointing selectiu del text és el mateix raonament: guardar el que és barat de guardar i car de recalcular, i recalcular la resta.
4. Eficiència: MFU# sense fórmules
Model FLOPs Utilization = FLOPs útils / FLOPs teòrics del hardware.
- 35-55% és bo per a entrenament dens a gran escala.
- Per sota del 25%, tens un problema (comunicació, bombolles, kernels subòptims, càrrega de dades).
- Amb MoE i context llarg és normal tenir MFU més baixos.
On es perd eficiència: comunicació no solapada, bombolles de pipeline, kernels no fusionats, càrrega de dades, i fallades. En execucions de milers de GPUs durant setmanes, les fallades de hardware són esdeveniments diaris — cal checkpointing freqüent i asíncron, detecció automàtica de nodes defectuosos i reinici ràpid.
MFU: on es perd el temps
Tens un MFU del 18 %. Quin terme se'n menja la resta, i què arreglaries primer?
El temps d'un pas, terme a terme
- Càlcul útil al pic · 18 %
- Kernels per sota del pic · 22 %
- Forward repetit (checkpointing) · 13 %
- Bombolla del pipeline · 23 %
- Comunicació no amagada · 13 %
- Esperant dades · 3 %
- Fallades i reinicis · 7 %
Un pas de Llama 3.1 405B amb lots de 16 M tokens. La franja verda és l'MFU: el temps que el pas trigaria si totes les GPUs fessin només càlcul útil al pic.
Fracció del pic que fan els productes de matrius quan calculen.
Fracció del temps perduda, amortitzada.
- MFU
- 18,0 %
- temps per pas
- 13.327 ms
- ideal
- 2.397 ms
- el terme que domina
- Bombolla del pipeline
Els dos camins han de donar el mateix, i donen el mateix: 6 × 405 G × 16 M tokens són 3,89·10¹⁹ FLOPs útils per pas; dividits pels 13,33 s que triga i pel pic de 16.384 H100 80 GB, surt el 18,0 %. Els FLOPs del forward repetit no hi compten: són feina, però no són útils, i per això el checkpointing baixa l'MFU encara que la GPU vagi a ple.
| Terme | ms | del pas | què s'hi fa |
|---|---|---|---|
| Càlcul útil al pic | 2.397 | 18 % | — |
| Bombolla del pipeline | 3.107 | 23 % | més microlots, planificació intercalada o zero-bubble, menys etapes |
| Kernels per sota del pic | 2.930 | 22 % | kernels fusionats, FlashAttention, fp8, formes de matrius que omplin els tensor cores |
| Forward repetit (checkpointing) | 1.776 | 13 % | checkpointing selectiu: recalcular només el que és barat |
| Comunicació no amagada | 1.680 | 13 % | solapar amb el càlcul, TP dins del node, menys ZeRO-3 |
| Fallades i reinicis | 987 | 7 % | checkpoints més freqüents i asíncrons, detecció de nodes dolents |
| Esperant dades | 450 | 3 % | precàrrega asíncrona, més fils de lectura, dades preprocessades |
La bombolla surt de la fórmula que el widget de paral·lelisme simula: (p−1)/(m+p−1) = 30,4 % del temps del pipeline, que és 0,44 vegades la feina. El text diu que per sota del 25 % «tens un problema», i la taula diu quin: investiga en l'ordre de la columna del mig. Ara mateix, els tres primers són bombolla del pipeline, kernels per sota del pic, forward repetit (checkpointing). Arregla el primer —mou-ne el control— i mira com canvia l'ordre: cada terme que baixa fa que els altres pesin més, i per això a l'exercici 17.5 la resposta no és una llista fixa sinó un mètode.
Una cosa que la taula ensenya i que sorprèn: el checkpointing complet, que fa que el model hi càpiga, costa un 13 % del pas en FLOPs que no són útils.
5. El hardware, en una pàgina# sense fórmules
- GPU: memòria HBM (80-192+ GB), ample de banda (3-8 TB/s), FLOPs per format (bf16, fp8, fp4). L'ample de banda mana en inferència; els FLOPs manen en entrenament.
- Interconnexió: NVLink dins del node (centenars de GB/s per GPU), InfiniBand/RoCE entre nodes. La topologia determina quins paral·lelismes són viables on.
- Alternatives: TPU (pods amb interconnexió pròpia, excel·lents per a entrenament), acceleradors d'inferència especialitzats (Groq, Cerebras) amb latències molt baixes, i el cost energètic com a restricció cada cop més real.
El sostre del maquinari
Per què l'ample de banda mana quan generes i els FLOPs manen quan entrenes? Situa les dues operacions al diagrama.
- sostre de la H100 80 GB
- les altres GPUs del catàleg
- Decode, lot de 1
- Decode, lot de 64
- Entrenament, 8 k tokens
- colze de la GPU
- 296 FLOP/byte
- decode
- 1,0 FLOP/byte
- entrenament
- 3.641 FLOP/byte
- decode: del pic
- 0,3 %
El colze surt d'arithmeticIntensity: 990 TFLOP/s dividit per 3.350 GB/s. Els punts, de matmulIntensity: 2·m·k·n FLOPs d'un producte de matrius contra els bytes de llegir-ne les entrades i escriure'n la sortida.
El decode d'un token per seqüència multiplica un vector per la matriu de pesos: cada pes es llegeix per fer-ne dues operacions —una multiplicació i una suma— i ocupa dos bytes en bf16. Això és 1 FLOP per byte, no els «~2» que es diuen sovint: el 2 és per paràmetre, i en bf16 cada paràmetre són dos bytes. Queda 296 vegades per sota del colze, i la GPU fa com a molt el 0,34 % del que podria: l'ample de banda mana. Un pas d'entrenament amb 8 k tokens reutilitza cada pes 8 k vegades, i queda per sobre: els FLOPs manen.
Entremig hi ha el batching del fitxer 15: amb un lot de B seqüències, la intensitat del decode és gairebé B. Per arribar al colze d'aquesta GPU calen uns 512 tokens alhora. Canvia de GPU: el colze es mou —les GPUs noves guanyen potència més de pressa que ample de banda, i a cada generació és més fàcil quedar-s'hi per sota— i els punts es queden quiets.
6. Comprova que ho has entès#
0 de 5 resoltes
- Per què el data parallelism no resol el problema de memòria i ZeRO-3 sí?
- Per què el tensor parallelism es limita normalment a les GPUs d'un mateix node?
- Què és la bombolla de pipeline i com es redueix? Al widget de paral·lelisme, fes-la baixar per sota del 10 %.
Objectiu:Al widget de les cinc dimensions, amb almenys 2 etapes de pipeline, puja els microlots fins que la bombolla sigui menor del 10 %.
Vés al widget ↑pp = per defecte · m = per defecte - Calcula la memòria d'estats d'un model de 13B entrenat amb AdamW en precisió mixta, sense activacions. En GB.GBMarge acceptat: 2 %
- Tens un MFU del 18 %. Quines tres coses investigues primer, i com?
7. Recursos#
The Ultra-Scale Playbook: Training LLMs on GPU Clusters (Hugging Face) — el millor recurs obert que existeix sobre això.
Rajbhandari et al. (2019), ZeRO.
Shoeybi et al. (2019), Megatron-LM (tensor parallelism).
Narayanan et al. (2021), Efficient Large-Scale Language Model Training on GPU Clusters.
Llama 3 Herd of Models — secció d'infraestructura i fallades reals.