Den här bloggen är medförfattare av Sarah Jane Hong CSO, Darryl Barnhart CTO och Ian Thompson VD för Latent Space och Prem Ranga på AWS.
Latent rymd är en dold representation av abstrakta idéer som maskininlärningsmodeller (ML) lär sig. Till exempel är "hund", "blomma" eller "dörr" begrepp eller platser i latent utrymme. På Latent utrymme, vi arbetar på en motor som låter dig manipulera och utforska detta utrymme med både språk och visuella uppmaningar. Latent Space-teamet kommer från två områden som länge har haft lite överlappning: grafik och naturlig språkbehandling (NLP). Traditionellt har modaliteterna för bilder och text hanterats separat, var och en med sin egen historia av komplex, dyr och ömtålig funktionsteknik. NLP-uppgifter som dokumentförståelse eller frågesvar har vanligtvis inte haft något gemensamt med visionuppgifter som scenförståelse eller rendering, och vanligtvis använder vi väldigt olika tillvägagångssätt och modeller för varje uppgift. Men detta förändras snabbt.
Denna sammanslagning av modaliteter i ett enda delat latent utrymme låser upp en ny generation av kreativa och kommersiella applikationer, från spel till dokumentförståelse. Men att låsa upp dessa nya applikationer i en enda modell öppnar upp för nya skalningsutmaningar, vilket framhålls i "The Bitter Lesson" av Richard Sutton, och det spännande arbetet under de senaste åren med skalningslagar. För att göra detta möjligt arbetar Latent Space med spetsforskning för att smälta samman dessa modaliteter i en enda modell, men också för att skala och göra det effektivt. Det är här modellparallellism kommer in.
Amazon SageMakers unika automatiserade modellpartitionering och effektiva pipelining-strategi gjorde att vi kunde använda modellparallellism med liten ingenjörskonst, och vi skalade vår utbildning av modeller utöver 1 miljard parametrar (vi använder p4d.24xlarge A100-instanser), vilket är ett viktigt krav för oss. Dessutom observerade vi att när vi tränade med en träningsuppsättning med 16 noder, åtta GPU med SageMaker-modellens parallellitetsbibliotek, registrerade vi en 38% förbättring i effektivitet jämfört med våra tidigare träningskörningar.
Utmaningar med att träna storskaliga transformatorer
På Latent Space förenar vi språk och vision i transformatormodeller med miljarder parametrar för att stödja "out-of-distribution" användningsfall från en användares fantasi eller som skulle inträffa i den verkliga världen men inte i vår träningsdata. Vi hanterar utmaningarna som är inneboende i skalning till miljarder parametrar och vidare på två olika sätt:
Informationsinhämtningstekniker har länge varit en nyckelkomponent i sökmotorer och QA-uppgifter. Nyligen har spännande framsteg gjorts genom att kombinera klassisk IR-teknik med moderna transformatorer, specifikt för frågesvarsuppgifter där en modell tränas tillsammans med en neural retriever som lär sig att hämta relevanta dokument för att hjälpa till att svara på frågor. För en översikt, se det senaste arbetet från FAIR in Retrieval Augmented Generation: Effektivisering av skapandet av intelligenta bearbetningsmodeller för naturligt språk och Fusion-in-Decoder, Google Brain's RIKEoch Nvidias Neural Retriever för att svara på frågor.
Även om hämtningsförstärkta tekniker hjälper till med kostnader och effektivitet, kan vi fortfarande inte montera modellen på en enda GPU för vår största modell. Det betyder att vi måste använda modellparallellism för att träna upp den. På grund av vår hämtningsarkitekturs natur var det dock en utmaning att utforma vår modelluppdelning på grund av ömsesidiga beroenden mellan hämtade sammanhang över utbildningsingångar. Dessutom, även om vi bestämmer hur vi delar upp vår modell, var införandet av modellparallellism en betydande ingenjörsuppgift för oss att göra manuellt under hela vår forsknings- och utvecklingslivscykel.
SageMakers modellparallellismbibliotek
Modellparallellism är processen att dela upp en modell mellan flera enheter eller noder (som GPU-utrustade instanser) och skapa en effektiv pipeline för att träna modellen över dessa enheter för att maximera GPU-användningen. De modell parallellbibliotek i SageMaker gör modellparallellism mer tillgänglig genom att tillhandahålla automatiserad modelldelning, även kallad automatiserad modellpartitionering och sofistikerad pipelinekörningsplanering. Modelluppdelningsalgoritmerna kan optimera för hastighet eller minnesförbrukning. Biblioteket använder en partitioneringsalgoritm som balanserar minnet, minimerar kommunikationen mellan enheter och optimerar prestanda.
Automatiserad modellpartitionering
För vårt PyTorch-användningsfall kör modellens parallellbibliotek internt ett spårningssteg (i det första träningssteget) som konstruerar modellgrafen och bestämmer tensor- och parameterformerna. Den konstruerar sedan ett träd, som består av de kapslade nn.Module objekt i modellen, samt ytterligare data som samlats in från spårning, såsom mängden lagrad nn.Parameters, och körtid för varje nn.Module.
Biblioteket korsar sedan detta träd från roten och kör en partitioneringsalgoritm som balanserar beräkningsbelastning och minnesanvändning och minimerar kommunikationen mellan instanser. Om flera nn.Modules delar samma nn.Parameter, placeras dessa moduler på samma enhet för att undvika att bibehålla flera versioner av samma parameter. Efter att partitioneringsbeslutet har tagits laddas de tilldelade modulerna och vikterna till sina enheter.
Schemaläggning för pipelinekörning
En annan kärnfunktion i SageMaker distribuerade modell parallellbibliotek är rörledningar, som bestämmer i vilken ordning beräkningar görs och data bearbetas över enheter under modellträning. Pipelining bygger på att dela upp en minibatch i mikrobatcher, som matas in i träningspipelinen en efter en och följer ett körschema som definieras av bibliotekets körtid.
Microbatch-pipelinen säkerställer att alla GPU:er är fullt utnyttjade, vilket är något vi skulle behöva bygga själva, men med modellparallellismbiblioteket sker detta snyggt bakom kulisserna. Slutligen kan vi använda Amazon FSx, vilket är viktigt för att säkerställa att våra läshastigheter är höga med tanke på antalet filer som läses under träningen av en multimodal modell med hämtning.
Utbildningsarkitektur
Följande diagram visar hur vi ställer in vår utbildningsarkitektur. Våra primära mål var att förbättra träningshastigheten och minska kostnaderna. Bild- och språktransformatorerna vi tränar är mycket komplexa, med ett betydande antal lager och vikter inuti, som löper till miljarder parametrar, vilket gör att de inte kan passa in i minnet av en enda nod. Varje nod bär en delmängd av modellen, genom vilken dataflöden och transformationerna delas och kompileras. Vi ställer in 16 p4d.24xlarge instanser var och en med åtta GPU:er som använder följande arkitekturrepresentation:
När vi skalar upp våra modeller är en vanlig trend att ha allt lagrat i nätverkets vikter. Men av praktiska skäl vill vi utöka våra modeller för att lära oss hur man letar efter relevanta sammanhang för att hjälpa till med uppgiften att rendera. Detta gör att vi kan hålla nere våra serveringskostnader utan att kompromissa med bildkvaliteten. Vi använder en stor transformatorbaserad NLP-modell och som nämnts tidigare, observerade vi en 38% ökning av träningseffektiviteten med SageMaker-modellens parallellitetsbibliotek som visas av följande:
- Vi behöver en allreduce för varje beräkning i fallet med tensornivåparallellism. Detta tar O(log_2 n) parallella steg. Det är n maskiner som tar O(n) steg, för O(n log_2 n) totala operationer.
- För pipelineparallellism kräver vi O(1) parallella steg för att skicka data ner i pipelinen
- Givet 16 maskiner med åtta GPU:er har vi O(1)-kostnad för pipeline-parallell och O(log_2(8)) = O(3)-kostnad för djupgående modellparallell.
- I det här fallet ser vi att nätverkskostnaden reduceras till 1/3 genom att byta till pipeline parallell med vad vi använder med SageMaker modell parallellism, och den totala utbildningskostnaden minskar till 1/2 + 1/2 * 1/log_2(16) = 0.625 av den ursprungliga kostnaden vilket leder till en motsvarande effektivitetsförbättring.
I allmänhet, när behovet motiverar distribuerad utbildning (problem med skalning av modellstorlek eller utbildningsdata), kan vi följa en uppsättning bästa praxis för att avgöra vilken metod som fungerar bäst.
Bästa metoder för distribuerad utbildning
Baserat på vår erfarenhet föreslår vi att man börjar med ett parallellt tillvägagångssätt med distribuerad data. Distribuerad dataparallellism som t.ex SageMaker distribuerade dataparallellt bibliotek löser de flesta nätverksproblem med modellreplikor, så du bör passa in modeller i det minsta antalet noder och sedan replikera för att skala batchstorlek efter behov.
Om du får ont om minne under träning, som vi gjorde i det här scenariot, kanske du vill byta till ett modellparallellt tillvägagångssätt. Men överväg dessa alternativ innan du provar parallellträning:
- På NVIDIA Tensor Core-utrustad hårdvara, använd träning med blandad precision för att skapa snabbare och minska minnesförbrukningen.
- Minska batchstorleken (eller minska bildupplösningen eller NLP-sekvenslängden, om möjligt).
Dessutom föredrar vi modelldesigner som inte har batchnormalisering som beskrivs i Högpresterande storskalig bildigenkänning utan normalisering. Om det inte kan undvikas, se till att batchnormalisering synkroniseras mellan enheter. När du använder distribuerad träning är din batch uppdelad på GPU:er, så exakt batchstatistik kräver synkronisering mellan alla enheter. Utan detta kommer normaliseringen att få ökat fel och därmed försämra konvergensen.
Börja med modell parallell träning när du har följande begränsningar:
- Din modell passar inte på en enda enhet
- På grund av din modellstorlek står du inför begränsningar när det gäller att välja större batchstorlekar, till exempel om din modellvikter tar upp det mesta av ditt GPU-minne och du tvingas välja en mindre, suboptimal batchstorlek
När du optimerar för prestanda, gör följande:
- Använd pipelining för kommunikation mellan noder för att minimera latens och öka genomströmningen
- Håll rörledningarna så korta som möjligt för att minimera eventuella bubblor. Antalet mikrobatcher bör ställas in för att balansera beräkningseffektivitet med bubbelstorlek och vara åtminstone rörledningens längd. Om det behövs kan du bilda mikrobatcher på tokennivå enligt beskrivningen i TeraPipe: Token Level Pipeline Parallelism för träning av storskaliga språkmodeller
När du optimerar för kostnad, använd SageMaker-hanterade Spot Instances för träning. Detta kan optimera kostnaden för träningsmodeller med upp till 90 % jämfört med On-Demand-instanser. SageMaker hanterar Spot-avbrotten för din räkning.
Andra faktorer att tänka på:
- Inom en nod när det finns en snabb sammankoppling är det mer nyanserat. Om det finns gott om nätverkskapacitet inom nod, kan omblandning av data för mer optimal beräkning visa en fördel.
- Om aktiveringar är mycket större än vikttensorer kan en shard optimizer också hjälpa. Vänligen hänvisa till Noll för mer detaljer.
Följande tabell ger några vanliga scenarier för uppskalning av träning och hur du kan konfigurera dem på AWS.
| Scenario | När gäller det? | Lösning |
| Skalning från en enda GPU till många GPU:er | När mängden träningsdata eller storleken på modellen är för stor | Byt till en multi-GPU-instans som p3.16xlarge, som har åtta GPU:er, med data och bearbetning uppdelad över de åtta GPU:erna, och producerar en nästan linjär hastighetsökning under tiden det tar att träna din modell. |
| Skalning från en enda instans till flera instanser | När skalningsbehoven sträcker sig längre än att ändra instansstorleken | Skala antalet instanser med SageMaker Python SDK:s estimatorfunktion genom att ställa in din instance_type till p3.16xlarge och instance_count till 2. Istället för de åtta GPU:erna på en enda p3.16xlarge, har du 16 GPU:er över två identiska instanser. Överväg att använda SageMaker distribuerade dataparallellt bibliotek. |
| Att välja en modell för parallellt tillvägagångssätt för utbildning | När du stöter på minnesfel under träning | Byt till en parallell modell med hjälp av SageMaker distribuerade modell parallellbibliotek. |
| Nätverksprestanda för kommunikation mellan noder | För distribuerad träning med flera instanser (till exempel kommunikation mellan noderna i klustret när man gör en AllReduce-operation) | Dina instanser måste vara i samma region och samma tillgänglighetszon. När du använder SageMaker Python SDK hanteras detta åt dig. Din träningsdata bör också vara i samma tillgänglighetszon. Överväg att använda SageMaker distribuerade dataparallellt bibliotek. |
| Optimerad GPU, nätverk och lagring | För storskaliga distribuerade utbildningsbehov | Instanstypen p4d.24xlarge designades för snabb lokal lagring och ett snabbt nätverksbakplan med upp till 400 gigabit, och vi rekommenderar det starkt som det mest presterande alternativet för distribuerad träning. |
Slutsats
Med modellparallellbiblioteket i SageMaker får vi många av fördelarna ur lådan, såsom automatiserad modellpartitionering och effektiv pipelining. I det här inlägget delade vi våra utmaningar med vårt ML-användningsfall, våra överväganden om olika träningsmetoder och hur vi använde Amazon SageMakers modellparallellismbibliotek för att påskynda vår träning. Det bästa av allt är att det nu bara kan ta några timmar att anta bästa praxis för modellparallellism och prestandaförbättringar som beskrivs här. Om det här inlägget hjälper dig eller inspirerar dig att lösa ett problem vill vi gärna höra om det! Vänligen dela dina kommentarer och feedback.
Referensprojekt
För mer information, se följande:
Om författarna
Prem Ranga är en Enterprise Solutions Architect baserad i Atlanta, GA. Han är en del av Machine Learning Technical Field Community och älskar att arbeta med kunder på deras ML- och AI-resa. Prem brinner för robotik, är en autonom forskare och byggde också de Alexa-kontrollerade ölhällen i Houston och andra platser.
Sarah Jane Hong är medgrundare och Chief Science Officer på Latent Space. Hennes bakgrund ligger i skärningspunkten mellan människa-datorinteraktion och maskininlärning. Hon ledde tidigare NLP-forskning på Sonar (uppköpt av Marchex), som betjänar företag inom konversations-AI-området. Hon är också en uppskattad AR/VR-utvecklare, efter att ha fått utmärkelser och stipendier från Oculus, Mozilla Mixed Reality och Microsoft Hololens.
Darryl Barnhart är medgrundare och Chief Technology Officer på Latent Space. Han är en erfaren utvecklare med erfarenhet av GPU-acceleration, datorgrafik, storskalig data och maskininlärning. Andra passioner inkluderar matematik, spelutveckling och studiet av information.
Ian Thompson är grundare och VD på Latent Space. Ian är en ingenjör och forskare inspirerad av det "angränsande möjliga" - teknologier som kommer att ha stor inverkan på våra liv. För närvarande fokuserat på att förenkla och skala multimodal representationsinlärning för att hjälpa till att bygga säker och kreativ AI. Han har tidigare varit med och byggt företag inom grafik/virtuell verklighet (AltspaceVR, förvärvat av Microsoft) och utbildning/NLP (HSE).
