Hoe Latent Space de parallellismebibliotheek van Amazon SageMaker gebruikte om de grenzen van grootschalige transformatoren te verleggen

Bronknooppunt: 1204406

Deze blog is mede geschreven door Sarah Jane Hong CSO, Darryl Barnhart CTO en Ian Thompson CEO van Latent Space en Prem Ranga van AWS.

Latente ruimte is een verborgen weergave van abstracte ideeรซn die machine learning (ML)-modellen leren. 'Hond', 'bloem' of 'deur' zijn bijvoorbeeld concepten of locaties in de latente ruimte. Bij Latente ruimte, werken we aan een engine waarmee je deze ruimte kunt manipuleren en verkennen met zowel taal als visuele aanwijzingen. Het Latent Space-team komt uit twee vakgebieden die elkaar al lang weinig overlappen: grafische afbeeldingen en natuurlijke taalverwerking (NLP). Traditioneel werden de modaliteiten van afbeeldingen en tekst afzonderlijk behandeld, elk met hun eigen geschiedenis van complexe, dure en fragiele feature-engineering. NLP-taken zoals het begrijpen van documenten of het beantwoorden van vragen hebben meestal weinig gemeen met visuele taken zoals het begrijpen of weergeven van scรจnes, en meestal gebruiken we zeer verschillende benaderingen en modellen voor elke taak. Maar dit is snel aan het veranderen.

Deze samenvoeging van modaliteiten in een enkele gedeelde latente ruimte ontgrendelt een nieuwe generatie creatieve en commerciรซle toepassingen, van gaming tot het begrijpen van documenten. Maar het ontsluiten van deze nieuwe toepassingen in รฉรฉn enkel model opent nieuwe schaaluitdagingen, zoals benadrukt in "The Bitter Lesson" door Richard Sutton, en het spannende werk van de afgelopen jaren op het gebied van schaalwetten. Om dit mogelijk te maken, werkt Latent Space aan baanbrekend onderzoek om deze modaliteiten samen te smelten in รฉรฉn model, maar ook om op te schalen en dit efficiรซnt te doen. Dit is waar modelparallelisme om de hoek komt kijken.

Amazon Sage Maker's unieke geautomatiseerde modelpartitionering en efficiรซnte pipelining-benadering maakten onze adoptie van modelparallelisme mogelijk met weinig technische inspanning, en we hebben onze training van modellen opgeschaald tot voorbij 1 miljard parameters (we gebruiken de p4d.24xgrote A100-exemplaren), wat voor ons een belangrijke vereiste is. Verder hebben we vastgesteld dat we bij het trainen met een trainingsopstelling van 16 knooppunten en acht GPU's met de parallellismebibliotheek van het SageMaker-model een verbetering van 38% in efficiรซntie registreerden in vergelijking met onze eerdere trainingsruns.

Uitdagingen bij het trainen van grootschalige transformatoren

Bij Latent Space combineren we taal en visie in transformatormodellen met miljarden parameters om gebruiksscenario's uit de verbeelding van een gebruiker te ondersteunen die "uit de distributie" zijn of die in de echte wereld zouden voorkomen, maar niet in onze trainingsgegevens. We gaan op twee verschillende manieren om met de uitdagingen die inherent zijn aan het opschalen naar miljarden parameters en meer:

Technieken voor het ophalen van informatie zijn lange tijd een belangrijk onderdeel geweest van zoekmachines en QA-taken. Onlangs is er opwindende vooruitgang geboekt door klassieke IR-technieken te combineren met moderne transformatoren, met name voor het beantwoorden van vragen, waarbij een model samen met een neurale retriever wordt getraind die leert relevante documenten op te halen om vragen te helpen beantwoorden. Zie voor een overzicht het recente werk van FAIR in Retrieval Augmented Generation: stroomlijnen van de creatie van intelligente modellen voor natuurlijke taalverwerking en Fusie-in-decoder, Google Brains RIJKen Nvidia's Neurale Retriever voor het beantwoorden van vragen.

Hoewel retrieval-augmented-technieken helpen bij kosten en efficiรซntie, kunnen we het model nog steeds niet op een enkele GPU passen voor ons grootste model. Dit betekent dat we modelparallellisme moeten gebruiken om het te trainen. Vanwege de aard van onze ophaalarchitectuur was het ontwerpen van onze modelsplitsing echter een uitdaging vanwege de onderlinge afhankelijkheden tussen opgehaalde contexten in trainingsinvoer. Bovendien, zelfs als we bepalen hoe we ons model splitsen, was het introduceren van modelparallellisme een belangrijke technische taak die we handmatig moesten uitvoeren tijdens onze onderzoeks- en ontwikkelingslevenscyclus.

