Cum Amazon Search M5 a economisit 30% pentru costurile de formare LLM prin utilizarea AWS Trainium | Amazon Web Services

Timp de zeci de ani, Amazon a fost pionier și a inovat învățarea automată (ML), oferind experiențe încântătoare clienților săi. Încă din primele zile, Amazon a folosit ML pentru diverse cazuri de utilizare, cum ar fi recomandări de cărți, căutare și detectarea fraudelor. Similar cu restul industriei, progresele hardware-ului accelerat au permis echipelor Amazon să urmărească arhitecturi model folosind rețele neuronale și deep learning (DL).

Programul M5 din cadrul Amazon Search deține strategia de învățare prin descoperire pentru Amazon și construiește modele la scară largă în mai multe limbi, multi-locale, multi-entități, multitask și multi-modale, cum ar fi text, imagine și video. Programul M5 a servit înglobări universale și modele de fundație la scară largă pentru sute de echipe ML de pe Amazon, menținând în același timp controale stricte asupra optimizării costurilor. Pentru a realiza acest lucru, echipa M5 evaluează în mod regulat noi tehnici de reducere a costurilor.

La fel ca multe organizații ML, acceleratoarele sunt utilizate în mare măsură pentru a accelera instruirea și inferența DL. Când AWS a lansat acceleratoare special create cu prima lansare a Inferentia AWS în 2020, echipa M5 a început rapid utilizați-le pentru a implementa mai eficient sarcinile de producție, economisind atât costurile, cât și reducând latența. Anul trecut, AWS și-a lansat AWS Trainium acceleratoare, care optimizează performanța pe cost pentru dezvoltarea și construirea de modele DL de generație următoare. În această postare, discutăm despre modul în care M5 a reușit să reducă costul antrenării modelelor lor cu 30% și să împărtășim câteva dintre cele mai bune practici pe care le-am învățat pe parcurs.

Instanțele Trainium

Odată cu progresele în acceleratoarele create special, Amazon oferă, de asemenea, acceleratoare convingătoare sub formă de AWS Inferentia și Trainium. După cum sugerează și numele, aceste cipuri sunt optimizate pentru a depăși nevoile de inferență și, respectiv, sarcinile de lucru de antrenament. Pentru formarea la scară largă a modelelor de fundație care ating miliarde de parametri în dimensiune, Trainium Instanțele Trn1 și Trn1n sunt alegeri ideale datorită caracteristicilor lor. Instanțele Trn1 sunt alimentate de ultimă generație NeuronCore-v2, și au o cantitate copioasă de accelerator de calcul și memorie. Instanțele Trn1n pot fi, de asemenea, alese pentru o cantitate mai mare de lățime de bandă de rețea (1,600 Gbs), deci sunt potrivite ideal pentru antrenamentul performant, având în vedere optimizarea costurilor.

Pentru a utiliza acceleratoarele, aveți nevoie de un strat de software care să le susțină. Cu jetoane Trn și Inf, AWS Neuron SDK deblochează acceleratoarele Amazon create special cu ajutorul PyTorch XLA. PyTorch XLA convertește modul dornic al lui PyTorch în implementare bazată pe grafice în modul leneș. Aceste grafice sunt apoi utilizate și compilate în continuare pentru a fi utilizate cu acceleratorul. PyTorch Neuron (parte a SDK-ului Neuron) le permite utilizatorilor PyTorch să-și antreneze modelele pe Trainium NeuronCores cu câteva linii de cod.

Model și volum de muncă

Echipa M5 antrenează și implementează modele de bază și reprezentări universale pentru a ajuta diverse echipe de pe Amazon să aducă încântare Amazon.com Clienți. Un astfel de model este un model de codificare de text urmat de un perceptron multistrat (MLP) cu interacțiuni explicite sau implicite ale caracteristicilor definite de arhitectura rețelei neuronale cu sute de milioane de parametri antrenabili. Acest model este antrenat pe miliarde de jetoane și este folosit pentru a genera milioane de înglobări într-o setare offline de inferență în lot. Aceste încorporare sunt intrări către un serviciu Amazon de nivel 1 orientat către clienți.

