Denne blog er medforfatter af Sarah Jane Hong CSO, Darryl Barnhart CTO og Ian Thompson CEO for Latent Space og Prem Ranga fra AWS.
Latent rum er en skjult repræsentation af abstrakte ideer, som maskinlæringsmodeller (ML) lærer. For eksempel er "hund", "blomst" eller "dør" begreber eller placeringer i latent rum. På Latent rum, vi arbejder på en motor, der giver dig mulighed for at manipulere og udforske dette rum med både sproglige og visuelle prompter. Latent Space-teamet kommer fra to områder, der længe har haft lidt overlap: grafik og naturlig sprogbehandling (NLP). Traditionelt er modaliteterne af billeder og tekst blevet håndteret separat, hver med deres egen historie med kompleks, kostbar og skrøbelig feature engineering. NLP-opgaver som dokumentforståelse eller besvarelse af spørgsmål har normalt ikke haft meget til fælles med visionsopgaver som sceneforståelse eller gengivelse, og normalt bruger vi meget forskellige tilgange og modeller til hver opgave. Men dette ændrer sig hurtigt.
Denne sammensmeltning af modaliteter i et enkelt delt latent rum låser op for en ny generation af kreative og kommercielle applikationer, fra spil til dokumentforståelse. Men at låse op for disse nye applikationer i en enkelt model åbner op for nye skaleringsudfordringer, som fremhævet i "The Bitter Lesson" af Richard Sutton, og det spændende arbejde i de sidste par år med skaleringslove. For at gøre dette muligt arbejder Latent Space på banebrydende forskning for at fusionere disse modaliteter i en enkelt model, men også for at skalere og gøre det effektivt. Det er her modelparallelisme kommer ind.
Amazon SageMaker's unikke automatiserede modelopdeling og effektive pipelining-tilgang gjorde vores adoption af modelparallelisme mulig med en lille ingeniørindsats, og vi skalerede vores træning af modeller ud over 1 milliard parametre (vi bruger p4d.24xlarge A100-forekomster), hvilket er et vigtigt krav for os. Desuden observerede vi, at når vi trænede med en 16 node, otte GPU træningsopsætning med SageMaker model parallelism bibliotek, registrerede vi en 38% forbedring i effektivitet sammenlignet med vores tidligere træningsløb.
Udfordringer med at træne transformatorer i stor skala
Hos Latent Space fusionerer vi sprog og vision i transformermodeller med milliarder af parametre for at understøtte "ud af distribution"-brugstilfælde fra en brugers fantasi, eller som ville forekomme i den virkelige verden, men ikke i vores træningsdata. Vi håndterer de udfordringer, der er forbundet med skalering til milliarder af parametre og videre på to forskellige måder:
Informationssøgningsteknikker har længe været en nøglekomponent i søgemaskiner og QA-opgaver. På det seneste er der sket spændende fremskridt med at kombinere klassiske IR-teknikker med moderne transformatorer, specielt til spørgsmålsbesvarelsesopgaver, hvor en model trænes sammen med en neural retriever, der lærer at hente relevante dokumenter for at hjælpe med at besvare spørgsmål. For et overblik, se det seneste arbejde fra FAIR in Retrieval Augmented Generation: Strømlining af skabelsen af intelligente naturlige sprogbehandlingsmodeller , Fusion-in-dekoder, Google Brain's RIGE, og Nvidias Neural Retriever for besvarelse af spørgsmål.
Mens genfindingsforstærkede teknikker hjælper med omkostninger og effektivitet, er vi stadig ikke i stand til at tilpasse modellen på en enkelt GPU til vores største model. Det betyder, at vi skal bruge modelparallelisme til at træne det. Men på grund af arten af vores genfindingsarkitektur var det udfordrende at designe vores modelopdeling på grund af indbyrdes afhængighed mellem hentede kontekster på tværs af træningsinput. Desuden, selvom vi bestemmer, hvordan vi opdeler vores model, var indførelsen af modelparallelisme en betydelig ingeniøropgave for os at udføre manuelt på tværs af vores forsknings- og udviklingslivscyklus.
SageMaker model parallelisme bibliotek
Modelparallelisme er processen med at opdele en model mellem flere enheder eller noder (såsom GPU-udstyrede instanser) og skabe en effektiv pipeline til at træne modellen på tværs af disse enheder for at maksimere GPU-udnyttelsen. Det model parallelisme bibliotek i SageMaker gør modelparallelisme mere tilgængelig ved at tilbyde automatiseret modelopdeling, også kaldet automatiseret modelopdeling og sofistikeret planlægning af pipelinekørsel. Modelopdelingsalgoritmerne kan optimere for hastighed eller hukommelsesforbrug. Biblioteket bruger en partitioneringsalgoritme, der afbalancerer hukommelsen, minimerer kommunikationen mellem enheder og optimerer ydeevnen.
Automatiseret modelopdeling
Til vores PyTorch-brugstilfælde kører modellens parallelle bibliotek internt et sporingstrin (i det første træningstrin), der konstruerer modelgrafen og bestemmer tensor- og parameterformerne. Den konstruerer derefter et træ, som består af de indlejrede nn.Module objekter i modellen, samt yderligere data indsamlet fra sporing, såsom mængden af lagret nn.Parameters, og køretid for hver nn.Module.
Biblioteket krydser derefter dette træ fra roden og kører en partitioneringsalgoritme, der afbalancerer beregningsbelastning og hukommelsesbrug og minimerer kommunikation mellem instanser. Hvis flere nn.Moduler deler den samme nn.Parameter, placeres disse moduler på den samme enhed for at undgå at opretholde flere versioner af den samme parameter. Efter at partitioneringsbeslutningen er truffet, indlæses de tildelte moduler og vægte til deres enheder.
Planlægning af rørledningskørsel
En anden kernefunktion i SageMaker distribuerede model parallelbibliotek er rørledninger, som bestemmer rækkefølgen, hvori beregninger foretages, og data behandles på tværs af enheder under modeltræning. Pipelining er baseret på at opdele en mini-batch i mikrobatches, som føres ind i træningspipelinen én efter én og følger en køreplan defineret af bibliotekets runtime.
Microbatch-pipelinen sikrer, at alle GPU'erne udnyttes fuldt ud, hvilket er noget, vi selv skulle bygge, men med modelparallelismebiblioteket sker dette pænt bag kulisserne. Til sidst kan vi bruge Amazon FSx, hvilket er vigtigt for at sikre, at vores læsehastigheder er hurtige i betragtning af antallet af filer, der læses under træningen af en multimodal model med genfinding.
Træningsarkitektur
Følgende diagram viser, hvordan vi opsætter vores træningsarkitektur. Vores primære mål var at forbedre træningshastigheden og reducere omkostningerne. De billed- og sprogtransformatorer, vi træner, er meget komplekse, med et betydeligt stort antal lag og vægte indeni, der kører til milliarder af parametre, hvilket alle gør dem ude af stand til at passe ind i hukommelsen af en enkelt node. Hver node bærer en delmængde af modellen, gennem hvilken datastrømmene og transformationerne deles og kompileres. Vi opretter 16 p4d.24xlarge instanser hver med otte GPU'er ved hjælp af følgende arkitekturrepræsentation:
Når vi skalerer vores modeller op, er en almindelig tendens at have alt gemt i netværkets vægte. Men af praktiske årsager ønsker vi at udvide vores modeller for at lære, hvordan man leder efter relevante sammenhænge for at hjælpe med gengivelsesopgaven. Dette gør os i stand til at holde vores serveringsomkostninger nede uden at gå på kompromis med billedkvaliteten. Vi bruger en stor transformer-baseret NLP-model, og som nævnt før observerede vi en stigning på 38 % i træningseffektivitet med SageMaker-modellens parallelitetsbibliotek som vist af følgende:
- Vi har brug for en allreduce for hver beregning i tilfælde af tensor niveau parallelisme. Dette tager O(log_2 n) parallelle trin. Det vil sige n maskiner, der tager O(n) trin, for O(n log_2 n) samlede operationer.
- For pipelineparallelisme kræver vi O(1) parallelle trin for at sende data ned i pipelinen
- Givet 16 maskiner med otte GPU'er, har vi O(1)-omkostninger for pipeline-parallelle og O(log_2(8)) = O(3)-omkostninger for dybdegående modelparallelle.
- I dette tilfælde ser vi, at netværksomkostningerne reduceres til 1/3 ved at skifte til pipeline-parallel i forhold til, hvad vi bruger med SageMaker modelparallelisme, og de samlede træningsomkostninger reduceres til 1/2 + 1/2 * 1/log_2(16) = 0.625 af den oprindelige omkostning, hvilket fører til en tilsvarende effektivitetsforbedring.
Generelt kan vi, når behovet berettiger distribueret træning (problemer med skalering af modelstørrelse eller træningsdata), følge et sæt bedste praksis for at bestemme, hvilken tilgang der fungerer bedst.
Bedste praksis for distribueret træning
Baseret på vores erfaring foreslår vi at starte med en distribueret data parallel tilgang. Distribueret dataparallelisme som f.eks SageMaker distribuerede data parallelbibliotek løser de fleste netværksproblemer med modelreplikaer, så du skal passe modeller ind i det mindste antal noder og derefter replikere for at skalere batchstørrelse efter behov.
Hvis du løber tør for hukommelse under træning, som vi gjorde i dette scenarie, vil du måske skifte til en model parallel tilgang. Overvej dog disse alternativer, før du prøver model parallel træning:
- På NVIDIA Tensor Core-udstyret hardware, brug træning med blandet præcision at skabe speedup og reducere hukommelsesforbruget.
- Reducer batchstørrelsen (eller reducer billedopløsningen eller NLP-sekvenslængden, hvis det er muligt).
Derudover foretrækker vi modeldesign, der ikke har batchnormalisering som beskrevet i Højtydende billedgenkendelse i stor skala uden normalisering. Hvis det ikke kan undgås, skal du sikre dig, at batchnormalisering er synkroniseret på tværs af enheder. Når du bruger distribueret træning, er din batch opdelt på tværs af GPU'er, så nøjagtig batchstatistik kræver synkronisering på tværs af alle enheder. Uden dette vil normaliseringen have øget fejl og derved forringe konvergensen.
Start med model parallel træning, når du har følgende begrænsninger:
- Din model passer ikke på en enkelt enhed
- På grund af din modelstørrelse står du over for begrænsninger i at vælge større batchstørrelser, såsom hvis din modelvægt fylder det meste af din GPU-hukommelse, og du er tvunget til at vælge en mindre, suboptimal batchstørrelse
Når du optimerer til ydeevne, skal du gøre følgende:
- Brug pipelining til inter-node-kommunikation for at minimere latens og øge gennemløbet
- Hold rørledninger så korte som muligt for at minimere eventuelle bobler. Antallet af mikrobatcher bør indstilles til at balancere beregningseffektivitet med boblestørrelsen og være mindst rørledningens længde. Om nødvendigt kan du danne mikrobatcher på token-niveau som beskrevet i TeraPipe: Token Level Pipeline Parallelism til træning af store sprogmodeller
Når du optimerer for omkostninger, skal du bruge SageMaker administrerede Spot Instances til træning. Dette kan optimere omkostningerne ved træningsmodeller med op til 90 % i forhold til On-Demand-instanser. SageMaker administrerer Spot-afbrydelserne på dine vegne.
Andre faktorer at overveje:
- Inden for en node, når der er en hurtig sammenkobling, er det mere nuanceret. Hvis der er rigelig intra-node netværkskapacitet, kan det være en fordel at blande data for mere optimal beregning.
- Hvis aktiveringer er meget større end vægttensorer, kan en sharded optimizer også hjælpe. Vær sød at henvise til Nul for flere detaljer.
Følgende tabel viser nogle almindelige scenarier for træningsskalering, og hvordan du kan konfigurere dem på AWS.
| Scenario | Hvornår gælder det? | Løsning |
| Skalering fra en enkelt GPU til mange GPU'er | Når mængden af træningsdata eller størrelsen af modellen er for stor | Skift til en multi-GPU-instans som f.eks. p3.16xlarge, som har otte GPU'er, med data og behandling fordelt på de otte GPU'er, og frembring en næsten lineær hastighedsstigning i den tid, det tager at træne din model. |
| Skalering fra en enkelt instans til flere instanser | Når skaleringsbehovene rækker ud over at ændre instansstørrelsen | Skaler antallet af forekomster med SageMaker Python SDK's estimatorfunktion ved at indstille din instance_type til p3.16xlarge og instance_count til 2. I stedet for de otte GPU'er på en enkelt p3.16xlarge har du 16 GPU'er fordelt på to identiske forekomster. Overvej at bruge SageMaker distribuerede data parallelbibliotek. |
| Valg af en model parallel tilgang til træning | Når du støder på fejl i hukommelsen under træning | Skift til en model parallel tilgang ved hjælp af SageMaker distribueret model parallelbibliotek. |
| Netværksydelse til inter-node kommunikation | Til distribueret træning med flere instanser (f.eks. kommunikation mellem noderne i klyngen, når du udfører en AllReduce-operation) | Dine forekomster skal være i samme region og samme tilgængelighedszone. Når du bruger SageMaker Python SDK, håndteres dette for dig. Dine træningsdata bør også være i den samme tilgængelighedszone. Overvej at bruge SageMaker distribuerede data parallelbibliotek. |
| Optimeret GPU, netværk og lagring | Til store distribuerede træningsbehov | p4d.24xlarge-instanstypen er designet til hurtig lokal lagring og et hurtigt netværksbagplan med op til 400 gigabit, og vi anbefaler det stærkt som den mest effektive mulighed for distribueret træning. |
Konklusion
Med model parallelbiblioteket i SageMaker får vi mange af fordelene ud af boksen, såsom automatiseret modelopdeling og effektiv pipelining. I dette indlæg delte vi vores udfordringer med vores ML use case, vores overvejelser om forskellige træningstilgange, og hvordan vi brugte Amazon SageMaker model parallelism biblioteket til at fremskynde vores træning. Det bedste af det hele er, at det nu kun kan tage et par timer at vedtage bedste praksis for modelparallelisme og ydeevneforbedringer, der er beskrevet her. Hvis dette indlæg hjælper dig eller inspirerer dig til at løse et problem, vil vi meget gerne høre om det! Del venligst dine kommentarer og feedback.
Referencer
For mere information, se venligst følgende:
Om forfatterne
Prem Ranga er en Enterprise Solutions Architect baseret i Atlanta, GA. Han er en del af Machine Learning Technical Field Community og elsker at arbejde med kunder på deres ML- og AI-rejse. Prem er passioneret omkring robotteknologi, er en forsker i autonome køretøjer og byggede også de Alexa-kontrollerede Beer Pours i Houston og andre steder.
Sarah Jane Hong er medstifter og Chief Science Officer hos Latent Space. Hendes baggrund ligger i skæringspunktet mellem menneske-computer-interaktion og maskinlæring. Hun har tidligere ledet NLP-forskning hos Sonar (erhvervet af Marchex), som betjener virksomheder i det konverserende AI-rum. Hun er også en værdsat AR/VR-udvikler efter at have modtaget priser og stipendier fra Oculus, Mozilla Mixed Reality og Microsoft Hololens.
Darryl Barnhart er medstifter og Chief Technology Officer hos Latent Space. Han er en erfaren udvikler med erfaring i GPU-acceleration, computergrafik, data i stor skala og maskinlæring. Andre passioner omfatter matematik, spiludvikling og studiet af information.
Ian Thompson er stifter og administrerende direktør hos Latent Space. Ian er en ingeniør og forsker inspireret af de "tilstødende mulige" - teknologier, der er ved at have en stor indflydelse på vores liv. I øjeblikket fokuseret på at forenkle og skalere multimodal repræsentationslæring for at hjælpe med at opbygge sikker og kreativ AI. Han har tidligere været med til at bygge virksomheder indenfor grafik/virtuel virkelighed (AltspaceVR, opkøbt af Microsoft) og uddannelse/NLP (HSE).