De parallelle bibliotheek van het SageMaker-model

Modelparallellisme is het proces van het opsplitsen van een model over meerdere apparaten of knooppunten (zoals met GPU uitgeruste instanties) en het creรซren van een efficiรซnte pijplijn om het model op deze apparaten te trainen om het GPU-gebruik te maximaliseren. De model parallellisme bibliotheek in SageMaker maakt modelparallellisme toegankelijker door geautomatiseerde modelsplitsing te bieden, ook wel aangeduid als geautomatiseerde modelpartitionering en geavanceerde planning van pijplijnruns. De algoritmen voor het splitsen van modellen kunnen optimaliseren voor snelheid of geheugengebruik. De bibliotheek gebruikt een partitioneringsalgoritme dat het geheugen balanceert, de communicatie tussen apparaten minimaliseert en de prestaties optimaliseert.

Geautomatiseerde modelpartitionering

Voor onze PyTorch-gebruikscase voert de modelparallelle bibliotheek intern een traceerstap uit (in de eerste trainingsstap) die de modelgrafiek samenstelt en de tensor- en parametervormen bepaalt. Het construeert vervolgens een boom, die bestaat uit de geneste nn.Module objecten in het model, evenals aanvullende gegevens die zijn verzameld via tracering, zoals de hoeveelheid opgeslagen nn.Parameters, en looptijd voor elk nn.Module.

De bibliotheek doorloopt deze boom vervolgens vanaf de root en voert een partitioneringsalgoritme uit dat de rekenkracht en het geheugengebruik in evenwicht houdt, en de communicatie tussen instanties minimaliseert. Als meerdere nn.Modules dezelfde nn.Parameter delen, worden deze modules op hetzelfde apparaat geplaatst om te voorkomen dat meerdere versies van dezelfde parameter worden onderhouden. Nadat de indelingsbeslissing is genomen, worden de toegewezen modules en gewichten naar hun apparaten geladen.

Planning van pijplijnruns

Een ander kernkenmerk van de gedistribueerde modelparallelle bibliotheek van SageMaker is pijplijn loopt, die de volgorde bepalen waarin berekeningen worden gemaakt en gegevens worden verwerkt op verschillende apparaten tijdens modeltraining. Pipelining is gebaseerd op het splitsen van een mini-batch in microbatches, die รฉรฉn voor รฉรฉn in de trainingspijplijn worden ingevoerd en een uitvoeringsschema volgen dat is gedefinieerd door de looptijd van de bibliotheek.

De microbatch-pijplijn zorgt ervoor dat alle GPU's volledig worden benut, iets wat we zelf zouden moeten bouwen, maar met de modelparallelism-bibliotheek gebeurt dit netjes achter de schermen. Als laatste kunnen we gebruiken Amazon FSx, wat belangrijk is om ervoor te zorgen dat onze leessnelheden hoog zijn gezien het aantal bestanden dat wordt gelezen tijdens de training van een multimodaal model met ophalen.

architectuur trainen

Het volgende diagram geeft weer hoe we onze trainingsarchitectuur hebben opgezet. Onze primaire doelstellingen waren het verbeteren van de trainingssnelheid en het verlagen van de kosten. De beeld- en taaltransformatoren die we trainen, zijn zeer complex, met een aanzienlijk groot aantal lagen en gewichten erin, lopend naar miljarden parameters, waardoor ze allemaal niet in het geheugen van een enkel knooppunt passen. Elk knooppunt bevat een subset van het model, waardoor de gegevensstromen en de transformaties worden gedeeld en gecompileerd. Wij stellen 16 op p4d.24xgroot instanties met elk acht GPU's met behulp van de volgende architectuurrepresentatie:

Terwijl we onze modellen opschalen, is het een algemene trend dat alles wordt opgeslagen in de gewichten van het netwerk. Voor praktische doeleinden willen we onze modellen echter uitbreiden om te leren zoeken naar relevante contexten om te helpen bij de weergavetaak. Hierdoor kunnen we onze servicekosten laag houden zonder concessies te doen aan de beeldkwaliteit. We gebruiken een groot op transformatoren gebaseerd NLP-model en zoals eerder vermeld, zagen we een toename van 38% in trainingsefficiรซntie met de parallellismebibliotheek van het SageMaker-model, zoals blijkt uit het volgende:

  • We hebben een allreduce nodig voor elke berekening in het geval van parallellisme op tensorniveau. Hiervoor zijn O(log_2 n) parallelle stappen nodig. Dat zijn n machines die O(n) stappen nemen, voor O(n log_2 n) totale bewerkingen.
  • Voor pijplijnparallellisme hebben we O(1) parallelle stappen nodig voor het doorgeven van gegevens langs de pijplijn
  • Gegeven 16 machines met acht GPU's, hebben we O(1) kosten voor pijplijnparallel, en O(log_2(8)) = O(3) kosten voor dieptegewijs model parallel.
  • In dit geval zien we dat de netwerkkosten worden verlaagd tot 1/3e door over te schakelen naar parallelle pijplijn die we gebruiken met parallellisme van het SageMaker-model, en de totale trainingskosten worden verlaagd tot 1/2 + 1/2 * 1/log_2(16) = 0.625 van de oorspronkelijke kosten, wat leidt tot een overeenkomstige efficiรซntieverbetering.

