Distribuovaný trénink

Data, tensor a pipeline paralelismus, ZeRO/FSDP, komunikace.

Co se naučíš: Budeš rozumět tomu, jak se model rozdělí mezi tisíce karet a co je na tom drahé.

7 min čtení + cvičeníNavazuje na:🎚️ Optimizer a learning rate

Model 70 B potřebuje na trénink stovky gigabajtů paměti, do jedné karty se nevejde ani náhodou. Musí se rozdělit mezi desítky až tisíce GPU, a způsob rozdělení rozhoduje o tom, jestli trénink poběží efektivně, nebo se udusí komunikací. Tahle kapitola vysvětlí tři základní osy paralelismu a proč se v praxi kombinují.


Data parallel: nejjednodušší osa

Každá GPU má celou kopii modelu a dostane jinou část dávky. Po spočítání gradientů se gradienty všech karet sečtou (operace all-reduce) a všechny provedou stejný update.

GPU 0: model | data A ─┐
GPU 1: model | data B ─┼─→ all-reduce gradientů → všechny mají stejný update
GPU 2: model | data C ─┘
  • Plus: triviální, škáluje propustnost přímo úměrně.
  • Minus: nepomůže s pamětí, model, gradienty i stavy optimizeru jsou na každé kartě celé.

ZeRO / FSDP: data parallel, který šetří paměť

Nápad: když má každá karta stejnou kopii, je to plýtvání. Rozděl ji.

stupeň 1: rozděl mezi karty stavy optimizeru      → úspora ~4×
stupeň 2: + gradienty                             → úspora ~8×
stupeň 3: + samotné váhy                          → úspora úměrná počtu karet

Ve stupni 3 (FSDP) má každá karta jen svůj kousek vah. Když je potřeba spočítat vrstvu, karty si chybějící části na chvíli vymění, spočítají a zase zahodí. Platí se za to komunikací, ale umožní to trénovat modely, které by se jinak nevešly nikam.

Tohle je dnes výchozí volba pro většinu tréninků, které nejsou úplně extrémní.


Tensor parallel: rozříznout jednu vrstvu

Velké matice se rozdělí mezi karty po sloupcích nebo řádcích; každá spočítá svou část a výsledky se posbírají.

matice 4096 × 14336  →  GPU0: sloupce 1–7168 | GPU1: sloupce 7169–14336
  • Plus: zmenší paměť i výpočet na jednu vrstvu.
  • Minus: vyžaduje komunikaci uvnitř každé vrstvy, tedy mnohokrát za krok. Proto se používá jen v rámci jednoho serveru, kde jsou karty propojené rychlým NVLinkem, ne přes síť.

Pipeline parallel: rozdělit vrstvy

Karta 0 dostane vrstvy 1–8, karta 1 vrstvy 9–16 a tak dál. Data protékají jako po výrobní lince.

Problém jsou bubliny: než první dávka dorazí na poslední kartu, ostatní čekají. Řeší se rozsekáním dávky na mikrodávky, které se překrývají, ale nějaká režie zůstane vždycky.


3D paralelismus: kombinace všeho

Velké tréninky používají všechny tři osy najednou:

uvnitř serveru (8 GPU, NVLink):     tensor parallel
mezi servery v rámci skupiny:       pipeline parallel
mezi skupinami:                     data parallel + ZeRO

U MoE k tomu přibývá expert parallel, experti se rozdělí mezi karty a tokeny se k nim posílají podle rozhodnutí routeru.


Proč je komunikace to hlavní

NVLink uvnitř serveru:   stovky GB/s
síť mezi servery:        desítky GB/s (InfiniBand / RoCE)

Rozdíl je řádový, a proto platí základní pravidlo návrhu: co nejvíc komunikace nacpat dovnitř serveru a přes síť posílat co nejméně a co nejvzácněji. Špatně zvolené rozdělení umí snížit využití GPU z 50 % na 15 % a trénink pak stojí trojnásobek.


Co ještě běží na pozadí

  • Checkpointing. Stav se pravidelně ukládá, protože při tisících karet něco selže prakticky denně. Restart z checkpointu je běžná provozní událost, ne mimořádnost.
  • Deterministické pořadí dat, aby se trénink dal po restartu navázat.
  • Sledování: loss, norma gradientu, využití karet, propustnost tokenů. Loss spike se pozná během minut a řeší se návratem zpět.

Co z toho plyne pro tebe

Pokud netrénuješ základový model, nebudeš tohle nastavovat, ale vysvětluje ti to:

  • proč je pretraining doménou velkých hráčů (jde o infrastrukturní projekt, ne o skript),
  • proč je fine-tuning přes LoRA/QLoRA tak populární (vejde se na jednu kartu a nepotřebuje nic z výše uvedeného),
  • proč jsou velké modely drahé i na inferenci, když se musí rozdělit přes víc karet.

Cvičení

  1. Model se vejde na jednu kartu, ale trénink je pomalý. Který druh paralelismu nasadíš jako první a proč zrovna ten?
  2. Proč data parallel nesníží paměťové nároky na jednu kartu, zatímco ZeRO ano?
Náčrt řešení: rozbal, až si cvičení zkusíš sám
  1. Data parallel, protože je nejjednodušší a nejlevnější na komunikaci. Každá karta dostane vlastní dávku, spočítá gradienty a ty se zprůměrují. Komunikuje se jednou za krok a jen gradienty, takže se to dobře škáluje. Tensor a pipeline paralelismus se nasazují až tehdy, když se model na jednu kartu nevejde, protože oba přidávají komunikaci uvnitř kroku.
  2. Protože data parallel drží na každé kartě kompletní kopii vah, gradientů i stavů optimizeru. Zrychlí tedy průchod dat, ale paměťově nepomůže vůbec. ZeRO a FSDP tyhle tři věci rozdělí mezi karty a dotahují si jen to, co zrovna potřebují, čímž vymění část paměti za komunikaci. Odtud plyne, že ZeRO má smysl teprve tam, kde je propojení karet dost rychlé.

Shrnutí

  • Data parallel škáluje propustnost, ale nešetří paměť; ZeRO/FSDP ji rozdělí mezi karty.
  • Tensor parallel dělí jednotlivé vrstvy a vyžaduje rychlé propojení, používá se uvnitř serveru.
  • Pipeline parallel dělí vrstvy mezi uzly a bojuje s bublinami.
  • Rozhodující je komunikace: co nejvíc uvnitř serveru, co nejméně po síti.
Proč samotný data parallel nepomůže s modelem, který se nevejde do paměti karty?

Protože každá karta drží celou kopii modelu, gradientů i stavů optimizeru, dělí se jen data. Paměťově pomůže až ZeRO/FSDP, které tyto části rozdělí mezi karty a chybějící kousky si dočasně vyměňují.

Proč se tensor parallel používá jen uvnitř jednoho serveru?

Protože vyžaduje komunikaci uvnitř každé vrstvy, tedy mnohokrát během jednoho kroku. To dává smysl jen po rychlém propojení typu NVLink; přes běžnou síť by komunikace zabrala víc času než samotný výpočet.

Co jsou bubliny u pipeline paralelismu?

Prostoje, kdy karty čekají, než k nim dorazí data z předchozího stupně linky. Zmírňují se rozsekáním dávky na mikrodávky, které se překrývají, ale úplně odstranit se nedají.