Distribueret træning med Amazon EKS og Torch Distributed Elastic

Distribueret dyb læringsmodeltræning bliver stadig vigtigere, efterhånden som datastørrelser vokser i mange brancher. Mange applikationer inden for computersyn og naturlig sprogbehandling kræver nu træning af deep learning-modeller, som vokser eksponentielt i kompleksitet og ofte trænes med hundredvis af terabyte data. Så bliver det vigtigt at bruge en enorm cloud-infrastruktur til at skalere træningen af ​​så store modeller.

Udviklere kan bruge open source-rammer såsom PyTorch til nemt at designe intuitive modelarkitekturer. Det kan dog være en udfordring at skalere træningen af ​​disse modeller på tværs af flere noder på grund af øget orkestreringskompleksitet.

Distribueret modeltræning består hovedsageligt af to paradigmer:

• Model parallel – Ved model parallel træning er selve modellen så stor, at den ikke kan passe ind i hukommelsen på en enkelt GPU, og der skal flere GPU'er til for at træne modellen. Open AI's GPT-3-model med 175 milliarder parametre, der kan trænes (ca. 350 GB i størrelse) er et godt eksempel på dette.
• Data parallel – I data parallel træning kan modellen ligge i en enkelt GPU, men fordi dataene er så store, kan det tage dage eller uger at træne en model. Distribution af data på tværs af flere GPU-noder kan reducere træningstiden betydeligt.

I dette indlæg giver vi et eksempel på arkitektur til at træne PyTorch-modeller ved hjælp af Fakkelfordelt elastik framework på en distribueret data parallel måde ved hjælp af Amazon Elastic Kubernetes Service (Amazon EKS).

Forudsætninger

For at replikere resultaterne rapporteret i dette indlæg er den eneste forudsætning en AWS-konto. På denne konto opretter vi en EKS-klynge og en Amazon FSx til Luster filsystem. Vi skubber også containerbilleder til en Amazon Elastic Container Registry (Amazon ECR) lager på kontoen. Instruktioner til opsætning af disse komponenter leveres efter behov i hele indlægget.

EKS klynger

Amazon EKS er en administreret containertjeneste til at køre og skalere Kubernetes-applikationer på AWS. Med Amazon EKS kan du effektivt køre distribuerede træningsjob ved at bruge det nyeste Amazon Elastic Compute Cloud (Amazon EC2) forekomster uden at skulle installere, betjene og vedligeholde dit eget kontrolplan eller noder. Det er en populær orkestrator til maskinlæring (ML) og AI-arbejdsgange. En typisk EKS-klynge i AWS ser ud som den følgende figur.

Vi har udgivet et open source-projekt, AWS DevOps til EKS (aws-do-eks), som giver en stor samling af brugervenlige og konfigurerbare scripts og værktøjer til at klargøre EKS-klynger og køre distribuerede træningsjob. Dette projekt er bygget efter principperne i Lav Framework: Enkelhed, fleksibilitet og universalitet. Du kan konfigurere din ønskede klynge ved at bruge eks.konf fil, og start den derefter ved at køre eks-create.sh manuskript. Detaljerede instruktioner findes i GitHub repo.

Træn PyTorch-modeller ved hjælp af Torch Distributed Elastic

Torch Distributed Elastic (TDE) er et indbygget PyTorch-bibliotek til træning af store dybdelæringsmodeller, hvor det er afgørende at skalere beregningsressourcer dynamisk baseret på tilgængelighed. Det TorchElastic Controller til Kubernetes er en indbygget Kubernetes-implementering til TDE, der automatisk styrer livscyklussen for de pods og tjenester, der kræves til TDE-træning. Det giver mulighed for dynamisk skalering af computerressourcer under træning efter behov. Det giver også fejltolerant træning ved at genoprette job efter knudefejl.

I dette indlæg diskuterer vi trinene til at træne PyTorch EfficientNet-B7 , ResNet50 modeller, der bruger IMAGEnet data på en distribueret måde med TDE. Vi bruger PyTorch DistributedDataParallel API og Kubernetes TorchElastic-controlleren og køre vores træningsjob på en EKS-klynge, der indeholder flere GPU-noder. Følgende diagram viser arkitekturdiagrammet for denne modeltræning.

