Hur Amazon Search M5 sparade 30 % för LLM-utbildningskostnader genom att använda AWS Trainium | Amazon webbtjänster

I decennier har Amazon banat väg för och förnyat maskininlärning (ML), vilket ger sina kunder härliga upplevelser. Från de tidigaste dagarna har Amazon använt ML för olika användningsfall som bokrekommendationer, sökning och upptäckt av bedrägerier. I likhet med resten av branschen har framstegen med accelererad hårdvara gjort det möjligt för Amazon-team att bedriva modellarkitekturer med hjälp av neurala nätverk och djupinlärning (DL).

M5-programmet inom Amazon Search äger upptäcktsinlärningsstrategin för Amazon och bygger storskaliga modeller över flerspråkiga, multi-lokala, multi-entity, multitask och multimodala som text, bild och video. M5-programmet har betjänat universella inbäddningar och storskaliga grundmodeller till hundratals ML-team över hela Amazon samtidigt som det har upprätthållit strikta kontroller över kostnadsoptimering. För att uppnå detta utvärderar M5-teamet regelbundet nya tekniker för att minska kostnaderna.

Liksom många ML-organisationer används acceleratorer till stor del för att påskynda DL-träning och slutledning. När AWS lanserade specialbyggda acceleratorer med den första utgåvan av AWS slutledning 2020 började M5-teamet snabbt göra det använda dem för att mer effektivt distribuera produktionsbelastningar, vilket sparar både kostnader och minskar latens. Förra året lanserade AWS sin AWS Trainium acceleratorer, som optimerar prestanda per kostnad för att utveckla och bygga nästa generations DL-modeller. I det här inlägget diskuterar vi hur M5 kunde sänka kostnaden för att träna sina modeller med 30 %, och delar med oss ​​av några av de bästa metoderna vi lärde oss på vägen.

Trainium instanser

Med framstegen inom specialbyggda acceleratorer tillhandahåller Amazon också övertygande acceleratorer i form av AWS Inferentia och Trainium. Som deras namn antyder är dessa marker optimerade för att överträffa behoven av slutlednings- respektive träningsbelastningar. För storskalig utbildning av grundmodeller som når miljarder parametrar i storlek, Trainium Trn1- och Trn1n-instanser är idealiska val på grund av deras egenskaper. Trn1-instanser drivs av det senaste NeuronCore-v2, och har en stor mängd acceleratorberäkningar och minne. Trn1n-instanser kan också väljas för en större mängd nätverksbandbredd (1,600 XNUMX Gbs), så de är idealiska för prestandaträning med kostnadsoptimering i åtanke.

För att använda acceleratorer behöver du ett mjukvarulager för att stödja dem. Med Trn- och Inf-chips är AWS Neuron SDK låser upp Amazon specialbyggda acceleratorer med hjälp av PyTorch XLA. PyTorch XLA konverterar PyTorchs ivriga läge till grafbaserad implementering i lazy mode. Dessa grafer används sedan och sammanställs vidare för att användas med acceleratorn. PyTorch Neuron (en del av Neuron SDK) gör det möjligt för PyTorch-användare att träna sina modeller på Trainium NeuronCores med några rader kod.

Modell och arbetsbelastning

M5-teamet tränar och distribuerar grundläggande modeller och universella representationer för att hjälpa olika team över hela Amazon att skapa glädje till Amazon.com kunder. En sådan modell är en textkodarmodell följt av en multi-layer perceptron (MLP) med explicita eller implicita funktionsinteraktioner definierade av den neurala nätverksarkitekturen med hundratals miljoner inlärningsbara parametrar. Den här modellen tränas på miljarder tokens och används för att generera miljontals inbäddningar i en offline-batch-inferensinställning. Dessa inbäddningar är indata till en kundinriktad Amazon-tjänst på nivå 1.

Infrastrukturen för produktionsledningen använder AWS-batch med fair share köstrategier, med ett EFA-aktiverat trn1.32xlarge-kluster med flera noder som beräkning för modellträning. Funktionellt utför produktionspipelinen inkrementell modellträning, utvärdering av tränad modell och offline-batch slutledning på den tränade modellen, allt med PyTorch som det underliggande DL-biblioteket.

Mål

Att glädja våra kunder är en främsta grundsats. Med tanke på pipelinens kundinriktade karaktär är det avgörande att alla servicenivåavtal (SLA) uppfylls utan regressioner. Vi identifierade två kritiska acceptanskriterier för att anpassa vår befintliga GPU-produktionspipeline och överföra den till Trainium:

• Modellkvalitet – Kvaliteten på våra modeller påverkar kundupplevelsen direkt. Vi kräver att det ska vara mindre än 0.1 % skillnad i modellkvalitet mellan GPU och Trainium.
• Träningsgenomströmning – Vi tränar iterativt våra modeller med jämna mellanrum för att ge våra kunder den senaste upplevelsen. Vi kräver att modellkonvergens måste uppnås inom en fördefinierad tidsperiod (som 1 vecka) för att uppfylla våra produktions-SLA.

I följande avsnitt delar vi med oss ​​av vår resa att arbeta bakåt från detta kriterium och våra lärdomar för att stödja produktionsarbetsbelastningar i Amazon-skala.

Träningsmanus

Innan vi börjar med modellträning måste vi göra ändringar i träningsskriptet för att göra det XLA-kompatibelt. Med tanke på modellens storlek använder vi distribuerad dataparallell (DDP) för att träna modellen. DDP tillåter oss att öka genomströmningen av modellträning genom att skala upp antalet maskiner som används för att köra modellträning, utan några kodändringar. Vi följde instruktionerna i Neuron PyTorch MLP träningshandledning att lägga till XLA-specifika konstruktioner i våra träningsskript. Dessa kodändringar är enkla att implementera. Följande är några betydande tekniska lärdomar från övningen som avsevärt förbättrade vår modellgenomströmning:

• Placering av xm.mark_step() - xm.mark_step() sammanställer och kör de lättsamlade beräkningsgraferna. Åberopar mark_step för många gånger kommer att leda till ett större antal små grafer, medan att anropa det för få gånger leder till få men stora grafer. Beroende på din applikation kommer genomströmningen och implementeringen av din modellutbildning att variera beroende på din placering av xm.mark_step(). Vår implementering placerar en xm.mark_step() efter en fram- och bakåtpassning, och en efter optimeringssteget.
• Data loader inpackning med XLA multiprocessing enhet loader – Det här är ett kritiskt steg som lätt kan missas. Enhetsladdaren för flera processer torch_xla.distributed.parallel_loader.MpDeviceLoader laddar träningsdata på varje XLA-enhet med alternativ för att förladda och överlappa dataladdning med enhetskörningar för att förbättra genomströmningen. Enhetsladdaren anropar också xm.mark_step() och kan därför bygga grafer för dataladdning till enhet från värd.

Sammanställning för Trainium

Traditionellt innebär modellutvecklingscykeln med GPU:er att göra ändringar i modellen eller träningsskriptet och köra det direkt på GPU-enheten. Acceleratorer som Trainium som använder XLA kräver ytterligare ett steg innan modellträning kan köras på gaspedalen. XLA-beräkningsdiagram kan endast köras efter att de har kompilerats. Generellt finns det två sätt att utföra denna kompilering: Ahead of Time (AOT), där du spårar och kompilerar alla grafer först och sedan kör dem, eller Just In Time (JIT), där grafer spåras, kompileras och körs som de påträffas. Neuron SDK tillhandahåller båda dessa ur förpackningen. Vanligtvis utförs AOT-kompilering först. Grafer körs sedan efter denna kompilering. Om nya grafer påträffas, anropar Neuron runtime en JIT-kompilering innan de körs. För att utföra AOT-kompilering tillhandahåller Neuron SDK neuron_parallel_compile, ett kompileringsverktyg som extraherar grafer från en testkörning av träningsskriptet och utför parallell AOT-kompilering.

En viktig aspekt av AOT-kompilering är att säkerställa att inga nya beräkningsdiagram skapas under utbildningens gång. En källa till nya beräkningsgrafer (och därför omkompilering) är dynamiska former av träningsbatcherna under modellträning. Vi fann att användningen av statiska former och partier med fast storlek eliminerar träningstidssammanställningar och avsevärt förbättrar träningsgenomströmningen utan någon effekt på modellens noggrannhet. Genom att upprätthålla sådana begränsningar för träning, observerade vi att endast 4–5 steg av modellträning, ett steg av modellvalidering och kontroll av modellen en gång krävs för att spåra alla grafer under AOT-kompilering. Det är viktigt att notera att Neuron SDK ständigt utvecklas och i framtiden kommer att stödja dynamiska former också.

Dessutom lagras de sammanställda graferna i Neuron Persistent Cache på disk eller i en Amazon enkel lagringstjänst (Amazon S3) hink. Detta är särskilt användbart för produktionsbelastningar där modellarkitektur och utbildningskonfiguration inte ändras. Därför uppstår omkostnaden för sammanställningen bara en gång. Att använda cachen är lika enkelt som att ställa in en miljöflagga:

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

