Hvordan Amazon Search M5 sparte 30 % for LLM-opplæringskostnader ved å bruke AWS Trainium | Amazon Web Services

I flere tiår har Amazon vært banebrytende og innovert maskinlæring (ML), og bringer herlige opplevelser til kundene sine. Fra de tidligste dagene har Amazon brukt ML til ulike brukstilfeller som bokanbefalinger, søk og svindeloppdagelse. I likhet med resten av bransjen har fremskritt med akselerert maskinvare gjort det mulig for Amazon-team å forfølge modellarkitekturer ved å bruke nevrale nettverk og dyp læring (DL).

M5-programmet innen Amazon Search eier oppdagelseslæringsstrategien for Amazon og bygger storskalamodeller på tvers av flerspråklige, multi-lokale, multi-entiteter, multitask og multimodale som tekst, bilde og video. M5-programmet har tjent universelle innebygginger og storskala fundamentmodeller til hundrevis av ML-team på tvers av Amazon, samtidig som det har opprettholdt strenge kontroller over kostnadsoptimalisering. For å oppnå dette, evaluerer M5-teamet regelmessig nye teknikker for å redusere kostnadene.

Som mange ML-organisasjoner, brukes akseleratorer i stor grad for å akselerere DL-trening og slutninger. Da AWS lanserte spesialbygde akseleratorer med den første utgivelsen av AWS slutning i 2020 begynte M5-teamet raskt å gjøre det bruke dem til å distribuere produksjonsbelastninger mer effektivt, sparer både kostnader og reduserer ventetiden. I fjor lanserte AWS sin AWS Trainium akseleratorer, som optimerer ytelsen per kostnad for å utvikle og bygge neste generasjons DL-modeller. I dette innlegget diskuterer vi hvordan M5 klarte å redusere kostnadene for å trene modellene sine med 30 %, og deler noen av de beste fremgangsmåtene vi lærte underveis.

Trainium-forekomster

Med fremskrittene innen spesialbygde akseleratorer, tilbyr Amazon også overbevisende akseleratorer i form av AWS Inferentia og Trainium. Som navnene tilsier, er disse brikkene optimalisert for å overgå behovene til henholdsvis slutninger og treningsarbeidsmengder. For storskala opplæring av fundamentmodeller som når milliarder av parametere i størrelse, Trainium Trn1- og Trn1n-forekomster er ideelle valg på grunn av deres egenskaper. Trn1-forekomster er drevet av det siste NeuronCore-v2, og har en rikelig mengde akseleratorberegning og minne. Trn1n-forekomster kan også velges for en større mengde nettverksbåndbredde (1,600 Gbs), så de er ideelt egnet for effektiv trening med kostnadsoptimalisering i tankene.

For å bruke akseleratorer trenger du et programvarelag som støtter dem. Med Trn- og Inf-brikker er AWS Neuron SDK låser opp spesialbygde akseleratorer fra Amazon ved hjelp av PyTorch XLA. PyTorch XLA konverterer PyTorchs ivrige modus til grafbasert implementering i latmodus. Disse grafene brukes deretter og kompileres videre for å brukes med akseleratoren. PyTorch Neuron (en del av Neuron SDK) gjør at PyTorch-brukere kan trene modellene sine på Trainium NeuronCores med noen få linjer med kode.

Modell og arbeidsmengde

M5-teamet trener og distribuerer grunnleggende modeller og universelle representasjoner for å hjelpe ulike team over hele Amazon med å bringe glede til Amazon.com kunder. En slik modell er en tekstkodermodell etterfulgt av en flerlagsperceptron (MLP) med eksplisitte eller implisitte funksjonsinteraksjoner definert av den nevrale nettverksarkitekturen med hundrevis av millioner av trenbare parametere. Denne modellen er trent på milliarder av tokens, og brukes til å generere millioner av innebygginger i en offline batch-inferensinnstilling. Disse innebyggingene er innganger til en kundevendt tier-1 Amazon-tjeneste.

