Distribuert trening med Amazon EKS og Torch Distributed Elastic

Distribuert dyplæringsmodellopplæring blir stadig viktigere ettersom datastørrelsene vokser i mange bransjer. Mange applikasjoner innen datasyn og naturlig språkbehandling krever nå opplæring av dyplæringsmodeller, som vokser eksponentielt i kompleksitet og ofte trenes med hundrevis av terabyte med data. Det blir da viktig å bruke en enorm skyinfrastruktur for å skalere opplæringen av så store modeller.

Utviklere kan bruke åpen kildekode-rammeverk som PyTorch for enkelt å designe intuitive modellarkitekturer. Imidlertid kan det være utfordrende å skalere opplæringen av disse modellene på tvers av flere noder på grunn av økt kompleksitet i orkestreringen.

Distribuert modelltrening består hovedsakelig av to paradigmer:

• Modell parallell – I modellparallell trening er modellen i seg selv så stor at den ikke får plass i minnet til en enkelt GPU, og det trengs flere GPUer for å trene modellen. Open AIs GPT-3-modell med 175 milliarder trenbare parametere (ca. 350 GB i størrelse) er et godt eksempel på dette.
• Dataparallell – I dataparallell trening kan modellen ligge i en enkelt GPU, men fordi dataene er så store, kan det ta dager eller uker å trene en modell. Distribuering av data på tvers av flere GPU-noder kan redusere treningstiden betydelig.

I dette innlegget gir vi et eksempelarkitektur for å trene PyTorch-modeller ved å bruke Elastikk med lommelykt rammeverk på en distribuert dataparallell måte ved hjelp av Amazon Elastic Kubernetes-tjeneste (Amazon EKS).

Forutsetninger

For å gjenskape resultatene som er rapportert i dette innlegget, er den eneste forutsetningen en AWS-konto. I denne kontoen oppretter vi en EKS-klynge og en Amazon FSx for Luster filsystem. Vi skyver også containerbilder til en Amazon Elastic Container Registry (Amazon ECR) depot på kontoen. Instruksjoner for å sette opp disse komponentene er gitt etter behov i hele innlegget.

EKS-klynger

Amazon EKS er en administrert containertjeneste for å kjøre og skalere Kubernetes-applikasjoner på AWS. Med Amazon EKS kan du effektivt kjøre distribuerte opplæringsjobber ved å bruke det nyeste Amazon Elastic Compute Cloud (Amazon EC2)-forekomster uten å måtte installere, betjene og vedlikeholde ditt eget kontrollplan eller noder. Det er en populær orchestrator for maskinlæring (ML) og AI-arbeidsflyter. En typisk EKS-klynge i AWS ser ut som følgende figur.

Vi har lansert et åpen kildekode-prosjekt, AWS DevOps for EKS (aws-do-eks), som gir en stor samling av brukervennlige og konfigurerbare skript og verktøy for å klargjøre EKS-klynger og kjøre distribuerte opplæringsjobber. Dette prosjektet er bygget etter prinsippene til Gjør rammeverk: Enkelhet, fleksibilitet og universalitet. Du kan konfigurere ønsket klynge ved å bruke eks.conf fil og start den ved å kjøre eks-create.sh manus. Detaljerte instruksjoner er gitt i GitHub repo.

Tren PyTorch-modeller med Torch Distributed Elastic

Torch Distributed Elastic (TDE) er et innebygd PyTorch-bibliotek for opplæring av store dyplæringsmodeller der det er avgjørende å skalere beregningsressurser dynamisk basert på tilgjengelighet. De TorchElastic Controller for Kubernetes er en innebygd Kubernetes-implementering for TDE som automatisk administrerer livssyklusen til podene og tjenestene som kreves for TDE-trening. Det gir mulighet for dynamisk skalering av dataressurser under trening etter behov. Den gir også feiltolerant opplæring ved å gjenopprette jobber fra nodesvikt.