TorchElastic for Kubernetes består hovedsageligt af to komponenter: TorchElastic Kubernetes Controller (TEC) og parameterserveren (etcd). Controlleren er ansvarlig for overvågning og styring af træningsjob, og parameterserveren holder styr på arbejderknudepunkterne for distribueret synkronisering og peer-opdagelse.

For at træningspoderne kan få adgang til dataene, har vi brug for en delt datavolumen, der kan monteres af hver pod. Nogle muligheder for delte mængder igennem Container Storage Interface (CSI) drivere inkluderet i AWS DevOps til EKS er Amazon Elastic File System (Amazon EFS) og FSx for Luster.

Klynge opsætning

I vores klyngekonfiguration bruger vi én c5.2xlarge-instans til systempods. Vi bruger tre p4d.24xlarge forekomster som worker pods til at træne en EfficientNet-model. Til ResNet50-træning bruger vi p3.8xlarge instanser som worker pods. Derudover bruger vi et FSx-delt filsystem til at gemme vores træningsdata og modelartefakter.

AWS p4d.24xlarge instanser er udstyret med Elastisk stofadapter (EFA) for at give netværk mellem noder. Vi diskuterer EFA mere senere i indlægget. For at aktivere kommunikation gennem EFA, skal vi konfigurere klyngeopsætningen gennem en .yaml-fil. An eksempel fil findes i GitHub-lageret.

Når denne .yaml-fil er korrekt konfigureret, kan vi starte klyngen ved hjælp af scriptet i GitHub-repoen:

./eks-create.sh

Se i GitHub repo for detaljerede instruktioner.

Der er praktisk talt ingen forskel på at køre job på p4d.24xlarge og p3.8xlarge. Trinene beskrevet i dette indlæg fungerer for begge. Den eneste forskel er tilgængeligheden af ​​EFA på p4d.24xlarge forekomster. For mindre modeller som ResNet50 har standardnetværk sammenlignet med EFA-netværk minimal indflydelse på træningshastigheden.

FSx til Luster filsystem

FSx er designet til højtydende computerarbejdsbelastninger og giver forsinkelse på under millisekunder ved brug af solid-state-drevlagringsvolumener. Vi valgte FSx, fordi det gav bedre ydeevne, da vi skalerede til et stort antal noder. En vigtig detalje at bemærke er, at FSx kun kan eksistere i en enkelt tilgængelighedszone. Derfor bør alle noder, der får adgang til FSx-filsystemet, eksistere i den samme tilgængelighedszone som FSx-filsystemet. En måde at opnå dette på er at angive den relevante tilgængelighedszone i klyngen .yaml-filen for de specifikke nodegrupper, før klyngen oprettes. Alternativt kan vi ændre netværksdelen af ​​den automatiske skaleringsgruppe for disse noder, efter at klyngen er sat op, og begrænse den til at bruge et enkelt undernet. Dette kan nemt gøres på Amazon EC2-konsollen.

Forudsat at EKS-klyngen er oppe og køre, og subnet-id'et for tilgængelighedszonen er kendt, kan vi konfigurere et FSx-filsystem ved at give de nødvendige oplysninger i fsx.conf fil som beskrevet i readme og kører deploy.sh script i FSX folder. Dette opsætter den korrekte politik og sikkerhedsgruppe til at få adgang til filsystemet. Scriptet installerer også CSI driver til FSx som et daemonset. Endelig kan vi oprette FSx-vedvarende volumenkravet i Kubernetes ved at anvende en enkelt .yaml-fil:

kubectl apply -f fsx-pvc-dynamic.yaml

Dette opretter et FSx-filsystem i den tilgængelighedszone, der er angivet i fsx.conf fil, og opretter også et vedvarende volumenkrav fsx-pvc, som kan monteres af enhver af pods i klyngen på en read-write-many (RWX) måde.

