Kör PyTorch Lightning och inbyggd PyTorch DDP på ​​Amazon SageMaker Training, med Amazon Search

Så mycket data, så lite tid. Maskininlärningsexperter (ML), datavetare, ingenjörer och entusiaster har stött på detta problem över hela världen. Från naturlig språkbehandling till datorseende, tabeller till tidsserier och allt däremellan, det urgamla problemet med att optimera för hastighet när du kör data mot så många GPU:er som du kan få har inspirerat till otaliga lösningar. Idag är vi glada att kunna presentera funktioner för PyTorch-utvecklare som använder inbyggda ramverk med öppen källkod, som PyTorch Lightning och PyTorch DDP, som kommer att effektivisera deras väg till molnet.

Amazon SageMaker är en helt hanterad tjänst för ML, och SageMaker modellträning är en optimerad datormiljö för högpresterande träning i stor skala. SageMaker modellträning erbjuder en fjärrträningsupplevelse med ett sömlöst kontrollplan för att enkelt träna och reproducera ML-modeller till hög prestanda och låg kostnad. Vi är glada över att kunna presentera nya funktioner i SageMakers utbildningsportfölj som gör att köra PyTorch i skala ännu enklare och mer tillgängligt:

1. PyTorch Lightning kan nu integreras till SageMakers distribuerade data parallellbibliotek med endast en rad kodändring.
2. SageMaker modellutbildning har nu stöd för inbyggd PyTorch Distributed Data Parallel med NCCL-backend, vilket gör att utvecklare kan migrera till SageMaker enklare än någonsin tidigare.

I det här inlägget diskuterar vi dessa nya funktioner, och lär oss också hur Amazon Search har kört PyTorch Lightning med den optimerade distribuerade träningsbackend i SageMaker för att påskynda modellträningstiden.

Innan vi dyker in i fallstudien Amazon Search, för dem som inte är bekanta skulle vi vilja ge lite bakgrund om SageMakers distribuerade data parallellbibliotek. Under 2020 utvecklade och lanserade vi en anpassad klusterkonfiguration för distribuerad gradientnedstigning i skala som ökar den totala klustereffektiviteten, introducerad på Amazon Science som Sill. Genom att använda det bästa av både parameterservrar och ringbaserade topologier är SageMaker Distributed Data Parallel (SMDDP) optimerad för Amazon Elastic Compute Cloud (Amazon EC2) nätverkstopologi, inklusive EFA. För större klusterstorlekar kan SMDDP leverera 20–40 % genomströmningsförbättringar jämfört med Horovod (TensorFlow) och PyTorch Distributed Data Parallel. För mindre klusterstorlekar och modeller som stöds rekommenderar vi SageMaker Training Compiler, vilket kan minska den totala arbetstiden med upp till 50 %.

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

Amazon Search ansvarar för sök- och upptäcktsupplevelsen på Amazon.com. Det driver sökupplevelsen för kunder som letar efter produkter att köpa på Amazon. På en hög nivå bygger Amazon Search ett index för alla produkter som säljs på Amazon.com. När en kund anger en fråga använder Amazon Search sedan en mängd olika ML-tekniker, inklusive modeller för djupinlärning, för att matcha relevanta och intressanta produkter till kundfrågan. Sedan rangordnar den produkterna innan resultatet visas för kunden.

Amazon Search-forskare har använt PyTorch Lightning som ett av huvudramarna för att träna de djupinlärningsmodeller som driver sökrankningen på grund av dess extra användbarhetsfunktioner ovanpå PyTorch. SMDDP stöddes inte för modeller för djupinlärning skrivna i PyTorch Lightning före denna nya SageMaker-lansering. Detta hindrade Amazon Search-forskare som föredrar att använda PyTorch Lightning från att skala sin modellträning med hjälp av parallella datatekniker, vilket avsevärt saktar ner sin träningstid och hindrar dem från att testa nya experiment som kräver mer skalbar träning.

Teamets tidiga benchmarking-resultat visar 7.3 gånger snabbare träningstid för en provmodell när den tränas på åtta noder jämfört med en baslinje för träning med en nod. Baslinjemodellen som används i dessa benchmarking är ett flerlagers perceptron-neuralt nätverk med sju täta helt sammankopplade lager och över 200 parametrar. Följande tabell sammanfattar benchmarking-resultatet på ml.p3.16xlarge SageMaker-träningsinstanser.

Därefter dyker vi in ​​i detaljerna om de nya lanseringarna. Om du vill kan du gå igenom vår motsvarande exempel anteckningsbok.

Kör PyTorch Lightning med SageMakers distribuerade träningsbibliotek

Vi är glada att kunna meddela att SageMaker Data Parallel nu sömlöst integreras med PyTorch Lightning inom SageMaker-utbildningen.

