Kako je Amazon Search M5 prihranil 30 % pri stroških usposabljanja LLM z uporabo AWS Trainium | Spletne storitve Amazon

Amazon je desetletja pionir in inoviral strojno učenje (ML), ki svojim strankam prinaša čudovite izkušnje. Od najzgodnejših dni je Amazon uporabljal ML za različne primere uporabe, kot so priporočila knjig, iskanje in odkrivanje goljufij. Podobno kot v preostali industriji je napredek pospešene strojne opreme Amazonovim ekipam omogočil, da si prizadevajo za modelne arhitekture z uporabo nevronskih mrež in globokega učenja (DL).

Program M5 v okviru Amazon Search ima v lasti strategijo učenja odkrivanja za Amazon in gradi obsežne modele v večjezičnih, večlokalnih, večentitetnih, večopravilnih in večmodalnih modelih, kot so besedilo, slike in video. Program M5 streže univerzalne vdelave in obsežne modele temeljev stotinam ekip ML v Amazonu, hkrati pa ohranja strog nadzor nad optimizacijo stroškov. Da bi to dosegli, ekipa M5 redno ocenjuje nove tehnike za zmanjšanje stroškov.

Tako kot mnoge organizacije ML se tudi pospeševalniki večinoma uporabljajo za pospeševanje usposabljanja in sklepanja DL. Ko je AWS lansiral namenske pospeševalnike s prvo izdajo Sklepanje AWS leta 2020 je ekipa M5 hitro začela jih uporabite za učinkovitejšo razporeditev proizvodnih delovnih obremenitev, kar prihrani stroške in zmanjša zakasnitev. Lani je AWS lansiral svojo AWS Trainium pospeševalniki, ki optimizirajo zmogljivost glede na ceno za razvoj in izdelavo modelov DL naslednje generacije. V tej objavi razpravljamo o tem, kako je podjetju M5 uspelo znižati stroške za usposabljanje svojih modelov za 30 %, in delimo nekaj najboljših praks, ki smo se jih naučili na tej poti.

Primerki Trainium

Z napredkom namenskih pospeševalnikov Amazon ponuja tudi prepričljive pospeševalnike v obliki AWS Inferentia in Trainium. Kot pove njihova imena, so ti čipi optimizirani tako, da presegajo potrebe po sklepanju in vadbenih delovnih obremenitvah. Za obsežno usposabljanje modelov temeljev, ki dosegajo milijarde parametrov v velikosti, Trainium Primerki Trn1 in Trn1n so zaradi svojih lastnosti idealna izbira. Primerke Trn1 poganja najsodobnejša tehnologija NeuronCore-v2, in imajo veliko količino pospeševalnika izračuna in pomnilnika. Primerke Trn1n lahko izberete tudi za večjo količino omrežne pasovne širine (1,600 Gbs), zato so idealne za uspešno usposabljanje z optimizacijo stroškov v mislih.

Če želite uporabljati pospeševalnike, potrebujete programsko plast, ki jih podpira. S čipi Trn in Inf je AWS Neuron SDK odklene Amazonove namenske pospeševalnike s pomočjo PyTorch XLA. PyTorch XLA pretvori nestrpni način PyTorcha v implementacijo, ki temelji na grafih. Ti grafi se nato uporabijo in nadalje sestavijo za uporabo s pospeševalnikom. PyTorch Neuron (del Neuron SDK) omogoča uporabnikom PyTorcha, da usposobijo svoje modele na Trainium NeuronCores z nekaj vrsticami kode.

Model in delovna obremenitev

Ekipa M5 usposablja in uvaja temeljne modele in univerzalne predstavitve, da bi pomagala različnim ekipam po Amazonu pri razveseljevanju Amazon.com stranke. En tak model je model kodirnika besedila, ki mu sledi večplastni perceptron (MLP) z eksplicitnimi ali implicitnimi interakcijami funkcij, ki jih definira arhitektura nevronske mreže s stotinami milijonov parametrov, ki jih je mogoče učiti. Ta model je učen na milijardah žetonov in se uporablja za ustvarjanje milijonov vdelav v nastavitvah paketnega sklepanja brez povezave. Te vdelave so vhodni podatki za storitev Amazon Tier-1, usmerjeno v stranke.

