Kjør PyTorch Lightning og native PyTorch DDP på ​​Amazon SageMaker Training, med Amazon Search

Så mye data, så lite tid. Maskinlæringseksperter (ML), dataforskere, ingeniører og entusiaster har støtt på dette problemet over hele verden. Fra naturlig språkbehandling til datasyn, tabeller til tidsserier og alt derimellom, har det eldgamle problemet med å optimalisere for hastighet når du kjører data mot så mange GPUer som du kan få inspirert utallige løsninger. I dag er vi glade for å kunngjøre funksjoner for PyTorch-utviklere som bruker native open source-rammeverk, som PyTorch Lyn og PyTorch DDP, som vil strømlinjeforme deres vei til skyen.

Amazon SageMaker er en fullstendig administrert tjeneste for ML, og SageMaker modellopplæring er et optimalisert datamiljø for høyytelsestrening i skala. SageMaker modelltrening tilbyr en ekstern treningsopplevelse med et sømløst kontrollplan for enkelt å trene og reprodusere ML-modeller til høy ytelse og lave kostnader. Vi er glade for å kunngjøre nye funksjoner i SageMaker-treningsporteføljen som gjør det enda enklere og mer tilgjengelig å kjøre PyTorch i stor skala:

1. PyTorch Lightning kan nå integreres til SageMakers distribuerte data parallellbibliotek med bare én linje med kodeendring.
2. SageMaker modellopplæring har nå støtte for native PyTorch Distributed Data Parallel med NCCL backend, slik at utviklere kan migrere til SageMaker enklere enn noen gang før.

I dette innlegget diskuterer vi disse nye funksjonene, og lærer også hvordan Amazon Search har kjørt PyTorch Lightning med den optimaliserte distribuerte treningsbackend i SageMaker for å øke hastigheten på modelltreningstid.

Før du dykker inn i Amazon Search-casestudien, for de som ikke er kjent, vil vi gjerne gi litt bakgrunn om SageMakers distribuerte data parallellbibliotek. I 2020 utviklet og lanserte vi en tilpasset klyngekonfigurasjon for distribuert gradientnedstigning i skala som øker den totale klyngeeffektiviteten, introdusert på Amazon Science som Sild. Ved å bruke det beste fra både parameterservere og ringbaserte topologier, er SageMaker Distributed Data Parallel (SMDDP) optimalisert for Amazon Elastic Compute Cloud (Amazon EC2) nettverkstopologi, inkludert EFA. For større klyngestørrelser er SMDDP i stand til å levere 20–40 % gjennomstrømsforbedringer i forhold til Horovod (TensorFlow) og PyTorch Distributed Data Parallel. For mindre klyngestørrelser og støttede modeller anbefaler vi SageMaker Training Compiler, som kan redusere den totale jobbtiden med opptil 50 %.

Kundens søkelys: PyTorch Lightning på SageMakers optimaliserte backend med Amazon Search

Amazon Search er ansvarlig for søk og oppdagelsesopplevelse på Amazon.com. Det driver søkeopplevelsen for kunder som leter etter produkter å kjøpe på Amazon. På et høyt nivå bygger Amazon Search en indeks for alle produkter som selges på Amazon.com. Når en kunde skriver inn et søk, bruker Amazon Search en rekke ML-teknikker, inkludert dyplæringsmodeller, for å matche relevante og interessante produkter til kundens forespørsel. Deretter rangerer den produktene før den viser resultatene til kunden.

Amazon Search-forskere har brukt PyTorch Lightning som et av hovedrammene for å trene dyplæringsmodellene som driver søkerangeringen på grunn av dens ekstra brukervennlighetsfunksjoner på toppen av PyTorch. SMDDP ble ikke støttet for dyplæringsmodeller skrevet i PyTorch Lightning før denne nye SageMaker-lanseringen. Dette forhindret Amazon Search-forskere som foretrekker å bruke PyTorch Lightning fra å skalere modelltreningen ved hjelp av dataparallelle teknikker, noe som reduserer treningstiden betydelig og hindrer dem i å teste nye eksperimenter som krever mer skalerbar trening.