I dette innlegget diskuterer vi trinnene for å trene PyTorch EfficientNet-B7 og ResNet50 modeller som bruker IMAGEnet data på en distribuert måte med TDE. Vi bruker PyTorch DistribuertDataParallell API og Kubernetes TorchElastic-kontrolleren, og kjør treningsjobbene våre på en EKS-klynge som inneholder flere GPU-noder. Følgende diagram viser arkitekturdiagrammet for denne modellopplæringen.

TorchElastic for Kubernetes består hovedsakelig av to komponenter: TorchElastic Kubernetes Controller (TEC) og parameterserveren (etcd). Kontrolleren er ansvarlig for å overvåke og administrere treningsjobbene, og parameterserveren holder styr på arbeidernodene for distribuert synkronisering og peer-oppdagelse.

For at treningspodene skal få tilgang til dataene trenger vi et delt datavolum som kan monteres av hver pod. Noen alternativer for delte volumer gjennom Beholderlagringsgrensesnitt (CSI) drivere inkludert i AWS DevOps for EKS er Amazon elastisk filsystem (Amazon EFS) og FSx for Luster.

Klyngeoppsett

I klyngekonfigurasjonen vår bruker vi én c5.2xlarge-forekomst for systempoder. Vi bruker tre p4d.24xlarge-forekomster som worker pods for å trene en EfficientNet-modell. For ResNet50-opplæring bruker vi p3.8xlarge-forekomster som worker pods. I tillegg bruker vi et FSx delt filsystem for å lagre treningsdata og modellartefakter.

AWS p4d.24xlarge instanser er utstyrt med Elastisk stoffadapter (EFA) for å gi nettverk mellom noder. Vi diskuterer EFA mer senere i innlegget. For å aktivere kommunikasjon gjennom EFA, må vi konfigurere klyngeoppsettet gjennom en .yaml-fil. An eksempelfil er gitt i GitHub-depotet.

Etter at denne yaml-filen er riktig konfigurert, kan vi starte klyngen ved å bruke skriptet i GitHub-repoen:

./eks-create.sh

Referere til GitHub repo for detaljerte instruksjoner.

Det er praktisk talt ingen forskjell mellom å kjøre jobber på p4d.24xlarge og p3.8xlarge. Trinnene beskrevet i dette innlegget fungerer for begge. Den eneste forskjellen er tilgjengeligheten av EFA på p4d.24xlarge-forekomster. For mindre modeller som ResNet50 har standard nettverk sammenlignet med EFA-nettverk minimal innvirkning på treningshastigheten.

FSx for Luster filsystem

FSx er designet for høyytelses databehandlingsarbeidsbelastninger og gir forsinkelser på under millisekunder ved bruk av solid-state-disklagringsvolumer. Vi valgte FSx fordi det ga bedre ytelse da vi skalert til et stort antall noder. En viktig detalj å merke seg er at FSx bare kan eksistere i en enkelt tilgjengelighetssone. Derfor bør alle noder som har tilgang til FSx-filsystemet eksistere i samme tilgjengelighetssone som FSx-filsystemet. En måte å oppnå dette på er å spesifisere den relevante tilgjengelighetssonen i klyngen .yaml-filen for de spesifikke nodegruppene før du oppretter klyngen. Alternativt kan vi endre nettverksdelen av den automatiske skaleringsgruppen for disse nodene etter at klyngen er satt opp, og begrense den til å bruke et enkelt subnett. Dette kan enkelt gjøres på Amazon EC2-konsollen.

Forutsatt at EKS-klyngen er oppe og kjører, og subnett-IDen for tilgjengelighetssonen er kjent, kan vi sette opp et FSx-filsystem ved å oppgi nødvendig informasjon i fsx.conf fil som beskrevet i readme og kjører deploy.sh manus i FSX mappe. Dette setter opp riktig policy og sikkerhetsgruppe for tilgang til filsystemet. Skriptet installerer også CSI driver for FSx som et daemonsett. Til slutt kan vi opprette FSx vedvarende volumkrav i Kubernetes ved å bruke en enkelt .yaml-fil:

kubectl apply -f fsx-pvc-dynamic.yaml

Dette oppretter et FSx-filsystem i tilgjengelighetssonen spesifisert i fsx.conf fil, og oppretter også et vedvarende volumkrav fsx-pvc, som kan monteres av en hvilken som helst av podene i klyngen på en read-write-many (RWX) måte.

