Recente ontwikkelingen in deep learning hebben geleid tot steeds grotere modellen zoals GPT-3, BLOOM en OPT, waarvan sommige al meer dan 100 miljard parameters bevatten. Hoewel grotere modellen over het algemeen krachtiger zijn, vereist het trainen van dergelijke modellen aanzienlijke rekenkracht. Zelfs met het gebruik van geavanceerde gedistribueerde trainingsbibliotheken zoals FSDP en DeepSpeed, is het normaal dat voor trainingstaken honderden acceleratieapparaten nodig zijn voor meerdere weken of maanden achter elkaar.
Eind 2022 kondigde AWS de algemene beschikbaarheid aan van Amazon EC2 Trn1-instanties Mogelijk gemaakt door AWS Trainium—een speciaal gebouwde machine learning (ML)-versneller die is geoptimaliseerd om een krachtig, kosteneffectief en enorm schaalbaar platform te bieden voor het trainen van deep learning-modellen in de cloud. Trn1-instanties zijn verkrijgbaar in een aantal formaten (zie de volgende tabel), met maximaal 16 Trainium-versnellers per instantie.
| Instantiegrootte | Trainium-versnellers | Versnellergeheugen (GB) | vCPU's | Instantiegeheugen (GiB) | Netwerkbandbreedte (Gbps) |
| trn1.2xgroot | 1 | 32 | 8 | 32 | Tot 12.5 |
| trn1.32xgroot | 16 | 512 | 128 | 512 | 800 |
| trn1n.32xlarge (binnenkort beschikbaar) | 16 | 512 | 128 | 512 | 1600 |
Trn1-instanties kunnen worden ingezet als zelfstandige instanties voor kleinere trainingstaken, of in zeer schaalbare ultraclusters die gedistribueerde training over tienduizenden Trainium-versnellers ondersteunen. Alle Trn1-instanties ondersteunen de zelfstandige configuratie, terwijl Trn1-ultraclusters trn1.32xlarge- of trn1n.32xlarge-instanties vereisen. In een ultracluster bevinden meerdere Trn1-instanties zich op dezelfde locatie in een bepaalde AWS-beschikbaarheidszone en zijn verbonden met high-speed, low-latency, Elastic Fabric Adapter (EFA)-netwerken die 800 Gbps niet-blokkerende netwerkbandbreedte per instantie bieden voor collectieve computerbewerkingen . Het instantietype trn1n.32xlarge, dat begin 2023 wordt gelanceerd, verhoogt deze bandbreedte tot 1600 Gbps per instantie.
Veel zakelijke klanten kiezen ervoor om hun deep learning-workloads te implementeren met behulp van Kubernetes, de de facto standaard voor containerorkestratie in de cloud. AWS-klanten implementeren deze workloads vaak met behulp van Amazon Elastic Kubernetes-service (Amazone EKS). Amazon EKS is een beheerde Kubernetes-service die het maken, configureren, levenscyclus en monitoren van Kubernetes-clusters vereenvoudigt en tegelijkertijd de volledige flexibiliteit van upstream Kubernetes biedt.
Vandaag kondigen we met trots officiële ondersteuning aan voor gedistribueerde trainingstaken met behulp van Amazon EKS- en EC2 Trn1-instanties. Met deze aankondiging kunt u nu eenvoudig grootschalige gecontaineriseerde trainingstaken uitvoeren binnen Amazon EKS, terwijl u volledig profiteert van de prijs-prestatieverhouding, schaalbaarheid en het gebruiksgemak van Trn1-instanties.
Samen met deze aankondiging publiceren we ook een gedetailleerde zelfstudie die u door de stappen leidt die nodig zijn om een gedistribueerde trainingstaak met meerdere instanties (BERT fase 1 pre-training) uit te voeren met behulp van Amazon EKS- en Trn1-instanties. In dit bericht leer je over de oplossingsarchitectuur en bekijk je verschillende belangrijke stappen uit de zelfstudie. Verwijs naar de officiële tutorial-repository voor de volledige end-to-end workflow.
Om mee te gaan, een brede bekendheid met kern AWS-services zoals Amazon Elastic Compute-cloud (Amazon EC2) en Amazon EKS wordt geïmpliceerd, en basiskennis van diep leren en PyTorch zou nuttig zijn.
Oplossingsarchitectuur
Het volgende diagram illustreert de oplossingsarchitectuur.
De oplossing bestaat uit de volgende hoofdcomponenten:
- Een EKS-cluster
- Een EKS-knooppuntgroep bestaande uit trn1.32xlarge instanties
- De AWS-neuron SDK
- EKS-plug-ins voor Neuron en EFA
- An Amazon Elastic Container-register (Amazon ECR) Bewaarplaats
- Een afbeelding van een trainings container
- An Amazon FSx voor Luster bestandssysteem
- Een Volcano-batchplanner en etcd-server
- De TorchX universele job launcher
- De TorchX DDP-module voor Trainium
De kern van de oplossing wordt gevormd door een EKS-cluster dat u de kernfunctionaliteit van Kubernetes-beheer biedt via een EKS-service-eindpunt. Een van de voordelen van Amazon EKS is dat de service het besturingsvlak actief bewaakt en schaalt op basis van belasting, wat zorgt voor hoge prestaties voor grote werklasten, zoals gedistribueerde training. Binnen het EKS-cluster bevindt zich een knooppuntgroep die bestaat uit twee of meer op trn1.32xlarge Trainium gebaseerde instanties die zich in dezelfde beschikbaarheidszone bevinden.
De Neuron SDK is de softwarestack die de driver, compiler, runtime, framework-integratie (bijvoorbeeld PyTorch Neuron) en gebruikerstools biedt waarmee u toegang krijgt tot de voordelen van de Trainium-versnellers. Het stuurprogramma van het Neuron-apparaat werkt rechtstreeks op de EKS-knooppunten (Trn1-instanties) en biedt toegang tot de Trainium-chips vanuit de trainingscontainers die op de knooppunten worden gelanceerd. Neuron- en EFA-plug-ins zijn geïnstalleerd binnen het EKS-cluster om toegang te bieden tot de Trainium-chips en EFA-netwerkapparaten die nodig zijn voor gedistribueerde training.
Er wordt een ECR-repository gebruikt om de afbeeldingen van de trainingscontainer op te slaan. Deze afbeeldingen bevatten de Neuron SDK (exclusief het Neuron-stuurprogramma, dat rechtstreeks op de Trn1-instanties wordt uitgevoerd), het PyTorch-trainingsscript en de vereiste afhankelijkheden. Wanneer een trainingstaak op het EKS-cluster wordt gestart, worden de containerafbeeldingen eerst van Amazon ECR naar de EKS-knooppunten getrokken en worden de PyTorch-werkcontainers vervolgens vanuit de afbeeldingen geïnstantieerd.
Gedeelde opslag wordt geleverd met behulp van een krachtig FSx for Lustre-bestandssysteem dat zich in dezelfde beschikbaarheidszone bevindt als de trn1.32xlarge-instanties. Het maken en koppelen van het FSx for Luster-bestandssysteem aan het EKS-cluster wordt gemedieerd door de Amazon FSx voor Luster CSI-stuurprogramma. In deze oplossing wordt de gedeelde opslag gebruikt voor het opslaan van de trainingsdataset en alle logboeken of artefacten die tijdens het trainingsproces zijn gemaakt.
De oplossing gebruikt de TorchX universele job launcher om gedistribueerde trainingsbanen binnen Amazon EKS te lanceren. TorchX heeft twee belangrijke afhankelijkheden: de Volcano-batchplanner en de etcd-server. Volcano verzorgt de planning en wachtrij van trainingstaken, terwijl de etcd-server een sleutel-waardeopslag is die door TorchElastic wordt gebruikt voor synchronisatie en peer-ontdekking tijdens het opstarten van taken.
Wanneer een trainingstaak wordt gestart met TorchX, gebruikt het startcommando de meegeleverde door TorchX gedistribueerde DDP-module voor Trainium om de algehele trainingstaak te configureren en vervolgens de juiste torchrun-opdrachten uit te voeren op elk van de PyTorch-werkpods. Wanneer een taak wordt uitgevoerd, kan deze worden gecontroleerd met behulp van standaard Kubernetes-tools (zoals kubectl) of via standaard ML-toolsets zoals TensorBoard.
Overzicht oplossingen
Laten we eens kijken naar de belangrijke stappen van deze oplossing. In dit overzicht verwijzen we naar de Start een Multi-Node PyTorch Neuron Training Job op Trainium met behulp van TorchX en EKS zelfstudie op GitHub.
Maak een EKS-cluster
Om aan de slag te gaan met gedistribueerde trainingstaken in Amazon EKS met Trn1-instanties, maakt u eerst een EKS-cluster zoals beschreven in de zelfstudie op GitHub. Clustercreatie kan worden bereikt met behulp van standaardtools zoals eksctl en AWS CloudFormatie.
Maak een EKS-knooppuntgroep
Vervolgens moeten we een EKS-knooppuntgroep maken met twee of meer trn1.32xlarge-instanties in een ondersteunde regio. In de zelfstudie AWS CloudFormation wordt gebruikt om een Trainium-specifieke EC2-lanceringssjabloon te maken, wat ervoor zorgt dat de Trn1-instanties worden gestart met een geschikte Amazon Machine Image (AMI) en de juiste EFA-netwerkconfiguratie die nodig is om gedistribueerde training te ondersteunen. De AMI bevat ook het Neuron-apparaatstuurprogramma dat ondersteuning biedt voor de Trainium-versnellerchips. Met de eksctl Amazon EKS-beheertool, kunt u eenvoudig een Trainium-knooppuntgroep maken met behulp van een eenvoudig YAML-manifest dat verwijst naar de nieuw gemaakte startsjabloon. Bijvoorbeeld:
In het voorgaande manifest zijn verschillende kenmerken geconfigureerd om het gebruik van Trn1-instanties in het EKS-cluster mogelijk te maken. Eerst, metadata.region is ingesteld op een van de regio's die Trn1-instanties ondersteunt (momenteel us-east-1 en us-west-2). Vervolgens vereist Amazon EKS voor beschikbaarheidszones dat er twee beschikbaarheidszones worden opgegeven. Een van deze beschikbaarheidszones moet het gebruik van Trn1-instanties ondersteunen, terwijl de andere willekeurig kan worden gekozen. De tutorial laat zien hoe bepalen welke beschikbaarheidszones Trn1-instanties binnen uw AWS-account toestaan. Dezelfde Trn1-ondersteunende beschikbaarheidszone moet ook worden opgegeven met behulp van de availabiltyZones attribuut gekoppeld aan de EKS-knooppuntgroep. efaEnabled is ingesteld op true om de knooppunten te configureren met de juiste EFA-netwerkconfiguratie die vereist is voor gedistribueerde training. Als laatste de launchTemplate.id attribuut gekoppeld aan de knooppuntgroep verwijst naar de EC2-startsjabloon die in een eerdere stap via AWS CloudFormation is gemaakt.
Ervan uitgaande dat u de CloudFormation-sjabloon al hebt toegepast en de eksctl beheertool, kunt u een EKS-knooppuntgroep maken die geschikt is voor Trainium door de volgende code uit te voeren:
Installeer Kubernetes-plug-ins voor Trainium- en EFA-apparaten
Met de knooppuntengroep is de volgende stap het installeren van Kubernetes-plug-ins die ondersteuning bieden voor de Trainium-versnellers (via de Neuron-plug-in) en de EFA-apparaten (via de EFA-plug-in). Deze plug-ins kunnen met behulp van de standaard eenvoudig op het cluster worden geïnstalleerd kubectl beheertool zoals weergegeven in de zelfstudie.
Om de TorchX universele PyTorch-starter te gebruiken om gedistribueerde trainingstaken te starten, zijn twee vereisten vereist: de Volcano-batchplanner en de etcd-server. Net als de Neuron- en EFA-plug-ins kunnen we de kubectl tool om Volcano en de etcd-server op het EKS-cluster te installeren.
Koppel gedeelde opslag aan het EKS-cluster
In de zelfstudie wordt FSx for Luster gebruikt om een krachtig gedeeld bestandssysteem te bieden dat toegankelijk is voor de verschillende EKS-werknemerspods. Deze gedeelde opslag wordt gebruikt om de trainingsgegevensset te hosten, evenals eventuele artefacten en logboeken die tijdens het trainingsproces worden gemaakt. In de zelf studie wordt beschreven hoe u de gedeelde opslag maakt en koppelt aan het cluster met behulp van de Amazon FSx voor Luster CSI-stuurprogramma.
Maak een installatiekopie van een trainingscontainer
Vervolgens moeten we een afbeelding van een trainingscontainer maken die het PyTorch-trainingsscript en eventuele afhankelijkheden bevat. Een voorbeeld van Dockerfile is opgenomen in de zelfstudie, waarin het BERT pre-trainingsscript is opgenomen, samen met de bijbehorende softwareafhankelijkheden. De Dockerfile wordt gebruikt om de afbeelding van de trainingscontainer te bouwen en de afbeelding wordt vervolgens naar een ECR-repository gepusht van waaruit de PyTorch-medewerkers de afbeelding kunnen ophalen wanneer een trainingstaak op het cluster wordt gestart.
Stel de trainingsgegevens in
Voordat een trainingstaak wordt gestart, worden de trainingsgegevens eerst gekopieerd naar het gedeelde opslagvolume op FSx for Lustre. In de zelf studie wordt beschreven hoe u een tijdelijke Kubernetes-pod maakt die toegang heeft tot het gedeelde opslag volume, en hoe u zich aanmeldt bij de pod om de trainings gegevensset te downloaden en te extraheren met behulp van standaard Linux-shell-opdrachten.
Nu de verschillende infrastructuur- en softwarevereisten aanwezig zijn, kunnen we ons nu concentreren op de Trainium-aspecten van de oplossing.
Precompileer uw model
De Neuron SDK ondersteunt PyTorch via een integratielaag genaamd PyTorch-neuron. PyTorch Neuron werkt standaard met just-in-time-compilatie, waarbij de verschillende rekengrafieken van het neurale netwerk binnen een trainingstaak worden gecompileerd zoals ze tijdens het trainingsproces worden aangetroffen. Voor grotere modellen kan het handiger zijn om de meegeleverde te gebruiken neuron_parallel_compile tool om de verschillende rekengrafieken van tevoren te precompileren en in de cache op te slaan om compilatie van grafieken tijdens de training te voorkomen. Voordat de trainingstaak op het EKS-cluster wordt gestart, laat de tutorial zien hoe u eerst een precompilatietaak start via TorchX met behulp van de neuron_parallel_compile hulpmiddel. Na voltooiing van de precompilatietaak zal de Neuron-compiler alle rekengrafieken van het neurale netwerk hebben geïdentificeerd en gecompileerd, en deze in de cache hebben opgeslagen op het gedeelde opslagvolume voor later gebruik tijdens de eigenlijke BERT-pre-trainingstaak.
Start de gedistribueerde trainingstaak
Nadat de precompilatie is voltooid, wordt TorchX vervolgens gebruikt om een gedistribueerde trainingstaak van 64 werknemers te starten over twee trn1.32xlarge instanties, met 32 werknemers per instantie. We gebruiken 32 werkers per instantie omdat elke instantie van trn1.32xlarge 16 Trainium-versnellers bevat, waarbij elke versneller er 2 levert Neuronen kernen. Elke NeuronCore is uniek toegankelijk PyTorch XLA-apparaat in het trainingsschema. Een voorbeeld van een TorchX-startopdracht uit de zelfstudie ziet eruit als de volgende code:
De verschillende opdrachtregelargumenten in het voorgaande TorchX-commando worden in detail beschreven in de zelfstudie. De volgende argumenten zijn echter het belangrijkst bij het configureren van de trainingstaak:
- -cfg wachtrij=test – Specificeert de Volcano-wachtrij die moet worden gebruikt voor de trainingstaak
- -cfg afbeelding_repo - Specificeert de ECR-repository die moet worden gebruikt voor de TorchX-containerafbeeldingen
- –script_args - Specificeert alle argumenten die moeten worden doorgegeven aan het PyTorch-trainingsscript
- –nnodes en –nproc_per_node – Het aantal exemplaren en werkrollen per exemplaar dat moet worden gebruikt voor de trainingstaak
- -script – De naam van het PyTorch-trainingsscript dat moet worden gestart in de trainingscontainer
- -Beeld – Het pad naar de afbeelding van de trainingscontainer in Amazon ECR
- –bf16 – Het al dan niet inschakelen van het gegevenstype BF16
Bewaken van de opleidingsopdracht
Nadat de leeropdracht is gelanceerd, zijn er verschillende manieren waarop de functie kan worden gevolgd. De zelfstudie laat zien hoe u basisgegevens over trainingsscripts op de opdrachtregel kunt volgen met behulp van kubectl, hoe u de voortgang van het trainingsscript in TensorBoard visueel kunt volgen (zie de volgende schermafbeelding) en hoe u het gebruik van de Trainium-versneller kunt volgen met behulp van de neuron-top tool van de Neuron SDK.
De omgeving opruimen of hergebruiken
Wanneer de trainingstaak is voltooid, kan het cluster opnieuw worden gebruikt of opnieuw worden geconfigureerd voor extra trainingstaken. De EKS-knooppuntgroep kan bijvoorbeeld snel worden opgeschaald met behulp van de eksctl opdracht om trainingstaken te ondersteunen waarvoor extra Trn1-exemplaren nodig zijn. Evenzo kunnen de meegeleverde Dockerfile- en TorchX-startopdrachten eenvoudig worden aangepast om aanvullende deep learning-modellen te ondersteunen en trainingstopologieën te distribueren.
Als het cluster niet langer nodig is, bevat de zelf studie ook alle stappen die nodig zijn om de EKS-infrastructuur en gerelateerde bronnen te verwijderen.
Conclusie
In dit bericht hebben we onderzocht hoe Trn1-instanties en Amazon EKS een beheerd platform bieden voor krachtige, kosteneffectieve en enorm schaalbare gedistribueerde training van deep learning-modellen. We hebben ook een uitgebreide tutorial gedeeld die laat zien hoe je een real-world multi-instance gedistribueerde trainingstaak uitvoert in Amazon EKS met behulp van Trn1-instances, en benadrukte verschillende van de belangrijkste stappen en componenten in de oplossing. De inhoud van deze zelfstudie kan eenvoudig worden aangepast voor andere modellen en workloads, en biedt u een basisoplossing voor gedistribueerde training van deep learning-modellen in AWS.
Raadpleeg de Neuron documentatie.
Over de auteurs
Scott Perry is een Solutions Architect in het Annapurna ML-acceleratorteam bij AWS. Hij is gevestigd in Canada en helpt klanten bij het implementeren en optimaliseren van deep learning-trainingen en deductieworkloads met behulp van AWS Inferentia en AWS Trainium. Zijn interesses omvatten grote taalmodellen, diep versterkend leren, het internet der dingen en genomica.
Lorea Arrizabalaga is een Solutions Architect gericht op de Britse publieke sector, waar ze klanten helpt bij het ontwerpen van ML-oplossingen met Amazon SageMaker. Ze maakt ook deel uit van de Technical Field Community die zich toelegt op hardwareversnelling en helpt bij het testen en benchmarken van AWS Inferentia- en AWS Trainium-workloads.
- Door SEO aangedreven content en PR-distributie. Word vandaag nog versterkt.
- Platoblockchain. Web3 Metaverse Intelligentie. Kennis versterkt. Toegang hier.
- Bron: https://aws.amazon.com/blogs/machine-learning/scaling-distributed-training-with-aws-trainium-and-amazon-eks/
