Hogyan takarított meg 5%-ot az Amazon Search M30 az LLM képzési költségeiből az AWS Trainium | használatával Amazon webszolgáltatások

Az Amazon évtizedek óta úttörő és innovatív gépi tanulás (ML), amely elragadó élményeket kínál ügyfeleinek. Az Amazon a kezdetektől fogva használta az ML-t különféle felhasználási esetekre, például könyvajánlatokra, keresésre és csalások felderítésére. Az iparág többi részéhez hasonlóan a gyorsított hardver fejlődése lehetővé tette az Amazon csapatai számára, hogy neurális hálózatokat és mély tanulást (DL) használva modellarchitektúrákat dolgozzanak ki.

Az Amazon Search M5 programja birtokolja az Amazon felfedezés tanulási stratégiáját, és nagyszabású modelleket épít fel többnyelvű, több helyi, több entitásos, többfeladatos és multimodális, például szöveges, képi és videós modelleken. Az M5 program univerzális beágyazásokat és nagyméretű alapozási modelleket szolgált ki több száz ML-csapat számára az Amazonon, miközben szigorú ellenőrzést tartott a költségoptimalizálás felett. Ennek érdekében az M5 csapata rendszeresen értékeli az új technikákat a költségek csökkentése érdekében.

Mint sok ML szervezet, a gyorsítókat nagyrészt a DL képzés és következtetések felgyorsítására használják. Amikor az AWS piacra dobta a célra épített gyorsítókat az első kiadással AWS Inferentia 2020-ban az M5 csapata gyorsan elkezdett használja őket a termelési munkaterhelések hatékonyabb telepítéséhezköltséget takarít meg és csökkenti a késleltetést. Tavaly az AWS elindította AWS Trainium gyorsítók, amelyek költségenként optimalizálják a teljesítményt a következő generációs DL-modellek fejlesztéséhez és építéséhez. Ebben a bejegyzésben megvitatjuk, hogy az M5 hogyan tudta 30%-kal csökkenteni a modelljeik képzésének költségeit, és megosztunk néhány bevált gyakorlatot, amelyet az út során tanultunk.

Trainium példányok

A célzott gyorsítók fejlődésével az Amazon lenyűgöző gyorsítókat is kínál az AWS Inferentia és a Trainium formájában. Ahogy a nevük is sugallja, ezek a chipek úgy vannak optimalizálva, hogy meghaladják a következtetések és a képzési munkaterhelések igényeit. A több milliárd paramétert elérő alapmodellek nagyszabású képzéséhez Trainium Trn1 és Trn1n példányok tulajdonságaik miatt ideális választás. A Trn1 példányok a legkorszerűbbek NeuronCore-v2, és rengeteg gyorsítóval és memóriával rendelkezik. A Trn1n példányok nagyobb hálózati sávszélességhez (1,600 Gbs) is választhatók, így ideálisak a költségoptimalizálást szem előtt tartó teljesítményedzésekhez.

A gyorsítók használatához szükség van egy szoftverrétegre, amely támogatja őket. A Trn és Inf chipekkel a AWS Neuron SDK feloldja az Amazon erre a célra épített gyorsítóit a PyTorch XLA segítségével. A PyTorch XLA átalakítja a PyTorch lelkes üzemmódját lusta módú gráfalapú megvalósítássá. Ezeket a grafikonokat ezután a rendszer felhasználja, és tovább fordítja a gyorsítóval való használatra. A PyTorch Neuron (a Neuron SDK része) lehetővé teszi a PyTorch-felhasználók számára, hogy néhány soros kóddal Trainium NeuronCores rendszeren tanítsák modelljeiket.

Modell és munkaterhelés

Az M5 csapata alapmodelleket és univerzális képviseleteket képez ki és alkalmaz, hogy segítse az Amazon különböző csapatait abban, hogy örömet szerezzenek Amazon.com vásárlók. Az egyik ilyen modell a szövegkódoló modell, amelyet egy többrétegű perceptron (MLP) követ, explicit vagy implicit jellemző kölcsönhatásokkal a neurális hálózati architektúra által meghatározott több százmillió betanítható paraméterrel. Ezt a modellt több milliárd tokenre tanítják, és több millió beágyazás létrehozására használják offline kötegelt következtetési beállításban. Ezek a beágyazások bemenetek egy, az ügyfelek felé irányuló Tier-1 Amazon szolgáltatáshoz.