Teamets tidlige benchmarking-resultater viser 7.3 ganger raskere treningstid for en prøvemodell når de trenes på åtte noder sammenlignet med en enkelt-node treningsbaseline. Grunnlinjemodellen som brukes i disse benchmarkingen er et flerlags perceptron-nevralt nettverk med syv tette, fullt tilkoblede lag og over 200 parametere. Følgende tabell oppsummerer benchmarking-resultatet på ml.p3.16xlarge SageMaker-treningsforekomster.

Deretter dykker vi ned i detaljene om de nye lanseringene. Hvis du vil, kan du gå gjennom vår tilsvarende eksempel notatbok.

Kjør PyTorch Lightning med det distribuerte treningsbiblioteket fra SageMaker

Vi er glade for å kunngjøre at SageMaker Data Parallel nå sømløst integreres med PyTorch Lightning i SageMaker-trening.

PyTorch Lightning er et åpen kildekode-rammeverk som gir en forenkling for å skrive tilpassede modeller i PyTorch. På noen måter lik det Keras gjorde for TensorFlow, eller til og med uten tvil Hugging Face, gir PyTorch Lightning en API på høyt nivå med abstraksjoner for mye av funksjonaliteten på lavere nivå til selve PyTorch. Dette inkluderer definering av modellen, profilering, evaluering, beskjæring, modellparallellisme, hyperparameterkonfigurasjoner, overføringslæring og mer.

Tidligere var PyTorch Lightning-utviklere usikre på hvordan de sømløst skulle migrere treningskoden til høyytelses SageMaker GPU-klynger. I tillegg var det ingen måte for dem å dra nytte av effektivitetsgevinster introdusert av SageMaker Data Parallel.

For PyTorch Lightning, generelt sett, bør det være små eller ingen kodeendringer for å kjøre disse APIene på SageMaker Training. I eksempelnotatbøkene bruker vi DDP-strategi og DDPPplugin metoder.

Det er tre trinn for å bruke PyTorch Lightning med SageMaker Data Parallel som en optimalisert backend:

1. Bruk en støttet AWS Deep Learning Container (DLC) som basisbilde, eller eventuelt lag din egen beholder og installer SageMaker Data Parallel-backend selv. Sørg for at du har PyTorch Lightning inkludert i de nødvendige pakkene dine, for eksempel med en requirements.txt filen.
2. Gjør noen få mindre kodeendringer i treningsskriptet ditt som aktiverer den optimaliserte backend. Disse inkluderer:
   1. Importer SM DDP-biblioteket:

      import smdistributed.dataparallel.torch.torch_smddp
      

   2. Sett opp PyTorch Lightning-miljøet for SageMaker:

      from pytorch_lightning.plugins.environments.lightning_environment 
        import LightningEnvironment
      
      env = LightningEnvironment()
      env.world_size = lambda: int(os.environ["WORLD_SIZE"])
      env.global_rank = lambda: int(os.environ["RANK"])

   3. Hvis du bruker en versjon av PyTorch Lightning eldre enn 1.5.10, må du legge til noen flere trinn.
      1. Først legger du til miljøvariabelen:

         os.environ["PL_TORCH_DISTRIBUTED_BACKEND"] = "smddp"

      2. For det andre, sørg for at du bruker DDPPlugin, I stedet DDPStrategy. Hvis du bruker en nyere versjon, som du enkelt kan stille inn ved å plassere requirements.txt i source_dir for jobben din, så er dette ikke nødvendig. Se følgende kode:

         ddp = DDPPlugin(parallel_devices=[torch.device("cuda", d) for d in range(num_gpus)], cluster_environment=env)

   4. Du kan eventuelt definere prosessgruppens backend som "smddp" i DDPSTrategy gjenstand. Men hvis du bruker PyTorch Lightning med PyTorch DDP-backend, som også støttes, fjern ganske enkelt denne `process_group_backend` parameter. Se følgende kode:

      ddp = DDPStrategy(
        cluster_environment=env, 
        process_group_backend="smddp", 
        accelerator="gpu")

