Distribuerad träning med Amazon EKS och Torch Distributed Elastic

Utbildning för distribuerade djupinlärningsmodeller blir allt viktigare eftersom datastorlekarna växer i många branscher. Många applikationer inom datorseende och naturlig språkbehandling kräver nu träning av modeller för djupinlärning, som växer exponentiellt i komplexitet och ofta tränas med hundratals terabyte data. Det blir då viktigt att använda en stor molninfrastruktur för att skala utbildningen av så stora modeller.

Utvecklare kan använda ramverk med öppen källkod som PyTorch för att enkelt designa intuitiva modellarkitekturer. Men att skala utbildningen av dessa modeller över flera noder kan vara utmanande på grund av ökad orkestreringskomplexitet.

Distribuerad modellträning består huvudsakligen av två paradigm:

• Modell parallell – I modellparallellträning är själva modellen så stor att den inte får plats i minnet på en enda GPU, och flera GPU:er behövs för att träna modellen. Open AI:s GPT-3-modell med 175 miljarder träningsbara parametrar (cirka 350 GB i storlek) är ett bra exempel på detta.
• Data parallell – Vid dataparallellträning kan modellen ligga i en enda GPU, men eftersom datan är så stor kan det ta dagar eller veckor att träna en modell. Att distribuera data över flera GPU-noder kan minska träningstiden avsevärt.

I det här inlägget ger vi ett exempel på arkitektur för att träna PyTorch-modeller med hjälp av Torch Distributed Elastisk ramverk på ett distribuerat dataparallellt sätt med hjälp av Amazon Elastic Kubernetes-tjänst (Amazon EKS).

Förutsättningar

För att replikera resultaten som rapporteras i det här inlägget är den enda förutsättningen ett AWS-konto. I detta konto skapar vi ett EKS-kluster och ett Amazon FSx för Luster filsystem. Vi skjuter även containerbilder till en Amazon Elastic Container Registry (Amazon ECR) arkiv på kontot. Instruktioner för att ställa in dessa komponenter tillhandahålls efter behov under hela inlägget.

EKS-kluster

Amazon EKS är en hanterad containertjänst för att köra och skala Kubernetes-applikationer på AWS. Med Amazon EKS kan du effektivt köra distribuerade utbildningsjobb med det senaste Amazon Elastic Compute Cloud (Amazon EC2)-instanser utan att behöva installera, driva och underhålla ditt eget kontrollplan eller noder. Det är en populär orchestrator för maskininlärning (ML) och AI-arbetsflöden. Ett typiskt EKS-kluster i AWS ser ut som följande figur.

Vi har släppt ett projekt med öppen källkod, AWS DevOps för EKS (aws-do-eks), som tillhandahåller en stor samling lättanvända och konfigurerbara skript och verktyg för att tillhandahålla EKS-kluster och köra distribuerade utbildningsjobb. Detta projekt är byggt enligt principerna för Gör ramverk: Enkelhet, flexibilitet och universalitet. Du kan konfigurera ditt önskade kluster genom att använda eks.conf filen och starta den sedan genom att köra eks-create.sh manus. Detaljerade instruktioner finns i GitHub repo.

Träna PyTorch-modeller med Torch Distributed Elastic

Torch Distributed Elastic (TDE) är ett inbyggt PyTorch-bibliotek för utbildning av storskaliga modeller för djupinlärning där det är avgörande att skala beräkningsresurser dynamiskt baserat på tillgänglighet. De TorchElastic Controller för Kubernetes är en inbyggd Kubernetes-implementering för TDE som automatiskt hanterar livscykeln för de poddar och tjänster som krävs för TDE-utbildning. Det möjliggör dynamisk skalning av beräkningsresurser under träning efter behov. Den ger också feltolerant utbildning genom att återställa jobb från nodfel.

I det här inlägget diskuterar vi stegen för att träna PyTorch EfficientNet-B7 och ResNet50 modeller som använder IMAGEnet data på ett distribuerat sätt med TDE. Vi använder PyTorch DistribueradDataParallell API och Kubernetes TorchElastic-kontrollern och kör våra träningsjobb på ett EKS-kluster som innehåller flera GPU-noder. Följande diagram visar arkitekturdiagrammet för denna modellutbildning.