Infrastruktura za uporabo proizvodnega cevovoda Serija AWS z pravične strategije čakalne vrste, z uporabo gruče trn1.32xlarge z več vozlišči, ki podpira EFA, kot izračun za usposabljanje modela. Funkcionalno produkcijski cevovod izvaja inkrementalno usposabljanje modela, vrednotenje usposobljenega modela in paketno sklepanje brez povezave na usposobljenem modelu, pri čemer vse uporablja PyTorch kot osnovno knjižnico DL.

Cilji

Navdušiti naše stranke je glavno načelo. Glede na naravo cevovoda, ki se sooča s strankami, je ključnega pomena, da so vsi sporazumi o ravni storitev (SLA) izpolnjeni brez regresij. Identificirali smo dve kritični merili sprejemljivosti za prilagoditev našega obstoječega proizvodnega cevovoda GPE in prehod na Trainium:

• Kakovost modela – Kakovost naših modelov neposredno vpliva na uporabniško izkušnjo. Zahtevamo, da mora biti med GPU in Trainiumom manj kot 0.1 % razlike v kakovosti modela.
• Prepustnost usposabljanja – Naše modele redno usposabljamo, da našim strankam zagotovimo najbolj svežo izkušnjo. Zahtevamo, da mora biti konvergenca modela dosežena v vnaprej določenem časovnem obdobju (kot je 1 teden), da bi izpolnili naše proizvodne SLA.

V naslednjih razdelkih delimo našo pot dela nazaj glede na ta merila in naše učenje za podporo proizvodnim delovnim obremenitvam v obsegu Amazon.

Skripta za usposabljanje

Preden začnemo z usposabljanjem modela, moramo spremeniti skript za usposabljanje, da bo skladen z XLA. Glede na velikost modela uporabljamo porazdeljeno podatkovno paralelo (DDP) za usposabljanje modela. DDP nam omogoča, da povečamo prepustnost usposabljanja modela s povečanjem števila strojev, ki se uporabljajo za izvajanje usposabljanja modela, brez sprememb kode. Sledili smo navodilom v Vadnica za usposabljanje Neuron PyTorch MLP za dodajanje konstruktov, specifičnih za XLA, v naše skripte za usposabljanje. Te spremembe kode je enostavno izvesti. Sledi nekaj pomembnih tehničnih spoznanj iz vaje, ki so močno izboljšala prepustnost našega modela:

• Postavitev xm.mark_step() - xm.mark_step() prevaja in izvaja leno zbrane računske grafe. Priklicevanje mark_step prevečkrat bo privedlo do večjega števila majhnih grafov, medtem ko bo premalokrat priklic vodilo do malo, a velikih grafov. Odvisno od vaše aplikacije se bo prepustnost in izvedba vašega modelnega usposabljanja razlikovala glede na vašo umestitev xm.mark_step(). Naša izvedba postavlja eno xm.mark_step() po prehodu naprej in nazaj ter enega po koraku optimizacije.
• Ovijanje nalagalnika podatkov z nalagalnikom večprocesnih naprav XLA – To je kritičen korak, ki ga zlahka spregledamo. Nalagalnik večprocesnih naprav torch_xla.distributed.parallel_loader.MpDeviceLoader naloži podatke o usposabljanju v vsako napravo XLA z možnostmi prednalaganja in prekrivanja nalaganja podatkov z zagoni naprave za izboljšanje prepustnosti. Prikliče tudi nalagalnik naprave xm.mark_step() in je zato sposoben zgraditi grafe za nalaganje podatkov v napravo iz gostitelja.

Kompilacija za Trainium