I vores eksperiment brugte vi komplette ImageNet-data, som indeholder mere end 12 millioner træningsbilleder fordelt på 1,000 klasser. Data kan downloades fra ImageNet hjemmeside. Den originale TAR-bold har flere mapper, men til vores modeltræning er vi kun interesserede i ILSVRC/Data/CLS-LOC/, der inkluderer train , val undermapper. Før træning skal vi omarrangere billederne i val undermappe for at matche den mappestruktur, der kræves af PyTorch Billedmappe klasse. Dette kan gøres ved hjælp af en simpel Python script efter at dataene er kopieret til den vedvarende volumen i næste trin.

For at kopiere data fra en Amazon Simple Storage Service (Amazon S3) bucket til FSx-filsystemet, opretter vi et Docker-billede, der indeholder scripts til denne opgave. Et eksempel på Dockerfile og et shell-script er inkluderet i csi mappe i GitHub-repoen. Vi kan bygge billedet ved hjælp af build.sh script og skub det derefter til Amazon ECR ved hjælp af push.sh manuskript. Før vi bruger disse scripts, skal vi angive den korrekte URI for ECR-lageret i .env fil i rodmappen i GitHub-reposen. Når vi har skubbet Docker-billedet til Amazon ECR, kan vi starte en pod til at kopiere dataene ved at anvende den relevante .yaml-fil:

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

Poden kører automatisk scriptet data-prep.sh for at kopiere dataene fra Amazon S3 til det delte volumen. Fordi ImageNet-dataene har mere end 12 millioner filer, tager kopieringsprocessen et par timer. Python-scriptet imagenet_data_prep.py køres også for at omarrangere val datasæt som forventet af PyTorch.

Netværksacceleration

Vi kan bruge Elastic Fabric Adapter (EFA) i kombination med understøttede EC2-instanstyper for at accelerere netværkstrafikken mellem GPU-noderne i din klynge. Dette kan være nyttigt, når du kører store distribuerede træningsjob, hvor standardnetværkskommunikation kan være en flaskehals. Scripts til at implementere og teste EFA-enhedspluginnet i EKS-klyngen, som vi bruger her, er inkluderet i efa-device-plugin mappe i GitHub-reposen. For at aktivere et job med EFA i din EKS-klynge, ud over at klyngeknuderne har den nødvendige hardware og software, skal EFA-enhedspluginet installeres til klyngen, og din jobcontainer skal have kompatibel CUDA og NCCL versioner installeret.

For at demonstrere at køre NCCL-test og evaluere ydeevnen af ​​EFA på p4d.24xlarge-forekomster, skal vi først implementere Kubeflow MPI-operatøren ved at køre den tilsvarende deploy.sh script i mpi-operatør folder. Så kører vi deploy.sh script og opdater test-efa-nccl.yaml manifest så grænser og anmodninger om ressource vpc.amazonaws.com er indstillet til 4. De fire tilgængelige EFA-adaptere i p4d.24xlarge noderne bliver bundtet sammen for at give maksimal gennemstrømning.

Kør kubectl apply -f ./test-efa-nccl.yaml for at anvende testen og derefter vise logfilerne for testpoden. Følgende linje i log-output bekræfter, at EFA bruges:

NCCL INFO NET/OFI Selected Provider is efa

Testresultaterne skal ligne følgende output:

[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 observere i testresultaterne, at den maksimale gennemstrømning er omkring 42 GB/sek. og den gennemsnitlige busbåndbredde er cirka 8 GB.

Vi udførte også eksperimenter med en enkelt EFA-adapter aktiveret såvel som ingen EFA-adaptere. Alle resultater er opsummeret i følgende tabel.

Vi fandt også, at for relativt små modeller som ImageNet reducerer brugen af ​​accelereret netværk træningstiden pr. epoke kun med 5-8 % ved batchstørrelse på 64. For større modeller og mindre batchstørrelser, når øget netværkskommunikation af vægte er nødvendig , har brugen af ​​accelereret netværk større effekt. Vi observerede et fald i epoketræningstid med 15-18 % for træning af EfficientNet-B7 med batchstørrelse 1. Den faktiske effekt af EFA på din træning vil afhænge af størrelsen på din model.

GPU overvågning

Inden vi kører træningsjobbet, kan vi også sætte op amazoncloudwatch målinger til at visualisere GPU-udnyttelsen under træning. Det kan være nyttigt at vide, om ressourcerne bliver brugt optimalt eller potentielt identificerer ressourcesult og flaskehalse i træningsprocessen.

De relevante scripts til opsætning af CloudWatch er placeret i gpu-metrics folder. Først opretter vi et Docker-billede med amazon-cloudwatch-agent , nvidia-smi. Vi kan bruge Dockerfilen i gpu-metrics mappe for at oprette dette billede. Forudsat at ECR-registret allerede er indstillet i .env fil fra det forrige trin, kan vi bygge og skubbe billedet vha build.sh , push.sh. Herefter kører du deploy.sh script fuldfører automatisk opsætningen. Den starter et dæmonsæt med amazon-cloudwatch-agent og skubber forskellige metrics til CloudWatch. GPU-målingerne vises under CWAgent navneområde på CloudWatch-konsollen. Resten af ​​klyngemålingerne vises under ContainerInsights navneområde.

Model træning

Alle de nødvendige scripts til PyTorch-træning er placeret i elastikjob mappe i GitHub-reposen. Før vi påbegynder træningsjobbet, skal vi køre etcd server, som bruges af TEC til arbejderopdagelse og parameterudveksling. Det deploy.sh script i elasticjob mappe gør præcis det.

For at drage fordel af EFA i p4d.24xlarge-forekomster skal vi bruge et specifikt Docker-billede, der er tilgængeligt i Amazon ECR Public Gallery der understøtter NCCL-kommunikation gennem EFA. Vi skal bare kopiere vores træningskode til dette Docker-billede. Det Dockerfil under prøver mappen opretter et billede, der skal bruges, når du kører træningsjob på p4d-forekomster. Som altid kan vi bruge build.sh , push.sh scripts i mappen for at bygge og skubbe billedet.

imagenet-efa.yaml fil beskriver træningsjobbet. Denne .yaml-fil opsætter de ressourcer, der er nødvendige for at køre træningsjobbet og monterer også den vedvarende volumen med træningsdataene, der er opsat i det foregående afsnit.

Et par ting er værd at påpege her. Antallet af replikaer skal indstilles til antallet af tilgængelige noder i klyngen. I vores tilfælde satte vi dette til 3, fordi vi havde tre p4d.24xlarge noder. I den imagenet-efa.yaml fil, den nvidia.com/gpu parameter under ressourcer og nproc_per_node under args skal indstilles til antallet af GPU'er pr. node, hvilket i tilfælde af p4d.24xlarge er 8. Arbejderargumentet for Python-scriptet angiver også antallet af CPU'er pr. proces. Vi valgte dette til at være 4, fordi dette i vores eksperimenter giver optimal ydeevne, når du kører på p4d.24xlarge forekomster. Disse indstillinger er nødvendige for at maksimere brugen af ​​alle de tilgængelige hardwareressourcer i klyngen.

Når jobbet kører, kan vi observere GPU-brugen i CloudWatch for alle GPU'erne i klyngen. Det følgende er et eksempel fra et af vores træningsjob med tre p4d.24xlarge noder i klyngen. Her har vi valgt en GPU fra hver node. Med de tidligere nævnte indstillinger er GPU-bruget tæt på 100 % under træningsfasen af ​​epoken for alle noderne i klyngen.

For at træne en ResNet50-model ved hjælp af p3.8xlarge-instanser har vi brug for nøjagtig de samme trin som beskrevet for EfficientNet-træningen ved hjælp af p4d.24xlarge. Vi kan også bruge det samme Docker-billede. Som tidligere nævnt er p3.8xlarge-instanser ikke udstyret med EFA. For ResNet50-modellen er dette dog ikke en væsentlig ulempe. Det imagenet-fsx.yaml script, der leveres i GitHub-lageret, opsætter træningsjobbet med passende ressourcer til nodetypen p3.8xlarge. Jobbet bruger det samme datasæt fra FSx-filsystemet.

GPU-skalering

Vi kørte nogle eksperimenter for at observere, hvordan træningstiden skalerer for EfficientNet-B7-modellen ved at øge antallet af GPU'er. For at gøre dette ændrede vi antallet af replikaer fra 1 til 3 i vores trænings-.yaml-fil for hver træningskørsel. Vi observerede kun tiden for en enkelt epoke, mens vi brugte det komplette ImageNet-datasæt. Følgende figur viser resultaterne for vores GPU-skaleringseksperiment. Den røde stiplede linje repræsenterer, hvordan træningstiden skal gå ned fra en løbetur med 8 GPU'er ved at øge antallet af GPU'er. Som vi kan se, er skaleringen ret tæt på det forventede.

På samme måde opnåede vi GPU-skaleringsplottet til ResNet50-træning på p3.8xlarge-forekomster. I dette tilfælde ændrede vi replikaerne i vores .yaml-fil fra 1 til 4. Resultaterne af dette eksperiment er vist i følgende figur.

Ryd op

Det er vigtigt at skrue ned for ressourcer efter modeltræning for at undgå omkostninger forbundet med at køre inaktive tilfælde. Med hvert script, der opretter ressourcer, GitHub repo giver et matchende script til at slette dem. For at rydde op i vores opsætning skal vi slette FSx-filsystemet, før vi sletter klyngen, fordi den er forbundet med et undernet i klyngens VPC. For at slette FSx-filsystemet skal vi bare køre følgende kommando (indefra FSX folder):

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

Bemærk, at dette ikke kun vil slette den vedvarende volumen, det vil også slette FSx-filsystemet, og alle data på filsystemet vil gå tabt. Når dette trin er fuldført, kan vi slette klyngen ved at bruge følgende script i EKS folder:

./eks-delete.sh

Dette vil slette alle eksisterende pods, fjerne klyngen og slette den VPC, der blev oprettet i begyndelsen.

Konklusion

I dette indlæg beskrev vi de nødvendige trin for at køre PyTorch distribueret data parallel model træning på EKS klynger. Denne opgave kan virke skræmmende, men den AWS DevOps til EKS projekt oprettet af ML Frameworks-teamet hos AWS giver alle de nødvendige scripts og værktøjer til at forenkle processen og gøre distribueret modeltræning let tilgængelig.

For mere information om de teknologier, der bruges i dette indlæg, besøg Amazon EX , Fakkelfordelt elastik. Vi opfordrer dig til at anvende den her beskrevne tilgang til dine egne distribuerede træningstilfælde.

Ressourcer

Om forfatterne

Imran Younus er et Principal Solutions Architect for ML Frameworks-team hos AWS. Han fokuserer på maskinlæring i stor skala og deep learning-arbejdsbelastninger på tværs af AWS-tjenester som Amazon EKS og AWS ParallelCluster. Han har stor erfaring med anvendelser af Deep Leaning i Computer Vision og Industrial IoT. Imran opnåede sin ph.d. i højenergipartikelfysik, hvor han har været involveret i at analysere eksperimentelle data på peta-byte-skalaer.

Alex Iankoulski er en fuld stack software- og infrastrukturarkitekt, der kan lide at udføre dybt, praktisk arbejde. Han er i øjeblikket Principal Solutions Architect for Self-managed Machine Learning hos AWS. I sin rolle fokuserer han på at hjælpe kunder med containerisering og orkestrering af ML- og AI-arbejdsbelastninger på containerdrevne AWS-tjenester. Han er også forfatter til open source Lav rammer og en Docker-kaptajn, der elsker at anvende containerteknologier til at accelerere innovationstempoet og samtidig løse verdens største udfordringer. I løbet af de sidste 10 år har Alex arbejdet på at bekæmpe klimaændringer, demokratisere AI og ML, gøre rejser sikrere, sundhedspleje bedre og energismartere.