TorchElastic for Kubernetes består huvudsakligen av två komponenter: TorchElastic Kubernetes Controller (TEC) och parameterservern (etcd). Styrenheten ansvarar för att övervaka och hantera träningsjobben, och parameterservern håller reda på arbetarnoderna för distribuerad synkronisering och peer-upptäckt.

För att träningspodarna ska komma åt datan behöver vi en delad datavolym som kan monteras av varje pod. Några alternativ för delade volymer genom Gränssnitt för containerlagring (CSI) drivrutiner ingår i AWS DevOps för EKS är Amazon Elastic File System (Amazon EFS) och FSx för Luster.

Klusterinställningar

I vår klusterkonfiguration använder vi en c5.2xlarge-instans för systempods. Vi använder tre p4d.24xlarge-instanser som worker pods för att träna en EfficientNet-modell. För ResNet50-utbildning använder vi p3.8xlarge-instanser som worker pods. Dessutom använder vi ett FSx delat filsystem för att lagra vår träningsdata och modellartefakter.

AWS p4d.24xlarge instanser är utrustade med Elastisk tygadapter (EFA) för att tillhandahålla nätverk mellan noder. Vi diskuterar EFA mer senare i inlägget. För att möjliggöra kommunikation via EFA måste vi konfigurera klusterinställningen genom en .yaml-fil. En exempelfil finns i GitHub-förrådet.

Efter att denna .yaml-fil är korrekt konfigurerad kan vi starta klustret med skriptet som tillhandahålls i GitHub-repo:

./eks-create.sh

Referera till GitHub repo för detaljerade instruktioner.

Det är praktiskt taget ingen skillnad mellan att köra jobb på p4d.24xlarge och p3.8xlarge. Stegen som beskrivs i det här inlägget fungerar för båda. Den enda skillnaden är tillgängligheten av EFA på p4d.24xlarge-instanser. För mindre modeller som ResNet50 har standardnätverk jämfört med EFA-nätverk minimal inverkan på träningshastigheten.

FSx för Luster filsystem

FSx är designad för högpresterande datorarbetsbelastningar och ger en fördröjning på under millisekunder med hjälp av lagringsvolymer för solid-state-enheter. Vi valde FSx eftersom det gav bättre prestanda när vi skalade till ett stort antal noder. En viktig detalj att notera är att FSx bara kan existera i en enda tillgänglighetszon. Därför bör alla noder som har åtkomst till FSx-filsystemet finnas i samma tillgänglighetszon som FSx-filsystemet. Ett sätt att uppnå detta är att specificera den relevanta tillgänglighetszonen i klustrets .yaml-fil för de specifika nodgrupperna innan klustret skapas. Alternativt kan vi modifiera nätverksdelen av den automatiska skalningsgruppen för dessa noder efter att klustret har konfigurerats, och begränsa det till att använda ett enda subnät. Detta kan enkelt göras på Amazon EC2-konsolen.

Förutsatt att EKS-klustret är igång och subnät-ID för tillgänglighetszonen är känt, kan vi ställa in ett FSx-filsystem genom att tillhandahålla nödvändig information i fsx.conf fil enligt beskrivningen i readme och kör deploy.sh manus i FSX mapp. Detta ställer in rätt policy och säkerhetsgrupp för åtkomst till filsystemet. Skriptet installerar också CSI-drivrutinen för FSx som en demonset. Slutligen kan vi skapa FSx-beständiga volymanspråket i Kubernetes genom att använda en enda .yaml-fil:

kubectl apply -f fsx-pvc-dynamic.yaml

Detta skapar ett FSx-filsystem i tillgänglighetszonen som anges i fsx.conf fil och skapar också ett beständigt volymanspråk fsx-pvc, som kan monteras av vilken som helst av kapslarna i klustret på ett läs-skriv-många (RWX) sätt.