Tradicionalno razvojni cikel modela z grafičnimi procesorji vključuje spreminjanje modela ali skripta za usposabljanje in njegovo neposredno izvajanje na napravi grafične procesorske enote. Pospeševalniki, kot je Trainium, ki uporabljajo XLA, zahtevajo dodaten korak, preden lahko na pospeševalniku izvajate usposabljanje modela. Računske grafe XLA je mogoče zagnati šele, ko so že prevedeni. Na splošno obstajata dva načina za izvedbo te kompilacije: Ahead of Time (AOT), kjer najprej sledite in prevedete vse grafe in jih nato zaženete, ali Just In Time (JIT), kjer se grafi sledijo, sestavljajo in izvajajo, kot se se srečujejo. Neuron SDK ponuja oboje takoj po namestitvi. Običajno se najprej izvede prevajanje AOT. Grafi se nato izvajajo po tej kompilaciji. Če se najdejo novi grafi, izvajalno okolje Neuron prikliče kompilacijo JIT, preden jih zažene. Za izvedbo prevajanja AOT nudi Neuron SDK neuron_parallel_compile, pripomoček za prevajanje, ki izvleče grafe iz poskusnega zagona skripta za usposabljanje in izvede vzporedno prevajanje AOT.

Pomemben vidik kompilacije AOT je zagotoviti, da se med usposabljanjem ne ustvarijo novi računski grafi. Eden od virov novih računskih grafov (in s tem ponovnih kompilacij) so dinamične oblike učnih paketov med usposabljanjem modela. Ugotovili smo, da uporaba statičnih oblik in serij s fiksno velikostjo odpravlja zbiranje časa usposabljanja in močno izboljša prepustnost usposabljanja brez vpliva na natančnost modela. Z uveljavljanjem takšnih omejitev pri usposabljanju smo opazili, da je za sledenje vsem grafom med prevajanjem AOT potrebnih le 4–5 korakov usposabljanja modela, en korak validacije modela in enkratna kontrolna točka modela. Pomembno je omeniti, da se Neuron SDK nenehno razvija in bo v prihodnosti podpiral tudi dinamične oblike.

Poleg tega so sestavljeni grafi shranjeni v Neuron Persistent Cache na disku ali v Preprosta storitev shranjevanja Amazon (Amazon S3) vedro. To je še posebej uporabno za proizvodne delovne obremenitve, kjer se arhitektura modela in konfiguracija usposabljanja ne spremenita. Zato režijski stroški prevajanja nastanejo samo enkrat. Uporaba predpomnilnika je tako preprosta kot nastavitev zastavice okolja:

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

Prevajalnik Neuron ponuja tudi tri možnosti optimizacije na ravni prevajalnika (O1, O2, O3) za uravnoteženje časa prevajanja in pretoka modela. O1 omogoča osnovne optimizacije na računskem grafu in minimizira čas prevajanja, O3 zagotavlja izboljšano prepustnost izvajanja modela za ceno daljšega časa prevajanja, O2 (privzeta možnost) pa je ravnovesje med obema. Za naš primer uporabe smo uporabili optimizacijo O1 in opazili 86-odstotno zmanjšanje časa prevajanja brez sprememb meritev natančnosti modela, medtem ko smo opazili približno 5–7-odstotno zmanjšanje prepustnosti v primerjavi s privzeto optimizacijo (O2). Glede na primer uporabe lahko izberete različne stopnje optimizacije.

Če povzamemo, smo za prevajanje uporabili naslednje zastavice:

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

Združljivost kontrolnih točk

Ko je kompilacija uspešno končana, lahko nadaljujemo z usposabljanjem naših modelov na Trainiumu. Kot smo že omenili, naše modele postopno usposabljamo, kar pomeni, da naložimo predhodno usposobljeno kontrolno točko modela in nadaljujemo z usposabljanjem z novimi podatki. PyTorch in PyTorch XLA omogočata brezhibno prehajanje med pospeševalniki prek interoperabilnosti kontrolnih točk. Fleksibilnost premikanja med GPU in Trainiumom nam je omogočila nemoteno nalaganje prejšnjega modela GPE in usposabljanje na napravah Trainium. To je bilo ključnega pomena za zagotovitev, da lahko inicializiramo naš model z najboljšim predhodno usposobljenim modelom brez izpadov proizvodnje ali izgube natančnosti modela.

Ker je bil model GPE shranjen s standardnimi pripomočki za shranjevanje modelov PyTorch, smo lahko uporabili pripomoček za nalaganje kontrolnih točk PyTorch za nalaganje modela GPE v naprave Trainium.