I eksperimentet vårt brukte vi komplette ImageNet-data, som inneholder mer enn 12 millioner treningsbilder fordelt på 1,000 klasser. Dataene kan lastes ned fra ImageNet nettsted. Den originale TAR-ballen har flere kataloger, men for modellopplæringen vår er vi bare interessert i ILSVRC/Data/CLS-LOC/, som inkluderer train og val underkataloger. Før trening må vi omorganisere bildene i val underkatalog for å matche katalogstrukturen som kreves av PyTorch ImageFolder klasse. Dette kan gjøres ved hjelp av en enkel Python-manus etter at dataene er kopiert til det vedvarende volumet i neste trinn.

For å kopiere dataene fra en Amazon enkel lagringstjeneste (Amazon S3)-bøtte til FSx-filsystemet, lager vi et Docker-bilde som inkluderer skript for denne oppgaven. Et eksempel på Dockerfile og et shell-skript er inkludert i csi mappe i GitHub-repoen. Vi kan bygge bildet ved å bruke build.sh script og skyv det deretter til Amazon ECR ved hjelp av push.sh manus. Før vi bruker disse skriptene, må vi oppgi riktig URI for ECR-depotet i .env fil i rotmappen til GitHub-repoen. Etter at vi har sendt Docker-bildet til Amazon ECR, kan vi starte en pod for å kopiere dataene ved å bruke den relevante .yaml-filen:

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

Poden kjører skriptet automatisk data-prep.sh for å kopiere dataene fra Amazon S3 til det delte volumet. Fordi ImageNet-dataene har mer enn 12 millioner filer, tar kopieringsprosessen et par timer. Python-skriptet imagenet_data_prep.py kjøres også for å omorganisere val datasett som forventet av PyTorch.

Nettverksakselerasjon

Vi kan bruke Elastic Fabric Adapter (EFA) i kombinasjon med støttede EC2-forekomsttyper for å akselerere nettverkstrafikk mellom GPU-nodene i klyngen din. Dette kan være nyttig når du kjører store distribuerte opplæringsjobber der standard nettverkskommunikasjon kan være en flaskehals. Skript for å distribuere og teste EFA-enhetsplugin i EKS-klyngen som vi bruker her, er inkludert i efa-enhet-plugin mappen i GitHub-repoen. For å aktivere en jobb med EFA i EKS-klyngen din, i tillegg til at klyngenodene har nødvendig maskinvare og programvare, må EFA-enhetspluginen distribueres til klyngen, og jobbbeholderen må ha kompatible CUDA og NCCL versjoner installert.

For å demonstrere å kjøre NCCL-tester og evaluere ytelsen til EFA på p4d.24xlarge-forekomster, må vi først distribuere Kubeflow MPI-operatøren ved å kjøre den tilsvarende deploy.sh manus i mpi-operatør mappe. Så kjører vi deploy.sh skript og oppdater test-efa-nccl.yaml manifest så grenser og forespørsler om ressurs vpc.amazonaws.com er satt til 4. De fire tilgjengelige EFA-adapterne i p4d.24xlarge-nodene blir buntet sammen for å gi maksimal gjennomstrømning.

Kjør kubectl apply -f ./test-efa-nccl.yaml for å bruke testen og deretter vise loggene til testpoden. Følgende linje i loggutgangen bekrefter at EFA brukes:

NCCL INFO NET/OFI Selected Provider is efa

Testresultatene skal ligne på følgende 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 observere i testresultatene at maks gjennomstrømning er omtrent 42 GB/sek og gjennomsnittlig bussbåndbredde er omtrent 8 GB.

Vi utførte også eksperimenter med en enkelt EFA-adapter aktivert, så vel som ingen EFA-adaptere. Alle resultater er oppsummert i tabellen nedenfor.