Neuron-kompilatorn tillhandahåller också tre optimeringsalternativ på kompilatornivå (O1, O2, O3) för att balansera kompileringstid och modellkörning. O1 möjliggör kärnoptimeringar på beräkningsgrafen och minimerar kompileringstiden, O3 ger förbättrad modellkörningsgenomströmning till priset av högre kompileringstid, och O2 (standardalternativ) är en balans mellan de två. För vårt användningsfall använde vi O1-optimeringen och observerade en 86 % minskning av kompileringstiden utan någon förändring av modellens noggrannhetsmått, samtidigt som vi observerade ungefär en 5–7 % minskning av genomströmningen jämfört med standardoptimeringen (O2). Beroende på användningsfallet kan du välja olika optimeringsnivåer.

För att sammanfatta använde vi följande flaggor för sammanställning:

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

Kontrollpunktskompatibilitet

När sammanställningen är klar kan vi fortsätta att träna våra modeller på Trainium. Som tidigare nämnts tränar vi inkrementellt våra modeller, vilket innebär att vi laddar en tidigare tränad modellkontrollpunkt och fortsätter träna med ny data. PyTorch och PyTorch XLA möjliggör sömlös övergång mellan acceleratorer genom interoperabilitet för kontrollpunkter. Att ha flexibiliteten att flytta mellan GPU och Trainium gjorde det möjligt för oss att sömlöst ladda den tidigare GPU-modellen och träna på Trainium-maskiner. Detta var avgörande för att säkerställa att vi kan initiera vår modell med den bästa tidigare utbildade modellen utan produktionsstopp eller förlust av modellnoggrannhet.

Eftersom GPU-modellen sparades med hjälp av vanliga PyTorch-modellsparverktyg kunde vi använda PyTorch kontrollpunktsladdningsverktyg för att ladda GPU-modellen på Trainium-enheter.

Till exempel, på GPU/CPU, kan du spara modellen med följande kod:

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

Sedan laddar du tillbaka modellen 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å samma sätt kan du spara modellen på Trainium med följande kod:

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

Och ladda tillbaka modellen på GPU/CPU:

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

I själva verket, eftersom vi använder DDP för modellträning, är modellladdningen agnostisk för antalet maskiner som används för att träna den tidigare kontrollpunkten. Detta tillåter oss att horisontellt skala Trn1-flottan utan kodändringar eller negativa effekter på modellträning. Dessa PyTorch-baserade kontrollpunkter kan användas direkt eller till och med fackla-script för slutledningsanvändning på AWS Inferentia2 eller andra acceleratorer.

Driftstabilitet

Det kan inte understrykas nog att körning av arbetsbelastningar i produktionen kräver att flera SLA:er uppfylls. För vårt användningsfall, förutom SLA:erna för modellkvalitet och utbildningsgenomströmning, är det absolut nödvändigt att produktionspipelinen är driftsstabil, vilket innebär minimala stillestånd och störningar under modellträning, utvärdering och slutledning.

Precis som med den befintliga GPU-baserade pipelinen har vi lagt till många mekanismer för att göra pipelinen operationellt stabil. Innan vi startar modellträning kör vi flera hälsotester för att bedöma maskinernas hälsa. Dessa tester inkluderar i allmänhet enkla tensoroperationer för att verifiera tillståndet hos acceleratoranordningarna. Vi har observerat att för distribuerad utbildning är det viktigt att köra tester för att verifiera kollektiv kommunikation även mellan instanser. Vi använde NCCOM testsvit från Neuron SDK för att uppnå detta genom att köra en mängd olika operationer såsom all-gather, all-reduce och reduce-scatter.

Även efter att ha följt förslagen vi har nämnt har vi observerat att övergående problem är oundvikliga i alla pipeline, oavsett den underliggande acceleratorn. För att bygga upp motståndskraft i en träningspipeline rekommenderar vi att man bygger in mekanismer för att försöka igen för att lösa dessa potentiella problem. Vi använder AWS Batch automatiserade återförsök att prova jobb som stöter på ett övergående misslyckande under modellutbildning. Dessa omstarter kan bli kostsamma om ett misslyckande uppstår mot slutet av träningen. För att motverka detta problem har vi anpassat våra träningsskript för att ladda en tidigare tränad modellkontrollpunkt och fortsätta träningen från den punkten. Med denna funktionalitet kan vi aggressivt starta om misslyckade träningsjobb med minimal omkostnad.

Med dessa elasticitetsmekanismer på plats kunde vi uppnå 98.5 % framgångsfrekvenser för våra arbetsbelastningar på Trn1, jämförbara med vår befintliga GPU-pipeline framgångsfrekvens.

Resultat

