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

Så meget data, så lidt tid. Maskinlæringseksperter (ML), dataforskere, ingeniører og entusiaster er stødt på dette problem verden over. Fra naturlig sprogbehandling til computervision, tabelform til tidsserier og alt derimellem, har det ældgamle problem med at optimere til hastighed, når du kører data mod så mange GPU'er, som du kan få, inspireret til utallige løsninger. I dag er vi glade for at kunne annoncere funktioner til PyTorch-udviklere, der bruger native open source frameworks, som f.eks. PyTorch Lightning , PyTorch DDP, der vil strømline deres vej til skyen.

Amazon SageMaker er en fuldt administreret service til ML, og SageMaker modeltræning er et optimeret computermiljø til højtydende træning i skala. SageMaker modeltræning tilbyder en fjerntræningsoplevelse med et problemfrit kontrolplan for nemt at træne og reproducere ML-modeller til høj ydeevne og lave omkostninger. Vi er glade for at annoncere nye funktioner i SageMaker-træningsporteføljen, der gør det endnu nemmere og mere tilgængeligt at køre PyTorch i skala:

1. PyTorch Lightning kan nu integreres til SageMakers distribuerede data parallelle bibliotek med kun én linje kodeændring.
2. SageMaker model træning har nu understøttelse af native PyTorch Distributed Data Parallel med NCCL-backend, hvilket gør det muligt for udviklere at migrere til SageMaker nemmere end nogensinde før.

I dette indlæg diskuterer vi disse nye funktioner og lærer også, hvordan Amazon Search har kørt PyTorch Lightning med den optimerede distribuerede træningsbackend i SageMaker for at fremskynde modeltræningstiden.

Inden vi dykker ned i Amazon Search-casestudiet, for dem, der ikke er bekendte, vil vi gerne give lidt baggrund om SageMakers distribuerede data parallelle bibliotek. I 2020 udviklede og lancerede vi en tilpasset klyngekonfiguration til distribueret gradientnedstigning i skala, der øger den samlede klyngeeffektivitet, introduceret på Amazon Science som Sild. Ved at bruge det bedste fra både parameterservere og ringbaserede topologier er SageMaker Distributed Data Parallel (SMDDP) optimeret til Amazon Elastic Compute Cloud (Amazon EC2) netværkstopologi, herunder EFA. For større klyngestørrelser er SMDDP i stand til at levere 20-40% gennemløbsforbedringer i forhold til Horovod (TensorFlow) og PyTorch Distributed Data Parallel. Til mindre klyngestørrelser og understøttede modeller anbefaler vi SageMaker Training Compiler, hvilket er i stand til at reducere den samlede jobtid med op til 50 %.

Kundens fokus: PyTorch Lightning på SageMakers optimerede backend med Amazon Search

Amazon Search er ansvarlig for søge- og opdagelsesoplevelsen på Amazon.com. Det driver søgeoplevelsen for kunder, der leder efter produkter til at købe på Amazon. På et højt niveau opbygger Amazon Search et indeks for alle produkter, der sælges på Amazon.com. Når en kunde indtaster en forespørgsel, bruger Amazon Search derefter en række forskellige ML-teknikker, herunder deep learning-modeller, til at matche relevante og interessante produkter til kundens forespørgsel. Derefter rangerer den produkterne, før den viser resultaterne til kunden.

Amazon Search-forskere har brugt PyTorch Lightning som en af ​​hovedrammerne til at træne de dybe læringsmodeller, der driver søgerangeringen på grund af dens tilføjede brugervenlighedsfunktioner oven på PyTorch. SMDDP blev ikke understøttet for deep learning-modeller skrevet i PyTorch Lightning før denne nye SageMaker-lancering. Dette forhindrede Amazon Search-forskere, der foretrækker at bruge PyTorch Lightning, i at skalere deres modeltræning ved hjælp af dataparallelle teknikker, hvilket reducerede deres træningstid betydeligt og forhindrede dem i at teste nye eksperimenter, der kræver mere skalerbar træning.

Holdets tidlige benchmarking-resultater viser 7.3 gange hurtigere træningstid for en prøvemodel, når det trænes på otte knudepunkter sammenlignet med en træningsbaseline med en enkelt knude. Grundlinjemodellen, der bruges i disse benchmarking, er et flerlags perceptron-neuralt netværk med syv tætte fuldt forbundne lag og over 200 parametre. Følgende tabel opsummerer benchmarking-resultatet på ml.p3.16xlarge SageMaker-træningsforekomster.

Dernæst dykker vi ned i detaljerne om de nye lanceringer. Hvis du har lyst, kan du gå gennem vores tilsvarende eksempel notesbog.