A termelési csővezeték infrastruktúráját használja AWS köteg val vel fair share sorban állási stratégiák, egy EFA-kompatibilis többcsomópontos trn1.32xlarge fürtöt használva a modelltanítás számításaként. Funkcionálisan a termelési folyamat növekményes modell betanítást, a betanított modell értékelését és offline kötegelt következtetést hajt végre a betanított modellen, mindezt a PyTorch segítségével a mögöttes DL-könyvtárként.

Célok

Ügyfeleink megörvendeztetése az egyik legfontosabb alapelv. Tekintettel a folyamat ügyfélközpontú jellegére, kritikus fontosságú, hogy minden szolgáltatási szintű megállapodást (SLA) regresszió nélkül teljesítsenek. Két kritikus elfogadási feltételt határoztunk meg meglévő GPU-gyártási folyamatunk adaptálásához és Trainiumra való átállításához:

• Modell minőség – Modelljeink minősége közvetlenül befolyásolja a vásárlói élményt. Megköveteljük, hogy a GPU és a Trainium modell minősége között 0.1%-nál kevesebb különbség legyen.
• Képzési teljesítmény – Modelljeinket rendszeresen oktatjuk, hogy a legfrissebb élményt nyújthassuk ügyfeleinknek. Megköveteljük, hogy a modellkonvergenciát előre meghatározott időn belül (például 1 héten belül) el kell érni ahhoz, hogy teljesítsük a termelési szolgáltatásiszint-szerződéseinket.

A következő szakaszokban megosztjuk az e feltételektől való visszalépéssel kapcsolatos útjainkat, valamint az Amazon-méretű gyártási terhelések támogatásával kapcsolatos tanulságainkat.

Képzési forgatókönyv

Mielőtt elkezdené a modellképzést, módosítanunk kell a képzési szkriptet, hogy az XLA-kompatibilis legyen. Tekintettel a modell méretére, a modell betanításához elosztott párhuzamos adatokat (DDP) használunk. A DDP lehetővé teszi, hogy növeljük a modellképzés áteresztőképességét a modellképzés futtatásához használt gépek számának növelésével, kódmódosítás nélkül. Az utasításokat követtük Neuron PyTorch MLP képzési útmutató hogy XLA-specifikus konstrukciókat adjunk a képzési szkriptjeinkhez. Ezek a kódmódosítások egyszerűen végrehajthatók. Az alábbiakban felsorolunk néhány jelentős technikai tanulságot a gyakorlatból, amelyek nagymértékben javították modellünk teljesítményét:

• Az xm.mark_step() elhelyezése - xm.mark_step() összeállítja és lefuttatja a lustán összegyűjtött számítási gráfokat. Meghívás mark_step túl sokszor több kis grafikonhoz vezet, míg a túl kevés alkalom meghívása kevés, de nagy grafikonhoz vezet. Alkalmazásától függően a modellképzés teljesítménye és megvalósítása az Ön által elfoglalt helytől függően változik xm.mark_step(). Megvalósításunk egyet tesz xm.mark_step() egy előre és hátra lépés után, és egy az optimalizáló lépés után.
• Adatbetöltő burkolás XLA többfeldolgozó eszközbetöltővel – Ez egy kritikus lépés, amelyet könnyen el lehet hagyni. A többfeldolgozó eszköz betöltő torch_xla.distributed.parallel_loader.MpDeviceLoader minden XLA-eszközre betölti az edzési adatokat, és lehetőséget ad az előtöltésre és az adatbetöltés átfedésére az eszközfuttatásokkal az átviteli sebesség javítása érdekében. Az eszközbetöltő is meghív xm.mark_step() és ezért képes grafikonokat készíteni az adatoknak az eszközre történő betöltéséhez a gazdagépről.

Összeállítás a Trainium számára