Infrastructura pentru conducta de producție folosește Lot AWS cu strategii de cozi la cozile echitabile, folosind un cluster multi-nod trn1.32xlarge activat EFA ca calcul pentru antrenamentul modelului. Din punct de vedere funcțional, conducta de producție realizează instruirea incrementală a modelului, evaluarea modelului antrenat și inferența în lot offline pe modelul antrenat, toate folosind PyTorch ca bibliotecă DL de bază.

Goluri

Încântarea clienților noștri este un principiu primordial. Având în vedere natura conductei orientată către clienți, este esențial ca toate acordurile de nivel de serviciu (SLA) să fie îndeplinite fără regresii. Am identificat două criterii critice de acceptare pentru a adapta canalul nostru de producție GPU existent și pentru a-l trece la Trainium:

• Calitate model – Calitatea modelelor noastre afectează direct experiența clienților. Solicităm să existe o diferență de mai puțin de 0.1% în calitatea modelului între GPU și Trainium.
• Debitul de antrenament – Ne antrenăm periodic modelele în mod iterativ pentru a oferi clienților noștri cea mai proaspătă experiență. Solicităm ca convergența modelului să fie realizată într-o perioadă de timp predefinită (cum ar fi o săptămână) pentru a îndeplini SLA-urile noastre de producție.

În secțiunile următoare, împărtășim călătoria noastră de lucru înapoi pe baza acestor criterii și învățările noastre pentru a susține sarcinile de producție la scară Amazon.

Script de antrenament

Înainte de a începe cu antrenamentul modelului, trebuie să facem modificări la scriptul de antrenament pentru a-l face compatibil XLA. Având în vedere dimensiunea modelului, folosim date distribuite paralele (DDP) pentru a antrena modelul. DDP ne permite să creștem debitul de instruire a modelului prin mărirea numărului de mașini utilizate pentru a rula antrenamentul modelului, fără modificări de cod. Am urmat instrucțiunile furnizate în Tutorial de instruire Neuron PyTorch MLP pentru a adăuga constructe specifice XLA în scripturile noastre de antrenament. Aceste modificări de cod sunt ușor de implementat. Următoarele sunt câteva învățăminte tehnice semnificative din exercițiu care au îmbunătățit considerabil randamentul modelului nostru:

• Plasarea xm.mark_step() - xm.mark_step() compilează și rulează graficele de calcul adunate leneș. Invocarea mark_step de prea multe ori va duce la un număr mai mare de grafice mici, în timp ce invocarea lui de prea puține ori va duce la puține, dar grafice mari. În funcție de aplicația dvs., debitul și implementarea instruirii modelului dvs. vor varia în funcție de plasarea dvs xm.mark_step(). Implementarea noastră plasează unul xm.mark_step() după o pasă înainte și înapoi și una după pasul de optimizare.
• Încărcarea de date cu încărcătorul de dispozitive multiprocesare XLA – Acesta este un pas critic care poate fi ratat cu ușurință. Încărcătorul dispozitivului multiprocesare torch_xla.distributed.parallel_loader.MpDeviceLoader încarcă datele de antrenament pe fiecare dispozitiv XLA cu opțiuni pentru a preîncărca și suprapune încărcarea datelor cu rulările dispozitivului pentru a îmbunătăți debitul. De asemenea, încărcătorul dispozitivului invocă xm.mark_step() și, prin urmare, este capabil să construiască grafice pentru încărcarea datelor pe dispozitiv de la gazdă.

Compilare pentru Trainium