PyTorch Lightning är ett ramverk med öppen källkod som ger en förenkling för att skriva anpassade modeller i PyTorch. På vissa sätt liknande vad Keras gjorde för TensorFlow, eller till och med Hugging Face, tillhandahåller PyTorch Lightning ett högnivå-API med abstraktioner för mycket av funktionaliteten på lägre nivå i själva PyTorch. Detta inkluderar att definiera modellen, profilering, utvärdering, beskärning, modellparallellism, hyperparameterkonfigurationer, överföringsinlärning och mer.

Tidigare var PyTorch Lightning-utvecklare osäkra på hur de sömlöst skulle migrera sin träningskod till högpresterande SageMaker GPU-kluster. Dessutom fanns det inget sätt för dem att dra fördel av effektivitetsvinster introducerade av SageMaker Data Parallel.

För PyTorch Lightning, generellt sett, borde det finnas små eller inga kodändringar för att helt enkelt köra dessa API:er på SageMaker Training. I exemplet anteckningsböcker använder vi DDP-strategi och DDPPlugin metoder.

Det finns tre steg för att använda PyTorch Lightning med SageMaker Data Parallel som en optimerad backend:

1. Använd en stödd AWS Deep Learning Container (DLC) som din basbild, eller skapa din egen behållare och installera SageMaker Data Parallel-backend själv. Se till att du har PyTorch Lightning inkluderat i dina nödvändiga paket, till exempel med en requirements.txt fil.
2. Gör några mindre kodändringar i ditt träningsskript som möjliggör den optimerade backend. Dessa inkluderar:
   1. Importera SM DDP-biblioteket:

      import smdistributed.dataparallel.torch.torch_smddp
      

   2. Ställ in PyTorch Lightning-miljön för 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. Om du använder en version av PyTorch Lightning äldre än 1.5.10, måste du lägga till några fler steg.
      1. Lägg först till miljövariabeln:

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

      2. För det andra, se till att du använder DDPPlugin, Snarare än DDPStrategy. Om du använder en nyare version, som du enkelt kan ställa in genom att placera requirements.txt i source_dir för ditt jobb, då är detta inte nödvändigt. Se följande kod:

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

   4. Alternativt kan du definiera din processgrupps backend som "smddp" i DDPSTrategy objekt. Men om du använder PyTorch Lightning med PyTorch DDP-backend, som också stöds, ta helt enkelt bort denna `process_group_backend` parameter. Se följande kod:

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