Hagyományosan a GPU-kkal végzett modellfejlesztési ciklus magában foglalja a modell vagy a betanító szkript módosítását, és annak közvetlen futtatását a GPU-eszközön. Az XLA-t használó gyorsítóknak, mint például a Trainium, további lépésre van szükség, mielőtt a modelledzést le lehetne futtatni a gázpedálon. Az XLA számítási gráfok csak lefordításuk után futtathatók. Általában kétféleképpen hajthatja végre ezt az összeállítást: Ahead of Time (AOT), ahol először nyomon követi és összeállítja az összes grafikont, majd lefuttatja őket, vagy a Just In Time (JIT), ahol a grafikonok nyomon követése, összeállítása és futtatása történik. találkoznak. A Neuron SDK mindkettőt a dobozból biztosítja. Általában az AOT-összeállítást hajtják végre először. A grafikonok az összeállítás után futnak. Ha új gráfokkal találkozik, a Neuron futtatókörnyezet meghív egy JIT-fordítást, mielőtt futtatná őket. Az AOT-összeállítás végrehajtásához a Neuron SDK biztosítja neuron_parallel_compile, egy fordítási segédprogram, amely grafikonokat bont ki a betanító szkript próbafutásából, és párhuzamos AOT-fordítást hajt végre.

Az AOT összeállítás egyik fontos szempontja annak biztosítása, hogy a képzés során ne jöjjön létre új számítási grafikon. Az új számítási grafikonok (és így az újrafordítások) egyik forrása a képzési kötegek dinamikus alakzatai a modellképzés során. Azt találtuk, hogy a statikus formák és a rögzített méretű kötegek használata kiküszöböli a képzési idő összeállítását, és nagymértékben javítja az edzési teljesítményt anélkül, hogy ez befolyásolná a modell pontosságát. Az ilyen megszorítások betanításával azt tapasztaltuk, hogy mindössze 4-5 lépéses modelltanítás, egy lépés a modell validálása és a modell egyszeri ellenőrzése szükséges az összes grafikon nyomon követéséhez az AOT fordítás során. Fontos megjegyezni, hogy a Neuron SDK folyamatosan fejlődik, és a jövőben támogatni fogja a dinamikus formákat is.

Továbbá az összeállított grafikonokat a Neuron állandó gyorsítótár lemezen vagy egy Amazon egyszerű tárolási szolgáltatás (Amazon S3) vödör. Ez különösen hasznos az éles munkaterheléseknél, ahol a modell architektúrája és a képzési konfiguráció nem változik. Ezért az összeállítás általános költsége csak egyszer merül fel. A gyorsítótár használata olyan egyszerű, mint egy környezetjelző beállítása:

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

A Neuron fordító is hármat biztosít fordítói szintű optimalizálási lehetőségek (O1, O2, O3) a fordítási idő és a modell futási sebességének egyensúlyához. Az O1 lehetővé teszi az alapvető optimalizálást a számítási grafikonon, és minimálisra csökkenti a fordítási időt, az O3 jobb modellfutási teljesítményt biztosít magasabb fordítási idő árán, az O2 (alapértelmezett beállítás) pedig egyensúlyt teremt a kettő között. A mi felhasználási esetünkben az O1 optimalizálást használtuk, és 86%-os csökkenést figyeltünk meg a fordítási időben anélkül, hogy a modell pontossági mutatói változtak volna, miközben az átviteli sebességben körülbelül 5–7%-os csökkenést figyeltünk meg az alapértelmezett optimalizáláshoz (O2) képest. A használati esettől függően különböző optimalizálási szinteket választhat.

Összefoglalva, a következő zászlókat használtuk az összeállításhoz:

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

Ellenőrzőpont kompatibilitás

Ha az összeállítás sikeresen befejeződött, folytathatjuk modelljeink Trainiumon való betanítását. Ahogy korábban említettük, a modelljeinket fokozatosan oktatjuk, azaz betöltünk egy korábban betanított modellellenőrző pontot, és új adatokkal folytatjuk a képzést. A PyTorch és a PyTorch XLA zökkenőmentes átmenetet tesz lehetővé a gyorsítók között az ellenőrzőpontok együttműködésének köszönhetően. A GPU és a Trainium közötti mozgás rugalmassága lehetővé tette számunkra, hogy zökkenőmentesen betölthessük az előző GPU-modellt, és a Trainium gépeken oktassunk. Ez kulcsfontosságú volt annak biztosításához, hogy modellünket a legjobb, korábban betanított modellel inicializálhassuk anélkül, hogy a gyártási leállás vagy a modell pontossága csökkenne.

Mivel a GPU-modellt szabványos PyTorch modellmentő segédprogramokkal mentettük, a PyTorch ellenőrzőpont-betöltő segédprogrammal tudtuk betölteni a GPU-modellt a Trainium eszközökön.

Például GPU/CPU esetén a modellt a következő kóddal mentheti el:

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