În mod tradițional, ciclul de dezvoltare a modelului cu GPU-uri implică efectuarea de modificări la modelul sau scriptul de antrenament și rularea directă a acestuia pe dispozitivul GPU. Acceleratoarele precum Trainium care folosesc XLA necesită un pas suplimentar înainte ca antrenamentul modelului să poată fi rulat pe accelerator. Graficele de calcul XLA pot fi rulate numai după ce au fost compilate. În general, există două moduri de a efectua această compilare: Ahead of Time (AOT), în care urmăriți și compilați mai întâi toate graficele și apoi le executați, sau Just In Time (JIT), unde graficele sunt urmărite, compilate și rulate așa cum sunt sunt întâlnite. SDK-ul Neuron oferă ambele din cutie. De obicei, compilarea AOT este efectuată mai întâi. Graficele sunt apoi rulate după această compilare. Dacă sunt întâlnite grafice noi, timpul de execuție Neuron invocă o compilație JIT înainte de a le rula. Pentru a realiza compilarea AOT, Neuron SDK oferă neuron_parallel_compile, un utilitar de compilare care extrage grafice dintr-o rulare de probă a scriptului de antrenament și realizează o compilare AOT paralelă.

Un aspect important al compilarii AOT este acela de a se asigura că nu sunt create noi grafice de calcul pe parcursul instruirii. O sursă de noi grafice de calcul (și, prin urmare, recompilări) sunt formele dinamice ale loturilor de antrenament în timpul antrenamentului modelului. Am descoperit că utilizarea formelor statice și a loturilor de dimensiuni fixe elimină compilarea timpului de antrenament și îmbunătățește considerabil randamentul antrenamentului fără niciun efect asupra preciziei modelului. Prin aplicarea unor astfel de constrângeri asupra antrenamentului, am observat că doar 4-5 pași de antrenament al modelului, un pas de validare a modelului și verificarea modelului o singură dată sunt necesare pentru urmărirea tuturor graficelor în timpul compilării AOT. Este important să rețineți că Neuron SDK evoluează în mod constant și, în viitor, va accepta și forme dinamice.

În plus, graficele compilate sunt stocate în Neuron Persistent Cache pe disc sau într-un Serviciul Amazon de stocare simplă (Amazon S3) găleată. Acest lucru este util în special pentru sarcinile de producție în care arhitectura modelului și configurația de antrenament nu se modifică. Prin urmare, cheltuielile generale de compilare sunt suportate o singură dată. Utilizarea memoriei cache este la fel de simplă ca și setarea unui flag de mediu:

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

Compilatorul Neuron oferă și trei opțiuni de optimizare la nivel de compilator (O1, O2, O3) pentru a echilibra timpul de compilare și debitul rulării modelului. O1 permite optimizări de bază ale graficului de calcul și minimizează timpul de compilare, O3 oferă un randament îmbunătățit de rulare a modelului cu prețul unui timp de compilare mai mare, iar O2 (opțiune implicită) este un echilibru între cele două. Pentru cazul nostru de utilizare, am folosit optimizarea O1 și am observat o reducere de 86% a timpului de compilare fără nicio modificare a valorilor de precizie a modelului, observând în același timp o reducere de aproximativ 5-7% a debitului comparativ cu optimizarea implicită (O2). În funcție de caz de utilizare, puteți alege diferite niveluri de optimizare.

Pentru a rezuma, am folosit următoarele steaguri pentru compilare:

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

Compatibilitate punct de control

Când compilarea este finalizată cu succes, putem continua să ne antrenăm modelele pe Trainium. După cum am menționat mai devreme, ne antrenăm treptat modelele, ceea ce înseamnă că încărcăm un punct de control al modelului antrenat anterior și continuăm antrenamentul cu date noi. PyTorch și PyTorch XLA permit tranziția fără probleme între acceleratoare prin interoperabilitatea punctelor de control. Având flexibilitatea de a trece între GPU și Trainium, ne-a permis să încărcăm fără probleme modelul GPU anterior și să ne antrenăm pe mașinile Trainium. Acest lucru a fost esențial pentru a ne asigura că ne putem inițializa modelul cu cel mai bun model antrenat anterior, fără nicio întrerupere a producției sau pierderi de precizie a modelului.

Deoarece modelul GPU a fost salvat folosind utilitarele standard de salvare a modelului PyTorch, am putut folosi utilitarul de încărcare punct de control PyTorch pentru a încărca modelul GPU pe dispozitivele Trainium.

De exemplu, pe GPU/CPU, puteți salva modelul cu următorul cod:

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