Infrastrukturen for produksjonsrørledningen bruker AWS-batch med fair share køstrategier, ved å bruke en EFA-aktivert multi-node trn1.32xlarge cluster som beregning for modelltrening. Funksjonelt utfører produksjonsrørledningen inkrementell modellopplæring, evaluering av opplært modell og offline batch-slutning på den trente modellen, alt ved å bruke PyTorch som det underliggende DL-biblioteket.

Mål

Å glede kundene våre er en fremste grunnsetning. Gitt den kundevendte karakteren til rørledningen, er det avgjørende at alle servicenivåavtaler (SLAer) oppfylles uten regresjoner. Vi identifiserte to kritiske akseptkriterier for å tilpasse vår eksisterende GPU-produksjonspipeline og overføre den til Trainium:

• Modellkvalitet – Kvaliteten på modellene våre påvirker kundeopplevelsen direkte. Vi krever at det skal være mindre enn 0.1 % forskjell i modellkvalitet mellom GPU og Trainium.
• Treningsgjennomstrømning – Vi trener iterativt modellene våre med jevne mellomrom for å gi den ferskeste opplevelsen til kundene våre. Vi krever at modellkonvergens må oppnås innen en forhåndsdefinert tidsperiode (som 1 uke) for å oppfylle produksjons-SLAene våre.

I de følgende delene deler vi vår reise med å jobbe bakover fra dette kriteriet, og vår læring for å støtte produksjonsarbeidsbelastninger i Amazon-skala.

Treningsmanus

Før vi begynner med modellopplæring, må vi gjøre endringer i opplæringsskriptet for å gjøre det XLA-kompatibelt. Gitt størrelsen på modellen bruker vi distribuert dataparallell (DDP) for å trene modellen. DDP lar oss øke gjennomstrømningen av modelltrening ved å skalere opp antall maskiner som brukes til å kjøre modelltrening, uten kodeendringer. Vi fulgte instruksjonene gitt i Neuron PyTorch MLP opplæringsopplæring å legge til XLA-spesifikke konstruksjoner i treningsskriptene våre. Disse kodeendringene er enkle å implementere. Følgende er noen betydelige tekniske lærdommer fra øvelsen som i stor grad forbedret modellgjennomstrømningen vår:

• Plassering av xm.mark_step() - xm.mark_step() kompilerer og kjører de dovent innsamlede beregningsgrafene. Påkaller mark_step for mange ganger vil føre til et større antall små grafer, mens å påkalle det for få ganger vil føre til få, men store grafer. Avhengig av søknaden din, vil gjennomstrømningen og implementeringen av modellopplæringen din variere basert på plasseringen din xm.mark_step(). Vår implementering plasserer en xm.mark_step() etter en pasning forover og bakover, og en etter optimeringstrinnet.
• Datalasterinnpakning med XLA multiprosesseringsenhetslaster – Dette er et kritisk skritt som lett kan gå glipp av. Enhetslasteren for multiprosessering torch_xla.distributed.parallel_loader.MpDeviceLoader laster treningsdata på hver XLA-enhet med alternativer for å forhåndslaste og overlappe datalasting med enhetskjøringer for å forbedre gjennomstrømningen. Enhetslasteren påkaller også xm.mark_step() og er derfor i stand til å bygge grafer for datalasting til enhet fra vert.

Samling for Trainium

Tradisjonelt innebærer modellutviklingssyklusen med GPUer å gjøre endringer i modellen eller treningsskriptet og kjøre det direkte på GPU-enheten. Akseleratorer som Trainium som bruker XLA krever et ekstra trinn før modelltrening kan kjøres på gasspedalen. XLA-beregningsgrafer kan bare kjøres etter at de er kompilert. Generelt er det to måter å utføre denne kompileringen på: Ahead of Time (AOT), hvor du sporer og kompilerer alle grafer først og deretter kjører dem, eller Just In Time (JIT), hvor grafer spores, kompileres og kjøres etter hvert som de blir møtt. Neuron SDK gir begge disse ut av esken. Vanligvis utføres AOT-kompilering først. Grafer kjøres deretter etter denne kompileringen. Hvis det oppdages nye grafer, starter Neuron runtime en JIT-kompilering før de kjøres. For å utføre AOT-kompilering gir Neuron SDK neuron_parallel_compile, et kompileringsverktøy som trekker ut grafer fra en prøvekjøring av treningsskriptet og utfører parallell AOT-kompilering.