Kør PyTorch Lightning med SageMaker-distribuerede træningsbibliotek

Vi er glade for at kunne meddele, at SageMaker Data Parallel nu problemfrit integreres med PyTorch Lightning i SageMaker-uddannelsen.

PyTorch Lightning er en open source-ramme, der giver en forenkling til at skrive brugerdefinerede modeller i PyTorch. På nogle måder, der ligner, hvad Keras gjorde for TensorFlow, eller endda velsagtens Hugging Face, giver PyTorch Lightning en API på højt niveau med abstraktioner for meget af funktionaliteten på lavere niveau i selve PyTorch. Dette omfatter definering af modellen, profilering, evaluering, beskæring, modelparallelisme, hyperparameterkonfigurationer, overførselslæring og mere.

Tidligere var PyTorch Lightning-udviklere usikre på, hvordan de problemfrit skulle migrere deres træningskode til højtydende SageMaker GPU-klynger. Derudover var der ingen måde for dem at drage fordel af effektivitetsgevinster introduceret af SageMaker Data Parallel.

For PyTorch Lightning, generelt set, burde der være små eller ingen kodeændringer for blot at køre disse API'er på SageMaker Training. I eksemplet notesbøger bruger vi DDP-strategi , DDPPlugin metoder.

Der er tre trin til at bruge PyTorch Lightning med SageMaker Data Parallel som en optimeret backend:

1. Brug en understøttet AWS Deep Learning Container (DLC) som dit basisbillede, eller opret eventuelt din egen container og installer SageMaker Data Parallel-backend selv. Sørg for, at du har PyTorch Lightning inkluderet i dine nødvendige pakker, såsom med en requirements.txt fil.
2. Foretag et par mindre kodeændringer til dit træningsscript, der aktiverer den optimerede backend. Disse omfatter:
   1. Importer SM DDP-biblioteket:

      import smdistributed.dataparallel.torch.torch_smddp
      

   2. Konfigurer PyTorch Lightning-miljøet til 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 bruger en version af PyTorch Lightning ældre end 1.5.10, skal du tilføje et par trin mere.
      1. Tilføj først miljøvariablen:

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

      2. For det andet, sørg for at bruge DDPPlugin, Snarere end DDPStrategy. Hvis du bruger en nyere version, som du nemt kan indstille ved at placere requirements.txt i source_dir for dit job, så er dette ikke nødvendigt. 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 din procesgruppe-backend som "smddp" i DDPSTrategy objekt. Men hvis du bruger PyTorch Lightning med PyTorch DDP-backend, som også understøttes, skal du blot fjerne 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 distributionsmetode noteret i estimatoren, som f.eks distribution={"smdistributed":{"dataparallel":{"enabled":True} hvis du bruger Sild-backend, eller distribution={"pytorchddp":{"enabled":True}.

• For en komplet liste over passende parametre i distribution parameter, se vores dokumentation link..

Nu kan du starte dit SageMaker træningsjob! Du kan starte dit træningsjob via Python SDK, Boto3, SageMaker-konsollen, AWS kommandolinjegrænseflade (AWS CLI), og utallige andre metoder. Fra et AWS-perspektiv er dette en enkelt API-kommando: create-training-job. Uanset om du starter denne kommando fra din lokale terminal, en AWS Lambda funktion, en Amazon SageMaker Studio notebook, en KubeFlow-pipeline eller ethvert andet computermiljø er helt op til dig.

Bemærk venligst, at integrationen mellem PyTorch Lightning og SageMaker Data Parallel i øjeblikket kun understøttes for nyere versioner af PyTorch, fra 1.11. Derudover er denne udgivelse kun tilgængelig i AWS DLC'erne til SageMaker fra PyTorch 1.12. Sørg for at pege på dette billede som din base. I us-east-1, denne adresse 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'

Så kan du udvide din Docker-beholder ved at bruge dette som dit basisbillede, eller du kan overføre dette som en variabel til image_uri argument af SageMaker træningsvurdering.

Som et resultat vil du være i stand til at køre din PyTorch Lightning-kode på SageMaker Trainings optimerede GPU'er, med den bedste ydeevne tilgængelig på AWS.

Kør PyTorch Distributed Data Parallel på SageMaker

The biggest problem PyTorch Distributed Data Parallel (DDP) løser er bedragerisk simpelt: hastighed. En god distribueret træningsramme bør give stabilitet, pålidelighed og vigtigst af alt, fremragende ydeevne i skala. PyTorch DDP leverer på dette ved at give fakkeludviklere API'er til at replikere deres modeller over flere GPU-enheder, i både enkelt-node og multi-node indstillinger. Rammeværket styrer derefter sønderdeling af forskellige objekter fra træningsdatasættet til hver modelkopi, idet der tages et gennemsnit af gradienterne for hver af modelkopierne for at synkronisere dem ved hvert trin. Dette giver én model ved den samlede afslutning af den fulde træningskørsel. Følgende diagram illustrerer denne proces.

PyTorch DDP er almindelig i projekter, der bruger store datasæt. Den præcise størrelse af hvert datasæt vil variere meget, men en generel retningslinje er at skalere datasæt, beregne størrelser og modelstørrelser i lignende forhold. Også kaldet skaleringslove, den optimale kombination af disse tre er meget til debat og vil variere baseret på applikationer. Hos AWS kan vi, baseret på at arbejde med flere kunder, tydeligt se fordele ved dataparallelle strategier, når en samlet datasætstørrelse er mindst et par tiere af GB. Når datasættene bliver endnu større, er implementering af en eller anden type dataparallelstrategi en kritisk teknik til at fremskynde det overordnede eksperiment og forbedre din tid til værdi.

Tidligere manglede kunder, der brugte PyTorch DDP til distribueret træning i lokaler eller i andre computermiljøer, en ramme til nemt at migrere deres projekter til SageMaker Training for at drage fordel af højtydende GPU'er med et problemfrit kontrolplan. Specifikt skulle de enten migrere deres parallelle dataramme til SMDDP eller udvikle og teste funktionerne i PyTorch DDP på ​​SageMaker Training manuelt. I dag er SageMaker Training glade for at give en problemfri oplevelse for kunder, der onboarder deres PyTorch DDP-kode.

For at bruge dette effektivt behøver du ikke at foretage ændringer i dine træningsscripts.

Du kan se denne nye parameter i den følgende kode. I den distribution parameter, skal du blot tilføje pytorchddp og sæt aktiveret 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 konfiguration starter ved SageMaker Python SDK version 2.102.0 og PyTorch DLC's 1.11.

Til PyTorch DDP-udviklere, der er fortrolige med det populære fakkelløb ramme, er det nyttigt at vide, at dette ikke er nødvendigt på SageMaker træningsmiljøet, som allerede giver robust fejltolerance. For at minimere kodeomskrivninger kan du dog medbringe et andet launcher-script, der kører denne kommando, som dit indgangspunkt.

Nu kan PyTorch-udviklere nemt flytte deres scripts til SageMaker og sikre, at deres scripts og containere kan køre problemfrit på tværs af flere computermiljøer.

Dette forbereder dem til i fremtiden at drage fordel af SageMakers distribuerede træningsbiblioteker der giver AWS-optimerede træningstopologier til at levere op til 40 % fremskyndelsesforbedringer. For PyTorch-udviklere er dette en enkelt kodelinje! For PyTorch DDP-kode kan du blot indstille backend til smddp i initialiseringen (se Rediger et PyTorch-træningsscript), 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å indstille backend af DDPStrategy til smddp når du bruger Lightning. Dette kan føre til op til 40% samlede speedups til store klynger! For at lære mere om distribueret træning på SageMaker se vores on-demand webinar, støtte notesbøger, relevant dokumentationog papirer.

Konklusion

I dette indlæg introducerede vi to nye funktioner inden for SageMaker Training-familien. Disse gør det meget nemmere for PyTorch-udviklere at bruge deres eksisterende kode på SageMaker, både PyTorch DDP og PyTorch Lightning.

Vi viste også, hvordan Amazon Search bruger SageMaker Training til at træne deres deep learning-modeller, og i særdeleshed PyTorch Lightning med det SageMaker Data Parallel-optimerede kollektive bibliotek som backend. Overgangen til distribueret træning hjalp Amazon Search med at opnå 7.3 gange hurtigere togtider.

Om forfatterne

Emily Webber sluttede sig til AWS lige efter SageMaker lancerede, og har forsøgt at fortælle verden om det lige siden! Ud over at bygge nye ML-oplevelser for kunder, nyder Emily at meditere og studere tibetansk buddhisme.

Karan Dhiman er softwareudviklingsingeniør hos AWS med base i Toronto, Canada. Han er meget passioneret omkring Machine Learning-rum og opbygning af løsninger til at accelerere distribuerede computerarbejdsbelastninger.

Vishwa Karia er softwareudviklingsingeniør hos AWS Deep Engine. Hendes interesser ligger i skæringspunktet mellem maskinlæring og distribuerede systemer, og hun brænder også for at styrke kvinder inden for teknologi og kunstig intelligens.

Eiman Elnahrawy er en hovedsoftwareingeniør hos Amazon Search, der leder indsatsen for acceleration, skalering og automatisering af maskinlæring. Hendes ekspertise spænder over flere områder, herunder Machine Learning, Distributed Systems og Personalization.