Apoi încărcați modelul înapoi pe 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)

În mod similar, puteți salva modelul pe Trainium cu următorul cod:

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

Și încărcați modelul înapoi pe GPU/CPU:

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

De fapt, deoarece folosim DDP pentru antrenamentul modelului, încărcarea modelului este independentă de numărul de mașini utilizate pentru antrenamentul punctului de control anterior. Acest lucru ne permite să scalam pe orizontală flota Trn1 fără modificări de cod sau efecte adverse asupra antrenamentului de model. Aceste puncte de control bazate pe PyTorch pot fi utilizate direct sau chiar scriptate cu torță pentru cazuri de utilizare a inferenței pe AWS Inferentia2 sau alte acceleratoare.

Stabilitatea operațională

Nu se poate sublinia suficient că rularea sarcinilor de lucru în producție necesită îndeplinirea mai multor SLA. Pentru cazul nostru de utilizare, în afară de SLA-urile privind calitatea modelului și debitul de instruire, este imperativ ca conducta de producție să fie stabilă din punct de vedere operațional, ceea ce înseamnă timpi de nefuncționare și întreruperi minime în timpul antrenării, evaluării și deducerii modelului.

Ca și în cazul conductei existente bazate pe GPU, am adăugat numeroase mecanisme pentru a face conducta stabilă din punct de vedere operațional. Înainte de a începe formarea modelelor, rulăm mai multe teste de sănătate pentru a evalua starea de sănătate a mașinilor. Aceste teste includ, în general, operații simple de tensor pentru a verifica starea de sănătate a dispozitivelor de accelerație. Am observat că pentru instruirea distribuită, este important să rulăm teste pentru a verifica și comunicarea colectivă între instanțe. Noi am folosit Suita de testare NCCOM de la Neuron SDK pentru a realiza acest lucru, rulând o varietate de operațiuni, cum ar fi all-gather, all-reduce și reduce-scatter.

Chiar și după ce am urmat sugestiile pe care le-am menționat, am observat că problemele tranzitorii sunt inevitabile în orice conductă, indiferent de acceleratorul subiacent. Pentru a construi reziliența în orice canal de instruire, vă recomandăm să construiți mecanisme de reîncercare pentru a rezolva aceste probleme potențiale. Folosim Reîncercări automate AWS Batch pentru a reîncerca joburile care întâmpină o eroare tranzitorie în timpul antrenamentului de model. Aceste reporniri pot fi costisitoare dacă se întâlnește un eșec spre sfârșitul antrenamentului. Pentru a contracara această problemă, ne-am adaptat scripturile de antrenament pentru a încărca un punct de control model antrenat anterior și pentru a continua antrenamentul din acel punct. Cu această funcționalitate, suntem capabili să repornim agresiv joburile de antrenament eșuate, cu un cost minim.

Cu aceste mecanisme de reziliență în vigoare, am reușit să obținem rate de succes de 98.5% pentru sarcinile noastre de lucru pe Trn1, comparabile cu ratele de succes existente în pipeline GPU.

REZULTATE

Pentru a valida acuratețea modelelor noastre, am inițializat două modele din același punct de control GPU și am antrenat unul pe Trainium și celălalt pe un GPU comparabil. Ambele modele au fost antrenate cu aceiași hiperparametri de antrenament. Setul de date utilizat pentru calcularea valorilor este un set de date holdout și evaluăm acuratețea modelului pe acest set de date la fiecare N pași la nivel global. Axa X este pasul global, iar axa Y este precizia modelului. Am observat o diferență mai mică de 0.1% în precizia modelului în fiecare punct din graficul următor.

Mai mult, pentru a evalua rentabilitatea antrenamentului modelului, preferăm să comparăm timpul necesar pentru ceasul de perete pentru a ajunge la convergența modelului. Credem că aceasta oferă o perspectivă mai practică a economiilor de costuri în comparație cu măsuri precum costul pe token, FLOPS/dolar atins și alți factori. Având în vedere timpul de antrenament al trn1.32xl și comparabil Cloud Elastic de calcul Amazon (Amazon EC2), am observat că Trainium oferă un cost cu până la 30% mai ieftin pentru convergența modelului.

