ZeRO och Sharded Optimizers
ZeRO (Zero Redundancy Optimizer) eliminerar den slösaktiga minnesdupliceringen av dataparallellism genom att splittra optimerartillstånd, gradienter och vikter över GPU:er.
Översikt
It lets you train enormous models with the simplicity of data parallelism but a fraction of the per-GPU memory.
Djupdykning
I vanlig dataparallellism lagrar varje GPU en redundant fullständig kopia av optimerartillståndet, gradienter och parametrar, vilket är enormt slöseri, särskilt för Adam, där optimerartillståndet kan vara flera gånger storleken på själva modellen. ZeRO, introducerad av Microsoft i DeepSpeed, tar bort denna redundans genom att partitionera dessa tensorer över GPU:er så att varje enhet bara äger en del. ZeRO kommer i tre progressiva steg: Steg 1 skärvor optimizer tillstånd, Steg 2 lägger till gradient skärning, och Steg 3 skär själva parametrarna. Vid behov samlar GPU:er de saknade skivorna via kommunikation, beräknar och släpper dem sedan. Resultatet är dramatiskt lägre minne per GPU, vilket möjliggör träning av miljarder till biljoner parametrar, samtidigt som den enkla programmeringsmodellen för dataparallellism behålls.
Teknisk insikt
ZeRO byter ut extra kommunikation för minnesbesparingar. I steg 3, innan ett lagers framåtpassning, samlar en all-gather in lagrets fullständiga parametrar på varje GPU; efteråt kasseras de icke-ägda skivorna för att återta minnet. Gradienter är reducerade spridda så att varje GPU behåller endast gradientdelen som matchar parametrarna den äger. PyTorchs FSDP (Fully Sharded Data Parallel) implementerar samma idé inbyggt och lindar moduler för att skära och skära om i farten.
Strategisk inverkan
Cost and budget
Arkitekturbeslut driver prestanda och driftskostnader i flera år.
Clearer decisions
Teknisk utbildning hjälper team att välja rätt stack, inte bara den nyaste.
Quality control
Bättre tekniska val minskar tillförlitlighetsincidenter i produktionen.
Framtiden för ZeRO och Sharded Optimizers
Sharding håller på att bli standard för storskalig träning snarare än ett exotiskt alternativ. Förvänta dig djupare integrering med avlastning (skjuta skivor till CPU eller NVMe via ZeRO-Infinity), bättre överlappning av all-gather och reducer-scatter med beräkningar för att dölja deras kostnad, och kombinationer med tensor och pipeline-parallellism. När modellerna fortsätter att växa, är minneseffektiva sharded-optimerare centrala för att anpassa dem till realistiska hårdvarubudgetar.
Real-World Implementation
Använder DeepSpeed ZeRO Steg 2 för att finjustera en språkmodell med flera miljarder parametrar som annars skulle svämma över GPU-minnet.
Träning med PyTorch FSDP, som skär parametrar, gradienter och optimerartillstånd över GPU:er och samlar dem per lager på begäran.
Tillämpning av ZeRO-Offload för att skjuta optimeringstillstånd till CPU-minne, vilket låter en enda GPU träna en modell många gånger större än dess VRAM.
Skala en biljonparametermodell med ZeRO-Infinity genom att strömma parameterskärvor från NVMe-lagring när GPU och CPU-minne tar slut.
Risker & skyddsräcken
Att optimera ett riktmärke kan dölja bredare systemsvagheter.
Infrastruktur- och underhållskostnader underskattas ofta.
Säkerhets- och observerbarhetsluckor kan växa i takt med att systemen blir mer komplexa.
Färdplan för genomförande
Definiera latens-, kvalitet- och kostnadsmål före implementering.
Benchmark under realistiska belastnings- och dataförhållanden.
Instrumentövervakning för fel, drift och användarpåverkan.
Förbered återställnings- och incidentsvarsvägar innan skalning.
Fortsätt utforska
Free newsletter
Get the daily AI briefing
Three verified AI stories every weekday morning, written in plain English. Free forever, no ads.
One email each weekday. Unsubscribe in one click. We never sell or share your address.
Test yourself
Take the ZeRO and Sharded Optimizers quiz
Instant feedback on every answer, and a shareable certificate with a verifiable ID once you pass a course.
Support free AI education. AI Understanding is a 501(c)(3) nonprofit — no ads, no paywall, ever. Make a donation
Next guide
Lookahead och Lion Optimizers
Frequently asked questions
What is ZeRO and Sharded Optimizers?
ZeRO (Zero Redundancy Optimizer) eliminerar den slösaktiga minnesdupliceringen av dataparallellism genom att splittra optimerartillstånd, gradienter och vikter över GPU:er. Den låter dig träna enorma modeller med enkelheten av dataparallellism men en bråkdel av minnet per GPU.
Vilken redundans eliminerar ZeRO jämfört med vanlig dataparallellism?
Standarddataparallellism lagrar en fullständig kopia av optimerartillstånd, gradienter och vikter på varje GPU; ZeRO skär dessa så att varje GPU bara rymmer en bit.
Varför är optimerartillstånd ofta det största minnessvinet med Adam?
Adam upprätthåller löpande uppskattningar som första och andra moment per parameter, vilket i kombination med fp32 mastervikter kan dvärga modellens egen storlek.
Vad gör ZeRO Stage 3 som inte steg 1 och 2 gör?
Steg 1 skär optimeringsläge för skärvor, steg 2 lägger till gradienter och steg 3 går vidare genom att även skärpa modellparametrarna över GPU:er.
I ZeRO Steg 3, hur får en GPU alla parametrar den behöver för ett lagers framåtpassning?
Innan ett lager beräknas, sätter en all-gather ihop sina fullständiga parametrar på varje GPU; när de är klara frigörs de icke-ägda skivorna för att återta minnet.
Vilken PyTorch-funktion implementerar sönderdelning i ZeRO-stil?
PyTorchs Fully Sharded Data Parallel (FSDP) skär parametrar, gradienter och optimerartillstånd, samlar in och delar om dem i farten, vilket speglar ZeRO.