I vårt experiment använde vi komplett ImageNet-data, som innehåller mer än 12 miljoner träningsbilder uppdelade i 1,000 XNUMX klasser. Data kan laddas ner från ImageNet hemsida. Den ursprungliga TAR-kulan har flera kataloger, men för vår modellutbildning är vi bara intresserade av ILSVRC/Data/CLS-LOC/, som inkluderar train och val underkataloger. Innan träning måste vi ordna om bilderna i val underkatalog för att matcha katalogstrukturen som krävs av PyTorch Bildmapp klass. Detta kan göras med en enkel Python-manus efter att data har kopierats till den beständiga volymen i nästa steg.

Att kopiera data från en Amazon enkel lagringstjänst (Amazon S3) till FSx-filsystemet skapar vi en Docker-bild som innehåller skript för denna uppgift. Ett exempel på Dockerfile och ett skalskript ingår i CSI mapp i GitHub-repo. Vi kan bygga bilden med hjälp av build.sh skriptet och skjut det sedan till Amazon ECR med hjälp av push.sh manus. Innan vi använder dessa skript måste vi tillhandahålla rätt URI för ECR-förrådet i .env fil i rotmappen för GitHub-repo. När vi har skickat Docker-bilden till Amazon ECR kan vi starta en pod för att kopiera data genom att använda den relevanta .yaml-filen:

kubectl apply -f fsx-data-prep-pod.yaml

Podden kör skriptet automatiskt data-prep.sh för att kopiera data från Amazon S3 till den delade volymen. Eftersom ImageNet-data har mer än 12 miljoner filer tar kopieringsprocessen ett par timmar. Python-skriptet imagenet_data_prep.py körs också för att ordna om val dataset som förväntat av PyTorch.

Nätverksacceleration

Vi kan använda Elastic Fabric Adapter (EFA) i kombination med stödde EC2-instanstyper för att påskynda nätverkstrafiken mellan GPU-noderna i ditt kluster. Detta kan vara användbart när du kör stora distribuerade utbildningsjobb där standardnätverkskommunikation kan vara en flaskhals. Skript för att distribuera och testa EFA-enhetsplugin i EKS-klustret som vi använder här ingår i efa-device-plugin mapp i GitHub-repo. För att aktivera ett jobb med EFA i ditt EKS-kluster, förutom att klusternoderna har den nödvändiga hårdvaran och mjukvaran, måste EFA-enhetspluginen distribueras till klustret, och din jobbbehållare måste ha kompatibla CUDA och NCCL versioner installerad.

För att demonstrera att köra NCCL-tester och utvärdera prestandan för EFA på p4d.24xlarge-instanser måste vi först distribuera Kubeflow MPI-operatören genom att köra motsvarande deploy.sh manus i mpi-operatör mapp. Sedan kör vi deploy.sh skript och uppdatera test-efa-nccl.yaml uppenbara så gränser och önskemål om resurs vpc.amazonaws.com är inställda på 4. De fyra tillgängliga EFA-adaptrarna i p4d.24xlarge-noderna buntas ihop för att ge maximal genomströmning.

Körning kubectl apply -f ./test-efa-nccl.yaml för att tillämpa testet och sedan visa loggarna för testpodden. Följande rad i loggutgången bekräftar att EFA används:

NCCL INFO NET/OFI Selected Provider is efa

Testresultaten bör se ut som följande utdata:

[1,0]<stdout>:#                                                       out-of-place                       in-place
[1,0]<stdout>:#       size         count      type   redop     time   algbw   busbw  error     time   algbw   busbw  error
[1,0]<stdout>:#        (B)    (elements)                       (us)  (GB/s)  (GB/s)            (us)  (GB/s)  (GB/s)
[1,0]<stdout>:           8             2     float     sum    629.7    0.00    0.00  2e-07    631.4    0.00    0.00  1e-07
[1,0]<stdout>:          16             4     float     sum    630.5    0.00    0.00  1e-07    628.1    0.00    0.00  1e-07
[1,0]<stdout>:          32             8     float     sum    627.6    0.00    0.00  1e-07    628.2    0.00    0.00  1e-07
[1,0]<stdout>:          64            16     float     sum    633.6    0.00    0.00  1e-07    628.4    0.00    0.00  6e-08
[1,0]<stdout>:         128            32     float     sum    627.5    0.00    0.00  6e-08    632.0    0.00    0.00  6e-08
[1,0]<stdout>:         256            64     float     sum    634.5    0.00    0.00  6e-08    636.5    0.00    0.00  6e-08
[1,0]<stdout>:         512           128     float     sum    634.8    0.00    0.00  6e-08    635.2    0.00    0.00  6e-08
[1,0]<stdout>:        1024           256     float     sum    646.6    0.00    0.00  2e-07    643.6    0.00    0.00  2e-07
[1,0]<stdout>:        2048           512     float     sum    745.0    0.00    0.01  5e-07    746.0    0.00    0.01  5e-07
[1,0]<stdout>:        4096          1024     float     sum    958.2    0.00    0.01  5e-07    955.8    0.00    0.01  5e-07
[1,0]<stdout>:        8192          2048     float     sum    963.0    0.01    0.02  5e-07    954.5    0.01    0.02  5e-07
[1,0]<stdout>:       16384          4096     float     sum    955.0    0.02    0.03  5e-07    955.5    0.02    0.03  5e-07
[1,0]<stdout>:       32768          8192     float     sum    975.5    0.03    0.06  5e-07   1009.0    0.03    0.06  5e-07
[1,0]<stdout>:       65536         16384     float     sum   1353.4    0.05    0.09  5e-07   1343.5    0.05    0.09  5e-07
[1,0]<stdout>:      131072         32768     float     sum   1395.9    0.09    0.18  5e-07   1392.6    0.09    0.18  5e-07
[1,0]<stdout>:      262144         65536     float     sum   1476.7    0.18    0.33  5e-07   1536.3    0.17    0.32  5e-07
[1,0]<stdout>:      524288        131072     float     sum   1560.3    0.34    0.63  5e-07   1568.3    0.33    0.63  5e-07
[1,0]<stdout>:     1048576        262144     float     sum   1599.2    0.66    1.23  5e-07   1595.3    0.66    1.23  5e-07
[1,0]<stdout>:     2097152        524288     float     sum   1671.1    1.25    2.35  5e-07   1672.5    1.25    2.35  5e-07
[1,0]<stdout>:     4194304       1048576     float     sum   1785.1    2.35    4.41  5e-07   1780.3    2.36    4.42  5e-07
[1,0]<stdout>:     8388608       2097152     float     sum   2133.6    3.93    7.37  5e-07   2135.0    3.93    7.37  5e-07
[1,0]<stdout>:    16777216       4194304     float     sum   2650.9    6.33   11.87  5e-07   2649.9    6.33   11.87  5e-07
[1,0]<stdout>:    33554432       8388608     float     sum   3422.0    9.81   18.39  5e-07   3478.7    9.65   18.09  5e-07
[1,0]<stdout>:    67108864      16777216     float     sum   4783.2   14.03   26.31  5e-07   4782.6   14.03   26.31  5e-07
[1,0]<stdout>:   134217728      33554432     float     sum   7216.9   18.60   34.87  5e-07   7240.9   18.54   34.75  5e-07
[1,0]<stdout>:   268435456      67108864     float     sum    12738   21.07   39.51  5e-07    12802   20.97   39.31  5e-07
[1,0]<stdout>:   536870912     134217728     float     sum    24375   22.03   41.30  5e-07    24403   22.00   41.25  5e-07
[1,0]<stdout>:  1073741824     268435456     float     sum    47904   22.41   42.03  5e-07    47893   22.42   42.04  5e-07
[1,4]<stdout>:test-efa-nccl-worker-0:33:33 [4] NCCL INFO comm 0x7fd4a0000f60 rank 4 nranks 16 cudaDev 4 busId 901c0 - Destroy COMPLETE
[1,0]<stdout>:# Out of bounds values : 0 OK
[1,0]<stdout>:# Avg bus bandwidth    : 8.23785

Vi kan observera i testresultaten att den maximala genomströmningen är cirka 42 GB/sek och den genomsnittliga bussbandbredden är cirka 8 GB.