Concluzie

Există mulți factori de luat în considerare atunci când evaluați diferite acceleratoare pentru sarcinile dvs. de lucru DL. Unele dintre cele mai importante sunt calitatea modelului, randamentul, costul și disponibilitatea. Este esențial să vă asigurați că calitatea și randamentul modelului dvs. nu sunt sacrificate în funcție de acceleratorul pe care îl alegeți.

Datorită parteneriatului și colaborării noastre cu echipa Annapurna Neuron, echipa Amazon Search M5 a reușit să economisească până la 30% din costuri prin mutarea la Trainium. Echipa este capabilă să utilizeze Trainium și să obțină calitatea modelului și paritatea debitului cu acceleratoarele comparabile de pe piață. Interoperabilitatea punctelor de control și modificările minime ale codului cu suport pentru XLA i-au permis lui M5 să aleagă între mai multe acceleratoare pentru sarcinile lor de lucru. Acest lucru a permis echipei M5 să profite de puterea mare de calcul a Trainium și să creeze soluții agnostice de accelerator pentru a încânta clienții Amazon.com. Din punct de vedere operațional, Trainium s-a dovedit capabil să suporte servicii de nivel 1 la scara Amazon. Echipa M5 continuă să mute mai multe sarcini de lucru în Trainium pentru a oferi cele mai bune modele pentru Amazon la cele mai mici costuri.

Pe scurt, echipa M5 a reușit să efectueze instruire ML rentabilă, la nivel de producție, adăugând Trainium la flota de acceleratoare. Vă încurajăm să aruncați o privire la Trainium și la alte dispozitive Neuron, cum ar fi AWS Inferentia, pentru a beneficia de beneficiile siliciului Amazon creat special pentru sarcinile de lucru ML. Începeți cu ușurință cu unul dintre numeroasele tutoriale cu modele diferite, cum ar fi Llama 2, disponibil pe Trainium.

Despre Autori

Abhinandan Patni este inginer software senior la Amazon Search. El se concentrează pe construirea de sisteme și instrumente pentru instruire de deep learning distribuită scalabil și inferență în timp real.

James Park este arhitect de soluții la Amazon Web Services. El lucrează cu Amazon.com pentru a proiecta, construi și implementa soluții tehnologice pe AWS și are un interes deosebit pentru AI și învățarea automată. În timpul liber, îi place să caute noi culturi, experiențe noi și să fie la curent cu cele mai recente tendințe tehnologice. Îl poți găsi pe LinkedIn.

Jerry Mannil este inginer software la Amazon Search. Lucrează la îmbunătățirea eficienței, robusteței și scalabilității infrastructurii de instruire distribuite.

Ken Su este inginer software la Amazon Search. El lucrează la îmbunătățirea eficienței antrenamentului și a fluxului de lucru scalabil al instruirii distribuite. În afara serviciului, îi plac drumețiile și tenisul.

RJ este inginer în cadrul Amazon. El construiește și optimizează sisteme pentru sisteme distribuite pentru instruire și lucrează la optimizarea sistemelor de adoptare pentru a reduce latența pentru ML Inference. În afara muncii, el explorează utilizarea AI generativă pentru a construi rețete alimentare.

• Distribuție de conținut bazat pe SEO și PR. Amplifică-te astăzi.
• PlatoData.Network Vertical Generative Ai. Împuterniciți-vă. Accesați Aici.
• PlatoAiStream. Web3 Intelligence. Cunoștințe amplificate. Accesați Aici.
• PlatoESG. carbon, CleanTech, Energie, Mediu inconjurator, Solar, Managementul deșeurilor. Accesați Aici.
• PlatoHealth. Biotehnologie și Inteligență pentru studii clinice. Accesați Aici.
• Sursa: https://aws.amazon.com/blogs/machine-learning/how-amazon-search-m5-saved-30-for-llm-training-cost-by-using-aws-trainium/