3. Sørg for at du har en distribusjonsmetode notert i estimatoren, for eksempel distribution={"smdistributed":{"dataparallel":{"enabled":True} hvis du bruker Sild-backend, eller distribution={"pytorchddp":{"enabled":True}.

• For en fullstendig liste over passende parametere i distribution parameter, se vår dokumentasjon her..

Nå kan du starte din SageMaker-treningsjobb! Du kan starte treningsjobben din via Python SDK, Boto3, SageMaker-konsollen, AWS kommandolinjegrensesnitt (AWS CLI), og utallige andre metoder. Fra et AWS-perspektiv er dette en enkelt API-kommando: create-training-job. Enten du starter denne kommandoen fra din lokale terminal, en AWS Lambda funksjon, en Amazon SageMaker Studio en bærbar PC, en KubeFlow-pipeline eller et annet datamiljø er helt opp til deg.

Vær oppmerksom på at integrasjonen mellom PyTorch Lightning og SageMaker Data Parallel for øyeblikket kun støttes for nyere versjoner av PyTorch, fra og med 1.11. I tillegg er denne utgivelsen kun tilgjengelig i AWS DLC-er for SageMaker fra og med PyTorch 1.12. Pass på at du peker på dette bildet som din base. I us-east-1, denne adressen er som følger:

ecr_image = '763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:1.12.0-gpu-py38-cu113-ubuntu20.04-sagemaker'

Deretter kan du utvide Docker-beholderen din ved å bruke dette som basisbildet ditt, eller du kan sende dette som en variabel til image_uri argument av SageMaker treningsvurdering.

Som et resultat vil du kunne kjøre PyTorch Lightning-koden din på SageMaker Trainings optimaliserte GPUer, med den beste ytelsen tilgjengelig på AWS.

Kjør PyTorch Distributed Data Parallel på SageMaker

Det største problemet PyTorch Distributed Data Parallel (DDP) løser er villedende enkelt: hastighet. Et godt distribuert treningsrammeverk skal gi stabilitet, pålitelighet, og viktigst av alt, utmerket ytelse i skala. PyTorch DDP leverer dette ved å gi fakkelutviklere API-er for å replikere modellene deres over flere GPU-enheter, i både enkelt-node og multi-node innstillinger. Rammeverket administrerer deretter skjæring av forskjellige objekter fra treningsdatasettet til hver modellkopi, og beregner et gjennomsnitt av gradientene for hver av modellkopiene for å synkronisere dem ved hvert trinn. Dette gir én modell ved total fullføring av hele treningsløpet. Følgende diagram illustrerer denne prosessen.

PyTorch DDP er vanlig i prosjekter som bruker store datasett. Den nøyaktige størrelsen på hvert datasett vil variere mye, men en generell retningslinje er å skalere datasett, beregne størrelser og modellstørrelser i lignende forhold. Også kalt skaleringslover, den optimale kombinasjonen av disse tre er veldig oppe for debatt og vil variere basert på applikasjoner. Hos AWS, basert på å jobbe med flere kunder, kan vi tydelig se fordeler med dataparallelle strategier når en samlet datasettstørrelse er minst noen få titalls GB. Når datasettene blir enda større, er implementering av en eller annen type dataparallellstrategi en kritisk teknikk for å øke hastigheten på det generelle eksperimentet og forbedre tiden til verdi.

Tidligere manglet kunder som brukte PyTorch DDP for distribuert opplæring i lokaler eller i andre datamiljøer et rammeverk for enkelt å migrere prosjektene sine til SageMaker Training for å dra nytte av høyytelses GPUer med et sømløst kontrollplan. Spesifikt måtte de enten migrere dataparallelle rammeverk til SMDDP, eller utvikle og teste egenskapene til PyTorch DDP på ​​SageMaker Training manuelt. I dag er SageMaker Training glad for å gi en sømløs opplevelse for kunder som bruker PyTorch DDP-koden.

For å bruke dette effektivt, trenger du ikke gjøre noen endringer i treningsskriptene dine.

Du kan se denne nye parameteren i følgende kode. I distribution parameter, bare legg til pytorchddp og sett aktivert som true.

estimator = PyTorch(
    base_job_name="pytorch-dataparallel-mnist",
    source_dir="code",
    entry_point = "my_model.py",
    ... 
    # Training using SMDataParallel Distributed Training Framework
    distribution = {"pytorchddp": {"enabled": "true"}}
)

Denne nye konfigurasjonen starter på SageMaker Python SDK versjoner 2.102.0 og PyTorch DLCs 1.11.

For PyTorch DDP-utviklere som er kjent med det populære fakkelløp rammeverket, er det nyttig å vite at dette ikke er nødvendig i SageMaker treningsmiljø, som allerede gir robust feiltoleranse. For å minimere kodeomskrivinger kan du imidlertid ta med et annet oppstartsskript som kjører denne kommandoen som inngangspunkt.

Nå kan PyTorch-utviklere enkelt flytte skriptene sine til SageMaker, og sikre at skriptene og beholderne deres kan kjøre sømløst på tvers av flere datamiljøer.

Dette forbereder dem til i fremtiden å dra nytte av SageMakers distribuerte treningsbiblioteker som gir AWS-optimaliserte treningstopologier for å levere opptil 40 % fremskyndelsesforbedringer. For PyTorch-utviklere er dette en enkelt linje med kode! For PyTorch DDP-kode kan du ganske enkelt sette backend til smddp i initialiseringen (se Endre et PyTorch-treningsskript), som vist i følgende kode:

import smdistributed.dataparallel.torch.torch_smddp
import torch.distributed as dist
dist.init_process_group(backend='smddp')

Som vi så ovenfor, kan du også sette bakenden av DDPStrategy til smddp når du bruker Lightning. Dette kan føre til opptil 40 % generelle hastighetsøkninger for store klynger! For å lære mer om distribuert opplæring på SageMaker se vårt on-demand webinar, støtte notatbøker, aktuell dokumentasjonog papirer.

konklusjonen

I dette innlegget introduserte vi to nye funksjoner i SageMaker Training-familien. Disse gjør det mye enklere for PyTorch-utviklere å bruke sin eksisterende kode på SageMaker, både PyTorch DDP og PyTorch Lightning.

Vi viste også hvordan Amazon Search bruker SageMaker Training for å trene deres dyplæringsmodeller, og spesielt PyTorch Lightning med det SageMaker Data Parallel-optimaliserte kollektive biblioteket som backend. Å gå til distribuert trening generelt hjalp Amazon Search med å oppnå 7.3 ganger raskere togtider.

Om forfatterne

Emily Webber ble med i AWS like etter at SageMaker ble lansert, og har prøvd å fortelle verden om det siden den gang! Utenom å bygge nye ML-opplevelser for kunder, liker Emily å meditere og studere tibetansk buddhisme.

Karan Dhiman er programvareutviklingsingeniør ved AWS, basert i Toronto, Canada. Han er veldig lidenskapelig opptatt av maskinlæringsrom og å bygge løsninger for å akselerere distribuert dataarbeid.

Vishwa Karia er programvareutviklingsingeniør hos AWS Deep Engine. Hennes interesser ligger i skjæringspunktet mellom maskinlæring og distribuerte systemer, og hun brenner også for å styrke kvinner innen teknologi og AI.

Eiman Elnahrawy er en hovedprogramvareingeniør hos Amazon Search som leder arbeidet med maskinlæringsakselerasjon, skalering og automatisering. Hennes ekspertise spenner over flere områder, inkludert maskinlæring, distribuerte systemer og personalisering.