För att validera noggrannheten hos våra modeller initierade vi två modeller från samma GPU-kontrollpunkt och tränade en på Trainium och den andra på en jämförbar GPU. Båda modellerna tränades med samma träningshyperparametrar. Datauppsättningen som används för beräkning av mätvärden är en datauppsättning som håller ut, och vi utvärderar modellens noggrannhet på denna datauppsättning var N globala steg. X-axeln är det globala steget och Y-axeln är modellens noggrannhet. Vi observerade mindre än 0.1 % skillnad i modellnoggrannheten vid varje punkt i följande graf.

Dessutom, för att utvärdera kostnadseffektiviteten för modellträningen, föredrar vi att jämföra väggklockans tid det tar att nå modellkonvergens. Vi tror att detta ger en mer praktisk syn på kostnadsbesparingar jämfört med åtgärder som kostnad per token, uppnådda FLOPS/dollar och andra faktorer. Med tanke på träningstiden för trn1.32xl och jämförbar Amazon Elastic Compute Cloud (Amazon EC2) har vi observerat att Trainium erbjuder upp till 30 % billigare kostnad för modellkonvergens.

Slutsats

Det finns många faktorer att ta hänsyn till när du utvärderar olika acceleratorer för dina DL-arbetsbelastningar. Några av de viktigaste är modellkvalitet, genomströmning, kostnad och tillgänglighet. Det är ytterst viktigt att säkerställa att din modellkvalitet och genomströmning inte offras baserat på den accelerator du väljer.

Tack vare vårt partnerskap och samarbete med Annapurna Neuron-teamet har Amazon Search M5-teamet kunnat spara upp till 30 % i kostnader genom att flytta till Trainium. Teamet kan använda Trainium och uppnå modellkvalitet och genomströmningsparitet med jämförbara acceleratorer på marknaden. Kontrollpunkters interoperabilitet och minimala kodändringar med stöd för XLA har gjort det möjligt för M5 att välja mellan flera acceleratorer för sina arbetsbelastningar. Detta har gjort det möjligt för M5-teamet att dra fördel av Trainiums stora beräkningskraft och bygga acceleratoragnostiska lösningar för att glädja Amazon.com-kunder. Från en operativ synpunkt har Trainium visat sig kunna stödja tier-1-tjänster i Amazon-skala. M5-teamet fortsätter att flytta fler arbetsbelastningar till Trainium för att tillhandahålla de bästa modellerna för Amazon till lägsta kostnad.

Sammanfattningsvis har M5-teamet kunnat utföra kostnadseffektiv ML-utbildning i produktionsklass genom att lägga till Trainium till flottan av acceleratorer. Vi uppmuntrar dig att ta en titt på Trainium och andra Neuron-enheter som AWS Inferentia för att skörda fördelarna med specialbyggt Amazon-kisel för ML-arbetsbelastningar. Kom igång enkelt med en av de många handledningarna med olika modeller, t.ex Llama 2, tillgänglig på Trainium.

Om författarna

Abhinandan Patni är senior mjukvaruingenjör på Amazon Search. Han fokuserar på att bygga system och verktyg för skalbar distribuerad djupinlärningsträning och realtidsinferens.

James Park är en lösningsarkitekt på Amazon Web Services. Han arbetar med Amazon.com för att designa, bygga och distribuera tekniklösningar på AWS och har ett särskilt intresse för AI och maskininlärning. På fritiden tycker han om att söka nya kulturer, nya upplevelser och att hålla sig uppdaterad med de senaste tekniktrenderna. Du kan hitta honom på LinkedIn.

Jerry Mannil är en mjukvaruingenjör på Amazon Search. Han arbetar med att förbättra effektiviteten, robustheten och skalbarheten hos den distribuerade utbildningsinfrastrukturen.

Ken Su är en mjukvaruingenjör på Amazon Search. Han arbetar med att förbättra träningseffektiviteten och skalbart distribuerat träningsarbetsflöde. Utanför jobbet gillar han vandring och tennis.

RJ är ingenjör inom Amazon. Han bygger och optimerar system för distribuerade system för utbildning och arbetar med att optimera adopterande system för att minska latensen för ML Inference. Utanför arbetet utforskar han att använda Generativ AI för att bygga matrecept.

• SEO-drivet innehåll och PR-distribution. Bli förstärkt idag.
• PlatoData.Network Vertical Generative Ai. Styrka dig själv. Tillgång här.
• PlatoAiStream. Web3 Intelligence. Kunskap förstärkt. Tillgång här.
• Platoesg. Kol, CleanTech, Energi, Miljö, Sol, Avfallshantering. Tillgång här.
• PlatoHealth. Biotech och kliniska prövningar Intelligence. Tillgång här.
• Källa: https://aws.amazon.com/blogs/machine-learning/how-amazon-search-m5-saved-30-for-llm-training-cost-by-using-aws-trainium/