Na primer, na GPU/CPU lahko shranite model z naslednjo kodo:

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

Nato model naložite nazaj v 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)

Podobno lahko shranite model v Trainium z naslednjo kodo:

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

In naložite model nazaj na GPE/CPE:

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

Pravzaprav, ker uporabljamo DDP za usposabljanje modela, je nalaganje modela neodvisno od števila strojev, uporabljenih za usposabljanje prejšnje kontrolne točke. To nam omogoča horizontalno povečanje flote Trn1 brez sprememb kode ali škodljivih učinkov na usposabljanje modela. Te kontrolne točke, ki temeljijo na PyTorchu, je mogoče neposredno uporabiti ali celo skriptirati z baklo za primere uporabe sklepanja v AWS Inferentia2 ali drugih pospeševalnikih.

Stabilnost delovanja

Ni dovolj poudariti, da izvajanje delovnih obremenitev v proizvodnji zahteva izpolnjevanje več SLA. Za naš primer uporabe je poleg SLA za kakovost modela in prepustnost usposabljanja nujno, da je proizvodni cevovod operativno stabilen, kar pomeni minimalne izpade in motnje med usposabljanjem modela, vrednotenjem in sklepanjem.

Tako kot pri obstoječem cevovodu, ki temelji na GPU, smo dodali številne mehanizme, da zagotovimo stabilnost cevovoda. Preden začnemo z usposabljanjem modela, izvedemo več testov razumnosti, da ocenimo zdravje strojev. Ti testi na splošno vključujejo preproste tenzorske operacije za preverjanje zdravja pospeševalnih naprav. Opazili smo, da je za porazdeljeno usposabljanje pomembno izvajati teste, da se preveri tudi kolektivna komunikacija med primerki. Uporabili smo Testna zbirka NCCOM iz Neuron SDK, da bi to dosegli, z izvajanjem različnih operacij, kot so all-gather, all-reduce in reduce-scatter.

Tudi po upoštevanju predlogov, ki smo jih omenili, smo opazili, da so prehodne težave neizogibne v katerem koli cevovodu, ne glede na osnovni pospeševalnik. Za vzpostavitev odpornosti v katerem koli cevovodu usposabljanja priporočamo vgradnjo mehanizmov za ponovni poskus za rešitev teh morebitnih težav. Uporabljamo AWS Batch samodejni ponovni poskusi za ponovni poskus opravil, ki naletijo na prehodno napako med usposabljanjem modela. Ti ponovni zagoni so lahko dragi, če se proti koncu usposabljanja pojavi napaka. Da bi odpravili to težavo, smo prilagodili naše skripte za usposabljanje, da naložijo predhodno usposobljeno kontrolno točko modela in nadaljujejo z usposabljanjem od te točke. S to funkcionalnostjo lahko agresivno znova zaženemo neuspela usposabljanja z minimalnimi stroški.

S temi vzpostavljenimi mehanizmi odpornosti smo lahko dosegli 98.5-odstotno stopnjo uspešnosti za naše delovne obremenitve na Trn1, kar je primerljivo z našimi obstoječimi stopnjami uspešnosti cevovoda GPU.

Rezultati

Da bi preverili točnost naših modelov, smo inicializirali dva modela z iste kontrolne točke GPE in enega usposobili na Trainiumu, drugega pa na primerljivem GPE. Oba modela sta bila usposobljena z enakimi hiperparametri usposabljanja. Nabor podatkov, ki se uporablja za izračun meritev, je nabor podatkov o zadržanju in ocenimo natančnost modela na tem naboru podatkov vsakih N globalnih korakov. Os X je globalni korak, os Y pa je natančnost modela. Opazili smo manj kot 0.1 % razlike v točnosti modela na vsaki točki v naslednjem grafu.

Poleg tega za oceno stroškovne učinkovitosti usposabljanja modela raje primerjamo čas stenske ure, ki je potreben za dosego konvergence modela. Menimo, da to zagotavlja bolj praktičen pogled na prihranke stroškov v primerjavi z ukrepi, kot so cena na žeton, doseženi FLOPS/dolar in drugi dejavniki. Upoštevajoč čas treninga trn1.32xl in primerljivo Amazonski elastični računalniški oblak (Amazon EC2) smo opazili, da Trainium ponuja do 30 % cenejše stroške za konvergenco modela.