Ezután visszatölti a modellt a Trainiumra:

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)

Hasonlóképpen mentheti a modellt a Trainiumon a következő kóddal:

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

És töltse vissza a modellt GPU-ra/CPU-ra:

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

Valójában mivel DDP-t használunk a modellképzéshez, a modellbetöltés nem veszi figyelembe az előző ellenőrzőpont betanításához használt gépek számát. Ez lehetővé teszi számunkra, hogy vízszintesen méretezzük a Trn1 flottát anélkül, hogy kódváltozásokat vagy negatív hatásokat gyakorolna a modellképzésre. Ezek a PyTorch-alapú ellenőrzőpontok közvetlenül használhatók, vagy akár torch-scripttel is használhatók következtetések levonásához az AWS Inferentia2-n vagy más gyorsítókon.

Működési stabilitás

Nem lehet elégszer hangsúlyozni, hogy a munkaterhelések éles környezetben való futtatásához több SLA teljesítéséhez is szükség van. A mi felhasználási esetünkben a modellminőség és a betanítási teljesítmény SLA-kon kívül elengedhetetlen, hogy a gyártási folyamat stabil legyen, ami minimális leállást és zavarokat jelent a modell betanítása, értékelése és következtetése során.

A meglévő GPU-alapú folyamathoz hasonlóan számos mechanizmust adtunk hozzá, hogy a folyamat stabil működését biztosítsuk. A modellképzés megkezdése előtt több józansági tesztet is lefuttatunk, hogy felmérjük a gépek állapotát. Ezek a tesztek általában egyszerű tenzorműveleteket tartalmaznak a gyorsítóeszközök állapotának ellenőrzésére. Megfigyeltük, hogy az elosztott képzéshez fontos tesztek futtatása a példányok közötti kollektív kommunikáció ellenőrzésére is. Használtuk a NCCOM tesztcsomag a Neuron SDK-ból ennek eléréséhez, különféle műveletek futtatásával, mint például az összes összegyűjtés, az összes csökkentése és a szórás csökkentése.

Még az általunk említett javaslatok követése után is megfigyeltük, hogy a tranziens problémák elkerülhetetlenek minden folyamatban, függetlenül a mögöttes gyorsítótól. Bármely képzési folyamat rugalmasságának növelése érdekében javasoljuk az újrapróbálkozási mechanizmusok beépítését a lehetséges problémák megoldására. Használjuk AWS Batch automatikus újrapróbálkozások újrapróbálni azokat a feladatokat, amelyek átmeneti meghibásodásba ütköznek a modellképzés során. Ezek az újraindítások költségesek lehetnek, ha az edzés vége felé hiba történik. A probléma megoldása érdekében a képzési szkriptjeinket úgy alakítottuk át, hogy betöltsenek egy korábban betanított modellellenőrző pontot, és ettől kezdve folytassák a képzést. Ezzel a funkcióval minimális ráfordítással, agresszíven újraindíthatjuk a sikertelen képzési feladatokat.

Ezekkel a rugalmassági mechanizmusokkal 98.5%-os sikerességi arányt tudtunk elérni a Trn1-en végzett munkaterheléseink során, ami összevethető a meglévő GPU-folyamat sikerességi arányával.

Eredmények

Modelljeink pontosságának ellenőrzése érdekében két modellt inicializáltunk ugyanarról a GPU-ellenőrző pontról, és az egyiket Trainiumon, a másikat pedig egy hasonló GPU-n tanítottuk. Mindkét modellt ugyanazokkal a hiperparaméterekkel képezték ki. A metrikák kiszámításához használt adatkészlet egy visszatartási adatkészlet, és minden N globális lépésben értékeljük a modell pontosságát ezen az adatkészleten. Az X-tengely a globális lépés, az Y-tengely pedig a modell pontossága. A következő grafikon minden pontján kevesebb mint 0.1%-os eltérést figyeltünk meg a modell pontosságában.

Továbbá a modellképzés költséghatékonyságának értékeléséhez előnyben részesítjük a modellkonvergencia eléréséhez szükséges falióra-idő összehasonlítását. Úgy gondoljuk, hogy ez gyakorlatiasabb képet ad a költségmegtakarításokról, mint például a tokenenkénti költség, az elért FLOPS/dollár és egyéb tényezők. Figyelembe véve a trn1.32xl edzésidőt és összehasonlítható Amazon rugalmas számítási felhő (Amazon EC2) esetében megfigyeltük, hogy a Trainium akár 30%-kal olcsóbban kínálja a modellkonvergenciát.