Et viktig aspekt ved AOT-kompilering er å sikre at ingen nye beregningsgrafer opprettes i løpet av opplæringen. En kilde til nye beregningsgrafer (og derfor rekompileringer) er dynamiske former for treningsgruppene under modelltrening. Vi fant ut at bruk av statiske former og partier med fast størrelse eliminerer samlinger av treningstid og forbedrer treningsgjennomstrømningen betraktelig uten noen innvirkning på modellens nøyaktighet. Ved å håndheve slike begrensninger på opplæring, observerte vi at bare 4–5 trinn med modelltrening, ett trinn med modellvalidering og kontroll av modellen én gang er nødvendig for å spore alle grafene under AOT-kompilering. Det er viktig å merke seg at Neuron SDK er i stadig utvikling, og i fremtiden vil støtte dynamiske former også.

Videre lagres de kompilerte grafene i Nevron Persistent Cache på disk eller i en Amazon enkel lagringstjeneste (Amazon S3) bøtte. Dette er spesielt nyttig for produksjonsarbeidsbelastninger der modellarkitektur og treningskonfigurasjon ikke endres. Derfor påløper overheaden for kompilering bare én gang. Å bruke cachen er like enkelt som å sette et miljøflagg:

export NEURON_COMPILE_CACHE_URL="s3://BUCKET/KEY"

Neuron-kompilatoren gir også tre optimeringsalternativer på kompilatornivå (O1, O2, O3) for å balansere kompileringstid og modellkjøring. O1 muliggjør kjerneoptimaliseringer på beregningsgrafen og minimerer kompileringstiden, O3 gir forbedret modellkjøringsgjennomstrømning på bekostning av høyere kompileringstid, og O2 (standardalternativ) er en balanse mellom de to. For vår brukstilfelle brukte vi O1-optimalisering og observerte en 86 % reduksjon i kompileringstid uten endringer i modellnøyaktighetsmålinger, mens vi observerte omtrent 5–7 % reduksjon i gjennomstrømming sammenlignet med standardoptimalisering (O2). Avhengig av brukstilfellet kan du velge ulike optimaliseringsnivåer.

For å oppsummere brukte vi følgende flagg for kompilering:

NEURON_CC_FLAGS="--target trn1 --auto-cast all --auto-cast-type bf16 --model-type transformer --optlevel O1"

Sjekkpunktkompatibilitet

Når kompileringen er fullført, kan vi fortsette å trene modellene våre på Trainium. Som nevnt tidligere trener vi modellene våre trinnvis, noe som betyr at vi laster inn et tidligere trent modellsjekkpunkt og fortsetter å trene med nye data. PyTorch og PyTorch XLA tillater sømløs overgang mellom akseleratorer gjennom sjekkpunktinteroperabilitet. Å ha fleksibiliteten til å flytte mellom GPU og Trainium gjorde det mulig for oss å sømløst laste den forrige GPU-modellen og trene på Trainium-maskiner. Dette var avgjørende for å sikre at vi kan initialisere modellen vår med den best trente modellen uten produksjonsstans eller tap i modellnøyaktighet.

Fordi GPU-modellen ble lagret ved å bruke standard PyTorch-modelllagringsverktøy, kunne vi bruke PyTorch-sjekkpunkt-lasteverktøyet for å laste GPU-modellen på Trainium-enheter.

For eksempel, på GPU/CPU, kan du lagre modellen med følgende kode:

torch.save(model.state_dict(), PATH)

Deretter laster du modellen tilbake på Trainium:

