DeepSpeed ZeRO-stadia: welke trainingsstatus wordt verdeeld?
Automatische vertaling
Dit artikel is automatisch vertaald vanuit de oorspronkelijke Engelse versie.
ZeRO-1 verdeelt de optimizerstatus over dataparallelle GPUs, ZeRO-2 verdeelt ook de gradiënten en ZeRO-3 ook de modelparameters. Test het laagste stadium waarin de werklast in het geheugen past met een acceptabele throughput. Een hoger stadium bespaart meer geheugen voor de modelstatus, maar kan meer communicatie vereisen.
Deze stadia veranderen waar de trainingsstatus wordt opgeslagen. Ze verminderen het geheugengebruik van activations niet automatisch.
Een geheugenvoorbeeld met expliciete aannames
Stel dat de training FP16-weights en -gradiënten gebruikt, plus FP32-master-weights voor Adam en twee optimizermomenten. Dit vereist per parameter 2 bytes voor weights, 2 voor gradiënten en 12 voor de optimizerstatus: 16 bytes in totaal.
Voor 7.5 miljard parameters op acht GPUs levert de berekening voor de modelstatus het volgende op:
| Stadium | Status verdeeld over GPUs | GB modelstatus per GPU |
|---|---|---|
| Gewoon dataparallelisme | Geen | 120 |
| ZeRO-1 | Optimizerstatus | 41.25 |
| ZeRO-2 | Optimizerstatus en gradiënten | 28.125 |
| ZeRO-3 | Alle drie de componenten | 15 |
Dit zijn decimale GB. Activations, tijdelijk verzamelde parameters, extra geheugenverbruik van de allocator en runtimebuffers zijn niet meegerekend. Een trainingsomgeving die andere opslagprecisies gebruikt, vereist een andere berekening.
Het ZeRO-artikel legt de verdeling van de status uit en analyseert de communicatie. ZeRO-3 reconstrueert benodigde parameters met all-gather-operaties; de analyse van volledige verdeling in het artikel komt uit op ongeveer 1.5 keer het communicatievolume van standaarddataparallelisme. Die verhouding geldt onder de aannames van de analyse en is geen garantie voor elk netwerk en elke configuratie.
Kies op basis van het geheugenprofiel
Als de parameters in het geheugen passen en de optimizerstatus de fout veroorzaakt, test dan stadium 1 of 2. Als volledig gerepliceerde parameters niet passen, test dan stadium 3. Als activations het meeste geheugen gebruiken, combineer het gekozen stadium dan met activation checkpointing of een kleinere trainingsbatch.
DeepSpeed ondersteunt ook het verplaatsen van status naar CPU en NVMe. Bekijk de huidige ZeRO-configuratie en meet de overdrachtskosten op de daadwerkelijke host en opslag. De extra capaciteit kan helpen, maar trage overdrachten kunnen de throughput verlagen.
Het ZeRO-onderdeel van de LLM Engineering Guide biedt een uitvoerbare berekening voor andere aantallen parameters en GPU-groepen.