Vi fant også at for relativt små modeller som ImageNet, reduserer bruken av akselerert nettverk treningstiden per epoke bare med 5–8 % ved batchstørrelse på 64. For større modeller og mindre batchstørrelser, når økt nettverkskommunikasjon av vekter er nødvendig , har bruken av akselerert nettverk større innvirkning. Vi observerte en reduksjon av epoketreningstid med 15–18 % for trening av EfficientNet-B7 med batchstørrelse 1. Den faktiske innvirkningen av EFA på treningen din vil avhenge av størrelsen på modellen din.

GPU-overvåking

Før vi kjører treningsjobben kan vi også sette opp Amazon CloudWatch målinger for å visualisere GPU-bruken under trening. Det kan være nyttig å vite om ressursene brukes optimalt eller potensielt identifisere ressurssult og flaskehalser i opplæringsprosessen.

De relevante skriptene for å sette opp CloudWatch er plassert i gpu-metrikk mappe. Først lager vi et Docker-bilde med amazon-cloudwatch-agent og nvidia-smi. Vi kan bruke Dockerfilen i gpu-metrics mappe for å lage dette bildet. Forutsatt at ECR-registeret allerede er satt i .env fil fra forrige trinn, kan vi bygge og skyve bildet ved hjelp av build.sh og push.sh. Etter dette kjører du deploy.sh skriptet fullfører automatisk oppsettet. Den starter et daemonset med amazon-cloudwatch-agent og sender ulike beregninger til CloudWatch. GPU-beregningene vises under CWAgent navneområde på CloudWatch-konsollen. Resten av klyngeberegningene vises under ContainerInsights navnerom.

Modelltrening

Alle skriptene som trengs for PyTorch-trening er plassert i elastisk jobb mappen i GitHub-repoen. Før vi starter treningsjobben, må vi kjøre etcd server, som brukes av TEC for arbeideroppdagelse og parameterutveksling. De deploy.sh manus i elasticjob mappe gjør akkurat det.

For å dra nytte av EFA i p4d.24xlarge-forekomster, må vi bruke et spesifikt Docker-bilde som er tilgjengelig i Amazon ECR Public Gallery som støtter NCCL-kommunikasjon gjennom EFA. Vi trenger bare å kopiere treningskoden vår til dette Docker-bildet. De Dockerfile under prøver mappen oppretter et bilde som skal brukes når du kjører treningsjobb på p4d-forekomster. Som alltid kan vi bruke build.sh og push.sh skript i mappen for å bygge og skyve bildet.

De imagenet-efa.yaml fil beskriver treningsjobben. Denne .yaml-filen setter opp ressursene som trengs for å kjøre treningsjobben og monterer også det vedvarende volumet med treningsdataene som er satt opp i forrige seksjon.

Et par ting er verdt å påpeke her. Antall replikaer bør settes til antall tilgjengelige noder i klyngen. I vårt tilfelle satte vi dette til 3 fordi vi hadde tre p4d.24xlarge noder. I imagenet-efa.yaml fil, nvidia.com/gpu parameter under ressurser og nproc_per_node etter args bør settes til antall GPUer per node, som i tilfellet med p4d.24xlarge er 8. Arbeiderargumentet for Python-skriptet angir også antall CPUer per prosess. Vi valgte dette til å være 4 fordi dette i våre eksperimenter gir optimal ytelse når du kjører på p4d.24xlarge-forekomster. Disse innstillingene er nødvendige for å maksimere bruken av alle maskinvareressursene som er tilgjengelige i klyngen.

Når jobben kjører, kan vi observere GPU-bruken i CloudWatch for alle GPUene i klyngen. Følgende er et eksempel fra en av våre treningsjobber med tre p4d.24xlarge noder i klyngen. Her har vi valgt én GPU fra hver node. Med innstillingene nevnt tidligere, er GPU-bruken nær 100 % under treningsfasen av epoken for alle nodene i klyngen.

For å trene en ResNet50-modell ved bruk av p3.8xlarge-forekomster, trenger vi nøyaktig de samme trinnene som beskrevet for EfficientNet-treningen med p4d.24xlarge. Vi kan også bruke det samme Docker-bildet. Som nevnt tidligere, er ikke p3.8xlarge-forekomster utstyrt med EFA. For ResNet50-modellen er dette imidlertid ikke en betydelig ulempe. De imagenet-fsx.yaml skriptet gitt i GitHub-depotet setter opp treningsjobben med passende ressurser for nodetypen p3.8xlarge. Jobben bruker det samme datasettet fra FSx-filsystemet.