3. Se till att du har en distributionsmetod noterad i estimatorn, som t.ex distribution={"smdistributed":{"dataparallel":{"enabled":True} om du använder sillbackend, eller distribution={"pytorchddp":{"enabled":True}.

• För en fullständig lista över lämpliga parametrar i distribution parameter, se vår dokumentation här..

Nu kan du starta ditt SageMaker-utbildningsjobb! Du kan starta ditt träningsjobb via Python SDK, Boto3, SageMaker-konsolen, AWS-kommandoradsgränssnitt (AWS CLI), och otaliga andra metoder. Ur ett AWS-perspektiv är detta ett enda API-kommando: create-training-job. Oavsett om du startar detta kommando från din lokala terminal, en AWS Lambda funktion, en Amazon SageMaker Studio notebook, en KubeFlow-pipeline eller någon annan datormiljö är helt upp till dig.

Observera att integrationen mellan PyTorch Lightning och SageMaker Data Parallel för närvarande endast stöds för nyare versioner av PyTorch, från och med 1.11. Dessutom är den här utgåvan endast tillgänglig i AWS DLC för SageMaker från och med PyTorch 1.12. Se till att du pekar på den här bilden som bas. I us-east-1, denna adress är följande:

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

Sedan kan du utöka din Docker-behållare med denna som basbild, eller så kan du skicka den som en variabel till image_uri argument av SageMaker träningsuppskattare.

Som ett resultat kommer du att kunna köra din PyTorch Lightning-kod på SageMaker Trainings optimerade GPU:er, med den bästa prestandan som finns tillgänglig på AWS.

Kör PyTorch Distributed Data Parallel på SageMaker

Det största problemet PyTorch Distributed Data Parallel (DDP) löser är bedrägligt enkelt: hastighet. En bra distribuerad träningsram bör ge stabilitet, tillförlitlighet och viktigast av allt, utmärkt prestation i skala. PyTorch DDP levererar detta genom att förse torch-utvecklare med API:er för att replikera sina modeller över flera GPU-enheter, i både enkelnods- och multinodsinställningar. Ramverket hanterar sedan skärning av olika objekt från träningsdatauppsättningen till varje modellkopia, genom att snitta gradienterna för var och en av modellkopiorna för att synkronisera dem vid varje steg. Detta ger en modell vid det totala slutförandet av hela träningskörningen. Följande diagram illustrerar denna process.

PyTorch DDP är vanligt i projekt som använder stora datamängder. Den exakta storleken på varje datauppsättning kommer att variera kraftigt, men en allmän riktlinje är att skala datauppsättningar, beräkna storlekar och modellstorlekar i liknande förhållanden. Även kallad skalningslagar, den optimala kombinationen av dessa tre är mycket uppe för debatt och kommer att variera beroende på applikationer. På AWS, baserat på att arbeta med flera kunder, kan vi tydligt se fördelarna med dataparallella strategier när en total datauppsättningsstorlek är minst några tiotals GB. När datamängderna blir ännu större är implementering av någon typ av dataparallellstrategi en kritisk teknik för att påskynda det övergripande experimentet och förbättra din time to value.

Tidigare saknade kunder som använde PyTorch DDP för distribuerad utbildning i lokaler eller i andra datormiljöer ett ramverk för att enkelt migrera sina projekt till SageMaker Training för att dra fördel av högpresterande GPU:er med ett sömlöst kontrollplan. Specifikt behövde de antingen migrera sitt dataparallella ramverk till SMDDP eller utveckla och testa funktionerna hos PyTorch DDP på ​​SageMaker Training manuellt. Idag erbjuder SageMaker Training en sömlös upplevelse för kunder som använder sin PyTorch DDP-kod.

För att använda detta effektivt behöver du inte göra några ändringar i dina träningsskript.

Du kan se denna nya parameter i följande kod. I den distribution parameter, lägg till pytorchddp och aktiverad 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"}}
)

Denna nya konfiguration börjar vid SageMaker Python SDK version 2.102.0 och PyTorch DLC:s 1.11.

För PyTorch DDP-utvecklare som är bekanta med det populära torchrun ramverk är det bra att veta att detta inte är nödvändigt i SageMakers träningsmiljö, som redan ger robust feltolerans. Men för att minimera kodomskrivningar kan du ta med ett annat startskript som kör detta kommando som din ingångspunkt.

Nu kan PyTorch-utvecklare enkelt flytta sina skript till SageMaker, vilket säkerställer att deras skript och behållare kan köras sömlöst i flera datormiljöer.

Detta förbereder dem att i framtiden dra nytta av SageMakers distribuerade träningsbibliotek som tillhandahåller AWS-optimerade träningstopologier för att leverera upp till 40 % snabbhetsförbättringar. För PyTorch-utvecklare är detta en enda rad kod! För PyTorch DDP-kod kan du helt enkelt ställa in backend till smddp i initieringen (se Ändra ett PyTorch-träningsskript), som visas i följande kod:

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

Som vi såg ovan kan du också ställa in backend för DDPStrategy till smddp när du använder Lightning. Detta kan leda till upp till 40% totala hastigheter för stora kluster! För att lära dig mer om distribuerad utbildning på SageMaker se vårt on-demand webinar, stöd bärbara datorer, relevant dokumentationoch papper.

Slutsats

I det här inlägget introducerade vi två nya funktioner inom SageMaker Training-familjen. Dessa gör det mycket lättare för PyTorch-utvecklare att använda sin befintliga kod på SageMaker, både PyTorch DDP och PyTorch Lightning.

Vi visade också hur Amazon Search använder SageMaker Training för att träna sina djupinlärningsmodeller, och i synnerhet PyTorch Lightning med det SageMaker Data Parallel-optimerade kollektiva biblioteket som backend. Övergången till distribuerad utbildning hjälpte Amazon Search att uppnå 7.3 gånger snabbare tågtider.

Om författarna

Emily Webber gick med i AWS precis efter att SageMaker lanserades, och har försökt berätta för världen om det sedan dess! Förutom att bygga nya ML-upplevelser för kunder tycker Emily om att meditera och studera tibetansk buddhism.

Karan Dhiman är en mjukvaruutvecklingsingenjör på AWS, baserad i Toronto, Kanada. Han brinner mycket för maskininlärningsutrymme och att bygga lösningar för att accelerera distribuerade datorarbeten.

Vishwa Karia är en mjukvaruutvecklingsingenjör på AWS Deep Engine. Hennes intressen ligger i skärningspunkten mellan maskininlärning och distribuerade system och hon brinner också för att stärka kvinnor inom teknik och AI.

Eiman Elnahrawy är en huvudprogramvaruingenjör på Amazon Search som leder arbetet med acceleration, skalning och automatisering av maskininlärning. Hennes expertis spänner över flera områden, inklusive maskininlärning, distribuerade system och personalisering.