import torch_xla.core.xla_model as xm
xla_device = xm.xla_device()
model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(xla_device)

På samme måte kan du lagre modellen på Trainium med følgende kode:

import torch_xla.core.xla_model as xm
# automatically moves the data to CPU for the master device
xm.save(model.state_dict(), PATH) 

Og last modellen tilbake på GPU/CPU:

model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device) # can be any device

Faktisk, fordi vi bruker DDP for modelltrening, er modellbelastningen agnostisk for antall maskiner som ble brukt til å trene forrige sjekkpunkt. Dette gjør at vi kan skalere Trn1-flåten horisontalt uten kodeendringer eller negative effekter på modelltrening. Disse PyTorch-baserte sjekkpunktene kan brukes direkte eller til og med fakkelskript for slutningsbrukstilfeller på AWS Inferentia2 eller andre akseleratorer.

Driftsstabilitet

Det kan ikke understrekes nok at kjøring av arbeidsbelastninger i produksjon krever at flere SLAer oppfylles. For vår brukssituasjon, bortsett fra SLAene for modellkvalitet og opplæringsgjennomstrømning, er det avgjørende at produksjonsrørledningen er driftsstabil, noe som betyr minimal nedetid og forstyrrelser under modellopplæring, evaluering og konklusjon.

Som med den eksisterende GPU-baserte rørledningen, la vi til en rekke mekanismer for å gjøre rørledningen driftsstabil. Før vi starter modelltrening, kjører vi flere tilregnelighetstester for å vurdere helsen til maskinene. Disse testene inkluderer vanligvis enkle tensoroperasjoner for å verifisere helsen til akseleratorenhetene. Vi har observert at for distribuert opplæring er det viktig å kjøre tester for å bekrefte kollektiv kommunikasjon også mellom instanser. Vi brukte NCCOM testsuite fra Neuron SDK for å oppnå dette, og kjører en rekke operasjoner som alt-samler, all-reduser og reduser-spredning.

Selv etter å ha fulgt forslagene vi har nevnt, har vi observert at forbigående problemer er uunngåelige i enhver pipeline, uavhengig av den underliggende akseleratoren. For å bygge spenst i en hvilken som helst treningspipeline, anbefaler vi å bygge inn prøvemekanismer for å løse disse potensielle problemene. Vi bruker AWS Batch automatiserte gjenforsøk å prøve jobber som støter på en forbigående feil under modelltrening. Disse omstartene kan være kostbare hvis det oppstår en feil mot slutten av treningen. For å motvirke dette problemet har vi tilpasset treningsskriptene våre for å laste et tidligere trent modellsjekkpunkt og fortsette treningen fra det punktet. Med denne funksjonaliteten er vi i stand til aggressivt å starte mislykkede treningsjobber på nytt med minimal overhead.

Med disse elastisitetsmekanismene på plass, var vi i stand til å oppnå 98.5 % suksessrater for arbeidsbelastningene våre på Trn1, sammenlignbare med suksessratene våre for eksisterende GPU-pipeline.

Resultater

For å validere nøyaktigheten til modellene våre initialiserte vi to modeller fra samme GPU-sjekkpunkt, og trente en på Trainium og den andre på en sammenlignbar GPU. Begge modellene ble trent med de samme treningshyperparametrene. Datasettet som brukes til beregning av beregninger er et holdout-datasett, og vi evaluerer modellens nøyaktighet på dette datasettet for hvert N globale trinn. X-aksen er det globale trinnet, og Y-aksen er modellens nøyaktighet. Vi observerte mindre enn 0.1 % forskjell i modellnøyaktigheten på hvert punkt i følgende graf.

Videre, for å evaluere kostnadseffektiviteten til modelltreningen, foretrekker vi å sammenligne veggklokketiden det tar å nå modellkonvergens. Vi mener dette gir et mer praktisk syn på kostnadsbesparelser sammenlignet med tiltak som kostnad per token, oppnådde FLOPS/dollar og andre faktorer. Med tanke på treningstiden til trn1.32xl og sammenlignbar Amazon Elastic Compute Cloud (Amazon EC2) har vi observert at Trainium tilbyr opptil 30 % billigere kostnad for modellkonvergens.