GPU-skalering

Vi kjørte noen eksperimenter for å observere hvordan treningstiden skalerer for EfficientNet-B7-modellen ved å øke antall GPUer. For å gjøre dette endret vi antall replikaer fra 1 til 3 i trenings-.yaml-filen for hver treningskjøring. Vi observerte bare tiden for en enkelt epoke mens vi brukte hele ImageNet-datasettet. Følgende figur viser resultatene for GPU-skaleringseksperimentet vårt. Den røde stiplede linjen representerer hvordan treningstiden skal gå ned fra en løpetur med 8 GPUer ved å øke antall GPUer. Som vi kan se er skaleringen ganske nær det som er forventet.

På samme måte fikk vi GPU-skaleringsplottet for ResNet50-trening på p3.8xlarge-forekomster. For dette tilfellet endret vi replikaene i .yaml-filen vår fra 1 til 4. Resultatene av dette eksperimentet er vist i følgende figur.

Rydd opp

Det er viktig å spinne ned ressurser etter modelltrening for å unngå kostnader forbundet med å kjøre inaktive instanser. Med hvert skript som skaper ressurser, GitHub repo gir et samsvarende skript for å slette dem. For å rydde opp i oppsettet må vi slette FSx-filsystemet før vi sletter klyngen fordi den er knyttet til et undernett i klyngens VPC. For å slette FSx-filsystemet trenger vi bare å kjøre følgende kommando (fra innsiden av FSX mappe):

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

Merk at dette ikke bare vil slette det vedvarende volumet, det vil også slette FSx-filsystemet, og alle dataene på filsystemet vil gå tapt. Når dette trinnet er fullført, kan vi slette klyngen ved å bruke følgende skript i eks mappe:

./eks-delete.sh

Dette vil slette alle eksisterende pods, fjerne klyngen og slette VPC-en som ble opprettet i begynnelsen.

konklusjonen

I dette innlegget beskrev vi trinnene som trengs for å kjøre PyTorch-distribuerte dataparallelle modelltrening på EKS-klynger. Denne oppgaven kan virke skremmende, men AWS DevOps for EKS prosjekt opprettet av ML Frameworks-teamet hos AWS gir alle nødvendige skript og verktøy for å forenkle prosessen og gjøre distribuert modellopplæring lett tilgjengelig.

For mer informasjon om teknologiene som brukes i dette innlegget, besøk Amazon EX og Elastikk med lommelykt. Vi oppfordrer deg til å bruke tilnærmingen som er beskrevet her på dine egne distribuerte opplæringstilfeller.

Ressurser

Om forfatterne

Imran Younus er en Principal Solutions Architect for ML Frameworks-teamet ved AWS. Han fokuserer på maskinlæring i stor skala og arbeidsbelastninger for dyp læring på tvers av AWS-tjenester som Amazon EKS og AWS ParallelCluster. Han har lang erfaring med anvendelser av Deep Leaning i Computer Vision og Industrial IoT. Imran oppnådde sin doktorgrad i høyenergipartikkelfysikk hvor han har vært involvert i å analysere eksperimentelle data på peta-byte-skalaer.

Alex Iankoulski er en fullstack programvare- og infrastrukturarkitekt som liker å gjøre dypt, praktisk arbeid. Han er for tiden en hovedløsningsarkitekt for selvstyrt maskinlæring ved AWS. I sin rolle fokuserer han på å hjelpe kunder med containerisering og orkestrering av ML- og AI-arbeidsmengder på containerdrevne AWS-tjenester. Han er også forfatteren av åpen kildekode Gjør rammeverk og en Docker-kaptein som elsker å bruke containerteknologier for å akselerere innovasjonstakten samtidig som den løser verdens største utfordringer. I løpet av de siste 10 årene har Alex jobbet med å bekjempe klimaendringer, demokratisere AI og ML, gjøre reiser tryggere, helsevesenet bedre og energismartere.