Vi genomförde också experiment med en enda EFA-adapter aktiverad såväl som inga EFA-adaptrar. Alla resultat sammanfattas i följande tabell.

Vi fann också att för relativt små modeller som ImageNet minskar användningen av accelererat nätverk träningstiden per epok endast med 5–8 % vid batchstorlek på 64. För större modeller och mindre batchstorlekar, när ökad nätverkskommunikation av vikter behövs , har användningen av accelererat nätverk större inverkan. Vi observerade en minskning av epokträningstiden med 15–18 % för träning av EfficientNet-B7 med batchstorlek 1. Den faktiska effekten av EFA på din träning kommer att bero på storleken på din modell.

GPU-övervakning

Innan vi kör träningsjobbet kan vi även ställa upp amazoncloudwatch mätvärden för att visualisera GPU-användningen under träning. Det kan vara bra att veta om resurserna används optimalt eller potentiellt identifiera resurssvält och flaskhalsar i utbildningsprocessen.

De relevanta skripten för att ställa in CloudWatch finns i gpu-metrics mapp. Först skapar vi en Docker-bild med amazon-cloudwatch-agent och nvidia-smi. Vi kan använda Dockerfilen i gpu-metrics mapp för att skapa den här bilden. Förutsatt att ECR-registret redan är inställt i .env fil från föregående steg kan vi bygga och skjuta bilden med hjälp av build.sh och push.sh. Efter detta kör du deploy.sh skriptet slutför automatiskt installationen. Den startar en demonset med amazon-cloudwatch-agent och skickar olika mätvärden till CloudWatch. GPU-statistiken visas under CWAgent namnutrymme på CloudWatch-konsolen. Resten av klustermåtten visas under ContainerInsights namnområde.

Modellutbildning

Alla skript som behövs för PyTorch-träning finns i elastiskt jobb mapp i GitHub-repo. Innan vi startar utbildningsjobbet måste vi köra etcd server, som används av TEC för arbetarupptäckt och parameterutbyte. De deploy.sh manus i elasticjob mappen gör precis det.

För att dra fördel av EFA i p4d.24xlarge-instanser måste vi använda en specifik Docker-bild som är tillgänglig i Amazon ECR Public Gallery som stöder NCCL-kommunikation genom EFA. Vi behöver bara kopiera vår träningskod till denna Docker-bild. De Dockerfile under prover mappen skapar en bild som ska användas när du kör träningsjobb på p4d-instanser. Som alltid kan vi använda build.sh och push.sh skript i mappen för att bygga och pusha bilden.

Smakämnen imagenet-efa.yaml fil beskriver träningsjobbet. Denna .yaml-fil ställer in de resurser som behövs för att köra träningsjobbet och monterar även den beständiga volymen med träningsdata som ställts in i föregående avsnitt.

Ett par saker är värda att påpeka här. Antalet repliker bör ställas in på antalet tillgängliga noder i klustret. I vårt fall satte vi detta till 3 eftersom vi hade tre p4d.24xlarge noder. I den imagenet-efa.yaml fil, nvidia.com/gpu parameter under resurser och nproc_per_node under args bör ställas in på antalet GPU:er per nod, vilket i fallet med p4d.24xlarge är 8. Arbetarargumentet för Python-skriptet anger också antalet CPU:er per process. Vi valde detta till 4 eftersom detta i våra experiment ger optimal prestanda när du kör på p4d.24xlarge-instanser. Dessa inställningar är nödvändiga för att maximera användningen av alla tillgängliga hårdvaruresurser i klustret.

När jobbet körs kan vi observera GPU-användningen i CloudWatch för alla GPU:er i klustret. Följande är ett exempel från ett av våra utbildningsjobb med tre p4d.24xlarge noder i klustret. Här har vi valt en GPU från varje nod. Med de inställningar som nämnts tidigare är GPU-användningen nära 100 % under epokens träningsfas för alla noder i klustret.