konklusjonen

Det er mange faktorer å vurdere når du evaluerer forskjellige akseleratorer for DL-arbeidsbelastningene dine. Noen av de viktigste er modellkvalitet, gjennomstrømning, kostnad og tilgjengelighet. Det er viktig å sikre at modellens kvalitet og gjennomstrømning ikke blir ofret basert på akseleratoren du velger.

Takket være vårt partnerskap og samarbeid med Annapurna Neuron-teamet, har Amazon Search M5-teamet kunnet spare opptil 30 % i kostnader ved å flytte til Trainium. Teamet er i stand til å bruke Trainium og oppnå modellkvalitet og gjennomstrømningsparitet med sammenlignbare akseleratorer i markedet. Kontrollpunktinteroperabilitet og minimale kodeendringer med støtte for XLA har gjort det mulig for M5 å velge mellom flere akseleratorer for arbeidsbelastningen deres. Dette har gjort det mulig for M5-teamet å dra nytte av den store regnekraften til Trainium, og bygge akseleratoragnostiske løsninger for å glede Amazon.com-kunder. Fra et operativt synspunkt har Trainium vist seg i stand til å støtte tier-1-tjenester på Amazon-skala. M5-teamet fortsetter å flytte flere arbeidsmengder til Trainium for å tilby de beste modellene for Amazon til de laveste kostnadene.

Oppsummert har M5-teamet vært i stand til å utføre kostnadseffektiv ML-trening i produksjonsgrad ved å legge Trainium til flåten av akseleratorer. Vi oppfordrer deg til å ta en titt på Trainium og andre Neuron-enheter som AWS Inferentia for å høste fordelene av spesialbygd Amazon-silisium for ML-arbeidsbelastninger. Kom enkelt i gang med en av de mange veiledningene med forskjellige modeller, som Llama 2, tilgjengelig på Trainium.

Om forfatterne

Abhinandan Patni er senior programvareingeniør hos Amazon Search. Han fokuserer på å bygge systemer og verktøy for skalerbar distribuert dyp læringstrening og sanntidsslutning.

James Park er løsningsarkitekt hos Amazon Web Services. Han jobber med Amazon.com for å designe, bygge og distribuere teknologiløsninger på AWS, og har en spesiell interesse for AI og maskinlæring. På fritiden liker han å oppsøke nye kulturer, nye opplevelser og holde seg oppdatert med de nyeste teknologitrendene. Du kan finne ham på Linkedin.

Jerry Mannil er programvareingeniør hos Amazon Search. Han jobber med å forbedre effektiviteten, robustheten og skalerbarheten til den distribuerte treningsinfrastrukturen.

Ken Su er programvareingeniør hos Amazon Search. Han jobber med å forbedre treningseffektiviteten og skalerbar distribuert treningsarbeidsflyt. Utenfor jobben liker han fotturer og tennis.

RJ er ingeniør innen Amazon. Han bygger og optimerer systemer for distribuerte systemer for opplæring og jobber med å optimalisere adopterende systemer for å redusere latens for ML Inference. Utenfor jobben utforsker han å bruke Generative AI for å bygge matoppskrifter.

• SEO-drevet innhold og PR-distribusjon. Bli forsterket i dag.
• PlatoData.Network Vertical Generative Ai. Styrk deg selv. Tilgang her.
• PlatoAiStream. Web3 Intelligence. Kunnskap forsterket. Tilgang her.
• PlatoESG. Karbon, CleanTech, Energi, Miljø, Solenergi, Avfallshåndtering. Tilgang her.
• PlatoHelse. Bioteknologisk og klinisk etterretning. Tilgang her.
• kilde: https://aws.amazon.com/blogs/machine-learning/how-amazon-search-m5-saved-30-for-llm-training-cost-by-using-aws-trainium/