Következtetés

Számos tényezőt kell figyelembe venni a DL-munkaterhelések különböző gyorsítóinak értékelésekor. A legfontosabbak közül néhány a modell minősége, áteresztőképessége, költsége és elérhetősége. Kiemelkedően fontos annak biztosítása, hogy a választott gyorsító miatt a modell minősége és átvitele ne menjen feláldozásra.

Az Annapurna Neuron csapattal való partnerségünknek és együttműködésünknek köszönhetően az Amazon Search M5 csapata akár 30%-ot is megtakaríthatott a Trainiumba való átköltözéssel. A csapat képes a Trainium használatára, és a piacon elérhető hasonló gyorsítókkal modellminőséget és átviteli paritást elérni. A Checkpoint interoperabilitása és az XLA támogatásával járó minimális kódmódosítások lehetővé tették az M5 számára, hogy több gyorsító közül válasszon a munkaterheléséhez. Ez lehetővé tette az M5 csapatának, hogy kihasználja a Trainium nagy számítási teljesítményét, és gyorsító agnosztikus megoldásokat építsen az Amazon.com ügyfeleinek örömére. Működési szempontból a Trainium bizonyítottan képes az Amazon szintű 1. szintű szolgáltatások támogatására. Az M5 csapata továbbra is több feladatot helyez át a Trainiumra, hogy a legjobb modelleket biztosítsa az Amazon számára a legalacsonyabb költségek mellett.

Összefoglalva, az M5-ös csapat költséghatékony, gyártási szintű ML-képzést tudott végrehajtani azáltal, hogy a Trainiummal bővítette a gyorsítók flottáját. Javasoljuk, hogy vessen egy pillantást a Trainiumra és más Neuron-eszközökre, például az AWS Inferentiára, hogy kihasználhassa az ML-munkaterhelések céljára épített Amazon-szilícium előnyeit. Kezdje el könnyedén a sok oktatóanyag egyikével, amelyek különböző modelleket tartalmaznak, például Llama 2, elérhető a Trainiumon.

A szerzőkről

Abhinandan Patni az Amazon Search vezető szoftvermérnöke. A skálázható, elosztott mély tanulási képzéshez és a valós idejű következtetésekhez szükséges rendszerek és eszközök kiépítésére összpontosít.

James Parker az Amazon Web Services megoldástervezője. Együttműködik az Amazon.com-tal az AWS technológiai megoldások tervezésében, kiépítésében és üzembe helyezésében, és különösen érdeklődik az AI és a gépi tanulás iránt. Szabadidejében szívesen keres új kultúrákat, új élményeket, és naprakész marad a legújabb technológiai trendekkel. Megtalálhatod rajta LinkedIn.

Jerry Mannil az Amazon Search szoftvermérnöke. Az elosztott képzési infrastruktúra hatékonyságának, robusztusságának és méretezhetőségének javításán dolgozik.

Ken Su az Amazon Search szoftvermérnöke. Az edzés hatékonyságának és a méretezhető, elosztott képzési munkafolyamatnak a javításán dolgozik. Munkán kívül szeret túrázni és tenisztezni.

RJ mérnök az Amazonon belül. Rendszereket épít és optimalizál az elosztott rendszerekhez a képzéshez, és dolgozik az elfogadó rendszerek optimalizálásán az ML következtetés késleltetésének csökkentése érdekében. A munkán kívül a Generative AI segítségével kutatja az ételreceptek készítését.

• SEO által támogatott tartalom és PR terjesztés. Erősödjön még ma.
• PlatoData.Network Vertical Generative Ai. Erősítse meg magát. Hozzáférés itt.
• PlatoAiStream. Web3 Intelligence. Felerősített tudás. Hozzáférés itt.
• PlatoESG. Carbon, CleanTech, Energia, Környezet, Nap, Hulladékgazdálkodás. Hozzáférés itt.
• PlatoHealth. Biotechnológiai és klinikai vizsgálatok intelligencia. Hozzáférés itt.
• Forrás: https://aws.amazon.com/blogs/machine-learning/how-amazon-search-m5-saved-30-for-llm-training-cost-by-using-aws-trainium/