Over het algemeen kunnen we, wanneer de behoefte gedistribueerde training rechtvaardigt (problemen met schaalmodelgrootte of trainingsgegevens), een reeks best practices volgen om te bepalen welke aanpak het beste werkt.

Best practices voor gedistribueerde training

Op basis van onze ervaring raden we aan om te beginnen met een gedistribueerde gegevensparallelle aanpak. Gedistribueerde gegevensparallelisme zoals de SageMaker gedistribueerde parallelle gegevensbibliotheek lost de meeste netwerkproblemen met modelreplica's op, dus u moet modellen in het kleinste aantal knooppunten passen en vervolgens naar behoefte repliceren om de batchgrootte te schalen.

Als je tijdens de training geen geheugen meer hebt, zoals we in dit scenario hebben gedaan, wil je misschien overschakelen naar een parallelle modelbenadering. Overweeg echter deze alternatieven voordat u parallelle modeltraining probeert:

  • Gebruik op NVIDIA Tensor Core-uitgeruste hardware gemengde precisie training om versnelling te creรซren en het geheugenverbruik te verminderen.
  • Verklein de batchgrootte (of verminder de beeldresolutie of de lengte van de NLP-reeks, indien mogelijk).

Bovendien geven we de voorkeur aan modelontwerpen die geen batchnormalisatie hebben zoals beschreven in Hoogwaardige grootschalige beeldherkenning zonder normalisatie. Als dit niet kan worden vermeden, zorg er dan voor dat batchnormalisatie op alle apparaten wordt gesynchroniseerd. Wanneer u gedistribueerde training gebruikt, wordt uw batch verdeeld over GPU's, dus nauwkeurige batchstatistieken vereisen synchronisatie op alle apparaten. Zonder dit zal de normalisatie een grotere fout hebben en daardoor de convergentie belemmeren.

Begin met parallelle modeltraining als u de volgende beperkingen hebt:

  • Je model past niet op รฉรฉn apparaat
  • Vanwege de grootte van uw model wordt u geconfronteerd met beperkingen bij het kiezen van grotere batchgroottes, bijvoorbeeld als uw modelgewichten het grootste deel van uw GPU-geheugen in beslag nemen en u gedwongen wordt een kleinere, suboptimale batchgrootte te kiezen

Doe het volgende bij het optimaliseren voor prestaties:

  • Gebruik pipelining voor communicatie tussen knooppunten om latentie te minimaliseren en doorvoer te verhogen
  • Houd pijpleidingen zo kort mogelijk om eventuele luchtbellen te minimaliseren. Het aantal microbatches moet worden afgestemd om de rekenefficiรซntie in evenwicht te brengen met de bellengrootte, en moet ten minste de lengte van de pijplijn zijn. Indien nodig kunt u microbatches vormen op tokenniveau zoals beschreven in TeraPipe: Token Level Pipeline Parallellisme voor het trainen van grootschalige taalmodellen

Gebruik bij het optimaliseren voor kosten door SageMaker beheerde Spot Instances voor training. Dit kan de kosten van trainingsmodellen tot 90% optimaliseren ten opzichte van On-Demand-exemplaren. SageMaker beheert de Spot-onderbrekingen namens u.

Andere factoren om te overwegen:

  • Binnen een knooppunt als er een snelle verbinding is, is het genuanceerder. Als er voldoende intra-node netwerkcapaciteit is, kan het herschikken van gegevens voor meer optimale rekenkracht een voordeel opleveren.
  • Als activaties veel groter zijn dan gewicht tensoren, kan een sharded optimizer ook helpen. Raadpleeg Nul voor meer details.

De volgende tabel bevat enkele veelvoorkomende scenario's voor het opschalen van trainingen en hoe u deze op AWS kunt configureren.