För att träna en ResNet50-modell med p3.8xlarge-instanser behöver vi exakt samma steg som beskrivs för EfficientNet-utbildningen med p4d.24xlarge. Vi kan också använda samma Docker-bild. Som nämnts tidigare är p3.8xlarge-instanser inte utrustade med EFA. Men för ResNet50-modellen är detta inte en betydande nackdel. De imagenet-fsx.yaml skriptet som tillhandahålls i GitHub-förvaret ställer in träningsjobbet med lämpliga resurser för nodtypen p3.8xlarge. Jobbet använder samma datauppsättning från filsystemet FSx.

GPU-skalning

Vi körde några experiment för att observera hur träningstiden skalar för EfficientNet-B7-modellen genom att öka antalet GPU:er. För att göra detta ändrade vi antalet repliker från 1 till 3 i vår tränings-.yaml-fil för varje träningskörning. Vi observerade bara tiden för en enda epok när vi använde den kompletta ImageNet-datauppsättningen. Följande figur visar resultaten för vårt GPU-skalningsexperiment. Den röda streckade linjen representerar hur träningstiden ska gå ner från en löpning med 8 GPU:er genom att öka antalet GPU:er. Som vi kan se är skalningen ganska nära vad som förväntas.

På liknande sätt fick vi GPU-skalningsplotten för ResNet50-träning på p3.8xlarge-instanser. I det här fallet ändrade vi replikerna i vår .yaml-fil från 1 till 4. Resultaten av detta experiment visas i följande figur.

Städa upp

Det är viktigt att dra ner resurser efter modellträning för att undvika kostnader förknippade med att köra inaktiva instanser. Med varje skript som skapar resurser, GitHub repo tillhandahåller ett matchande skript för att radera dem. För att rensa upp våra inställningar måste vi ta bort FSx-filsystemet innan vi tar bort klustret eftersom det är associerat med ett subnät i klustrets VPC. För att ta bort FSx-filsystemet behöver vi bara köra följande kommando (inifrån FSX mapp):

kubectl delete -f fsx-pvc-dynamic.yaml
./delete.sh

Observera att detta inte bara tar bort den beständiga volymen, det kommer också att radera filsystemet FSx och all data på filsystemet kommer att gå förlorad. När detta steg är klart kan vi ta bort klustret genom att använda följande skript i exempel mapp:

./eks-delete.sh

Detta kommer att ta bort alla befintliga pods, ta bort klustret och ta bort VPC som skapades i början.

Slutsats

I det här inlägget beskrev vi stegen som behövs för att köra PyTorch-träning för distribuerad dataparallellmodell på EKS-kluster. Denna uppgift kan verka skrämmande, men AWS DevOps för EKS projekt skapat av ML Frameworks-teamet på AWS tillhandahåller alla nödvändiga skript och verktyg för att förenkla processen och göra distribuerad modellutbildning lättillgänglig.

För mer information om teknikerna som används i det här inlägget, besök Amazon EX och Torch Distributed Elastisk. Vi uppmuntrar dig att tillämpa det tillvägagångssätt som beskrivs här på dina egna användningsfall för distribuerad utbildning.

Resurser

Om författarna

Imran Younus är en Principal Solutions Architect för ML Frameworks-team på AWS. Han fokuserar på storskalig maskininlärning och arbetsbelastningar för djupinlärning över AWS-tjänster som Amazon EKS och AWS ParallelCluster. Han har lång erfarenhet av tillämpningar av Deep Leaning i datorseende och industriell IoT. Imran tog sin doktorsexamen i högenergipartikelfysik där han har varit involverad i att analysera experimentella data på peta-byte-skalor.

Alex Iankoulski är en mjukvaru- och infrastrukturarkitekt i full stack som gillar att göra djupgående, praktiskt arbete. Han är för närvarande en Principal Solutions Architect for Self-managed Machine Learning på AWS. I sin roll fokuserar han på att hjälpa kunder med containerisering och orkestrering av ML- och AI-arbetsbelastningar på containerdrivna AWS-tjänster. Han är också författare till öppen källkod Gör ramverk och en Docker-kapten som älskar att använda containerteknologier för att påskynda innovationstakten och samtidigt lösa världens största utmaningar. Under de senaste 10 åren har Alex arbetat med att bekämpa klimatförändringar, demokratisera AI och ML, göra resor säkrare, hälsovården bättre och energismartare.