zaključek

Pri ocenjevanju različnih pospeševalnikov za vaše delovne obremenitve DL je treba upoštevati veliko dejavnikov. Nekateri najpomembnejši so kakovost modela, prepustnost, cena in razpoložljivost. Najpomembneje je zagotoviti, da kakovost vašega modela in prepustnost nista žrtvovani glede na pospeševalnik, ki ga izberete.

Zahvaljujoč našemu partnerstvu in sodelovanju z ekipo Annapurna Neuron je ekipa Amazon Search M5 s prehodom na Trainium prihranila do 30 % stroškov. Ekipa lahko uporablja Trainium in doseže kakovost modela in pariteto prepustnosti s primerljivimi pospeševalniki na trgu. Interoperabilnost kontrolnih točk in minimalne spremembe kode s podporo za XLA so M5 omogočili izbiro med več pospeševalniki za svoje delovne obremenitve. To je ekipi M5 omogočilo, da izkoristi veliko računalniško moč Trainiuma in zgradi agnostične rešitve pospeševalnika, da navduši stranke Amazon.com. Z operativnega vidika se je izkazalo, da je Trainium sposoben podpirati storitve Tier-1 v obsegu Amazon. Ekipa M5 še naprej prenaša več delovnih obremenitev na Trainium, da zagotovi najboljše modele za Amazon po najnižjih stroških.

Če povzamemo, ekipa M5 je z dodajanjem Trainiuma v floto pospeševalnikov lahko izvedla stroškovno učinkovito usposabljanje ML proizvodnega razreda. Svetujemo vam, da si ogledate Trainium in druge naprave Neuron, kot je AWS Inferentia, da izkoristite prednosti namensko izdelanega Amazonovega silicija za delovne obremenitve ML. Začnite enostavno z eno od številnih vadnic, ki prikazujejo različne modele, kot je Llama 2, na voljo na Trainiumu.

O avtorjih

Abhinandan Patni je višji programski inženir pri Amazon Search. Osredotoča se na gradnjo sistemov in orodij za razširljivo porazdeljeno usposabljanje poglobljenega učenja in sklepanje v realnem času.

James Park je arhitekt rešitev pri Amazon Web Services. Sodeluje z Amazon.com pri oblikovanju, gradnji in uvajanju tehnoloških rešitev na AWS, posebej pa ga zanimata umetna inteligenca in strojno učenje. V prostem času uživa v iskanju novih kultur, novih izkušenj in spremljanju najnovejših tehnoloških trendov. Najdete ga na LinkedIn.

Jerry Mannil je programski inženir pri Amazon Search. Dela na izboljšanju učinkovitosti, robustnosti in razširljivosti infrastrukture porazdeljenega usposabljanja.

Ken Su je programski inženir pri Amazon Search. Dela na izboljšanju učinkovitosti usposabljanja in razširljivega porazdeljenega poteka usposabljanja. Izven službe ima rad pohodništvo in tenis.

RJ je inženir v Amazonu. Gradi in optimizira sisteme za porazdeljene sisteme za usposabljanje in dela na optimizaciji prevzemnih sistemov za zmanjšanje zakasnitve za sklepanje ML. Zunaj dela raziskuje uporabo Generative AI za izdelavo receptov za hrano.

• Distribucija vsebine in PR s pomočjo SEO. Okrepite se še danes.
• PlatoData.Network Vertical Generative Ai. Opolnomočite se. Dostopite tukaj.
• PlatoAiStream. Web3 Intelligence. Razširjeno znanje. Dostopite tukaj.
• PlatoESG. Ogljik, CleanTech, Energija, Okolje, sončna energija, Ravnanje z odpadki. Dostopite tukaj.
• PlatoHealth. Obveščanje o biotehnologiji in kliničnih preskušanjih. Dostopite tukaj.
• vir: https://aws.amazon.com/blogs/machine-learning/how-amazon-search-m5-saved-30-for-llm-training-cost-by-using-aws-trainium/