Scenario Wanneer is het van toepassing? Oplossing
Schalen van een enkele GPU naar meerdere GPU's Wanneer de hoeveelheid trainingsgegevens of de grootte van het model te groot is Stap over op een instantie met meerdere GPU's, zoals p3.16xlarge, die acht GPU's heeft, waarbij de gegevens en verwerking zijn verdeeld over de acht GPU's en een bijna lineaire versnelling produceert in de tijd die nodig is om uw model te trainen.
Schalen van een enkele instantie naar meerdere instanties Wanneer de schaalbehoeften verder gaan dan het wijzigen van de instantiegrootte Schaal het aantal instanties met de schattingsfunctie van de SageMaker Python SDK door uw instance_type in te stellen op p3.16xlarge en instance_count op 2. In plaats van de acht GPU's op een enkele p3.16xlarge, heeft u 16 GPU's over twee identieke instanties. Overweeg het gebruik van de SageMaker gedistribueerde parallelle gegevensbibliotheek.
Selecteren van een modelparallelle benadering voor training Wanneer u tijdens de training onvoldoende geheugenfouten tegenkomt Schakel over naar een modelparallelle benadering met behulp van de SageMaker gedistribueerde model parallelle bibliotheek.
Netwerkprestaties voor communicatie tussen knooppunten Voor gedistribueerde training met meerdere instanties (bijvoorbeeld communicatie tussen de knooppunten in het cluster bij het uitvoeren van een AllReduce-bewerking) Uw instanties moeten zich in dezelfde regio en dezelfde beschikbaarheidszone bevinden. Wanneer u de SageMaker Python SDK gebruikt, wordt dit voor u afgehandeld. Uw trainingsgegevens moeten zich ook in dezelfde beschikbaarheidszone bevinden. Overweeg het gebruik van de SageMaker gedistribueerde parallelle gegevensbibliotheek.
Geoptimaliseerde GPU, netwerk en opslag Voor grootschalige gedistribueerde trainingsbehoeften Het instantietype p4d.24xlarge is ontworpen voor snelle lokale opslag en een snelle netwerkbackplane met maximaal 400 gigabit, en we raden het ten zeerste aan als de meest performante optie voor gedistribueerde training.

Conclusie

Met de parallelle modelbibliotheek in SageMaker halen we veel van de voordelen uit de doos, zoals geautomatiseerde modelpartitionering en efficiรซnte pipelining. In dit bericht deelden we onze uitdagingen met onze ML-use-case, onze overwegingen over verschillende trainingsbenaderingen en hoe we de parallellismebibliotheek van het Amazon SageMaker-model gebruikten om onze training te versnellen. Het beste van alles is dat het nu slechts een paar uur kan duren om de hier beschreven best practices voor modelparallelisme en prestatieverbeteringen toe te passen. Als dit bericht je helpt of inspireert om een โ€‹โ€‹probleem op te lossen, horen we dat graag! Deel uw opmerkingen en feedback.

Referenties

Zie het volgende voor meer informatie:


Over de auteurs

Prem Ranga is een Enterprise Solutions Architect gevestigd in Atlanta, GA. Hij maakt deel uit van de Machine Learning Technical Field Community en werkt graag met klanten samen aan hun ML- en AI-traject. Prem is gepassioneerd door robotica, is een onderzoeker naar autonome voertuigen en bouwde ook de Alexa-gestuurde Beer Pours in Houston en andere locaties.

Sarah Jane Hong is de mede-oprichter en Chief Science Officer bij Latent Space. Haar achtergrond ligt op het snijvlak van mens-computerinteractie en machine learning. Eerder leidde ze NLP-onderzoek bij Sonar (overgenomen door Marchex), dat bedrijven bedient op het gebied van conversatie-AI. Ze is ook een gewaardeerde AR/VR-ontwikkelaar en heeft onderscheidingen en beurzen ontvangen van Oculus, Mozilla Mixed Reality en Microsoft Hololens.

Darryl Barnhart is de mede-oprichter en Chief Technology Officer bij Latent Space. Hij is een doorgewinterde ontwikkelaar met ervaring in GPU-versnelling, computergraphics, grootschalige data en machine learning. Andere passies zijn wiskunde, game-ontwikkeling en de studie van informatie.

Ian Thompson is de oprichter en CEO van Latent Space. Ian is een ingenieur en onderzoeker die zich laat inspireren door het 'aangrenzende mogelijke': technologieรซn die op het punt staan โ€‹โ€‹een grote impact op ons leven te hebben. Momenteel gericht op het vereenvoudigen en schalen van multimodaal representatief leren om veilige en creatieve AI te helpen bouwen. Eerder bouwde hij mee aan bedrijven in graphics/virtual reality (AltspaceVR, overgenomen door Microsoft) en onderwijs/NLP (HSE).

Bron: https://aws.amazon.com/blogs/machine-learning/how-latent-space-used-the-amazon-sagemaker-model-parallelism-library-to-push-the-frontiers-of-large-scale-transformers/

Tijdstempel:

Meer van AWS Blog over machine learning