Voer PyTorch Lightning en native PyTorch DDP uit op Amazon SageMaker Training, met Amazon Search

Zo veel data, zo weinig tijd. Machine learning (ML)-experts, datawetenschappers, ingenieurs en enthousiastelingen hebben dit probleem over de hele wereld ondervonden. Van natuurlijke taalverwerking tot computervisie, tabellen tot tijdreeksen en alles daartussenin, het eeuwenoude probleem van het optimaliseren voor snelheid bij het uitvoeren van gegevens tegen zoveel GPU's als je kunt hebben, heeft talloze oplossingen geïnspireerd. Vandaag kondigen we met genoegen functies aan voor PyTorch-ontwikkelaars die native open-source frameworks gebruiken, zoals: PyTorch Bliksem en PyTorch-DDP, dat hun pad naar de cloud zal stroomlijnen.

Amazon Sage Maker is een volledig beheerde service voor ML, en SageMaker-modeltraining is een geoptimaliseerde rekenomgeving voor hoogwaardige training op schaal. SageMaker-modeltraining biedt een trainingservaring op afstand met een naadloos besturingsvlak om ML-modellen eenvoudig te trainen en te reproduceren tegen hoge prestaties en lage kosten. We zijn verheugd om nieuwe functies in het SageMaker-trainingsportfolio aan te kondigen die het uitvoeren van PyTorch op grote schaal nog eenvoudiger en toegankelijker maken:

1. PyTorch Lightning kan nu worden geïntegreerd in: SageMaker's gedistribueerde parallelle databibliotheek met slechts één regel codewijziging.
2. SageMaker-modeltraining heeft nu ondersteuning voor native PyTorch Distributed Data Parallel met NCCL-backend, waardoor ontwikkelaars gemakkelijker dan ooit tevoren naar SageMaker kunnen migreren.

In dit bericht bespreken we deze nieuwe functies en leren we ook hoe Amazon Search PyTorch Lightning heeft uitgevoerd met de geoptimaliseerde gedistribueerde trainingsbackend in SageMaker om de trainingstijd van het model te versnellen.

Voordat we ingaan op de Amazon Search-casestudy, willen we voor degenen die niet bekend zijn, wat achtergrondinformatie geven over SageMaker's gedistribueerde parallelle databibliotheek. In 2020 hebben we een aangepaste clusterconfiguratie ontwikkeld en gelanceerd voor gedistribueerde gradiëntafdaling op schaal die de algehele clusterefficiëntie verhoogt, geïntroduceerd op Amazon Science als Haring. Door gebruik te maken van het beste van zowel parameterservers als ringgebaseerde topologieën, is SageMaker Distributed Data Parallel (SMDDP) geoptimaliseerd voor de Amazon Elastic Compute-cloud (Amazon EC2) netwerktopologie, inclusief: EFA. Voor grotere clustergroottes kan SMDDP 20-40% doorvoerverbeteringen leveren ten opzichte van Horovod (TensorFlow) en PyTorch Distributed Data Parallel. Voor kleinere clustergroottes en ondersteunde modellen raden we de: SageMaker-trainingscompiler, die in staat is om de totale werktijd tot 50% te verminderen.

Spotlight van de klant: PyTorch Lightning op de geoptimaliseerde backend van SageMaker met Amazon Search

Amazon Search is verantwoordelijk voor de zoek- en ontdekkingservaring op Amazon.com. Het maakt de zoekervaring mogelijk voor klanten die op zoek zijn naar producten om op Amazon te kopen. Op een hoog niveau bouwt Amazon Search een index op voor alle producten die op Amazon.com worden verkocht. Wanneer een klant een zoekopdracht invoert, gebruikt Amazon Search vervolgens een verscheidenheid aan ML-technieken, waaronder deep learning-modellen, om relevante en interessante producten af ​​te stemmen op de vraag van de klant. Vervolgens rangschikt het de producten voordat het de resultaten aan de klant laat zien.

Wetenschappers van Amazon Search hebben PyTorch Lightning gebruikt als een van de belangrijkste frameworks om de deep learning-modellen te trainen die de zoekrangschikking versterken vanwege de toegevoegde bruikbaarheidsfuncties bovenop PyTorch. SMDDP werd niet ondersteund voor deep learning-modellen die zijn geschreven in PyTorch Lightning vóór deze nieuwe lancering van SageMaker. Dit verhinderde Amazon Search-wetenschappers die liever PyTorch Lightning gebruiken om hun modeltraining te schalen met behulp van dataparallelle technieken, waardoor hun trainingstijd aanzienlijk werd vertraagd en ze geen nieuwe experimenten konden testen die meer schaalbare training vereisen.

De vroege benchmarkresultaten van het team laten een 7.3 keer snellere trainingstijd zien voor een voorbeeldmodel wanneer het wordt getraind op acht knooppunten in vergelijking met een basislijn voor training met één knooppunt. Het basismodel dat bij deze benchmarking wordt gebruikt, is een meerlagig perceptron-neuraal netwerk met zeven dichte, volledig verbonden lagen en meer dan 200 parameters. De volgende tabel geeft een overzicht van het benchmarkresultaat op ml.p3.16xgrote SageMaker-trainingsinstanties.

Vervolgens duiken we in de details van de nieuwe lanceringen. Als je wilt, kun je door onze corresponderende voorbeeld notitieboekje.

Voer PyTorch Lightning uit met de door SageMaker gedistribueerde trainingsbibliotheek

We zijn verheugd aan te kondigen dat SageMaker Data Parallel nu naadloos integreert met PyTorch Lightning binnen SageMaker-training.

PyTorch Lightning is een open-source framework dat een vereenvoudiging biedt voor het schrijven van aangepaste modellen in PyTorch. In sommige opzichten vergelijkbaar met wat Keras deed voor TensorFlow, of zelfs aantoonbaar Hugging Face, biedt PyTorch Lightning een API op hoog niveau met abstracties voor veel van de lagere functionaliteit van PyTorch zelf. Dit omvat het definiëren van het model, profilering, evaluatie, snoeien, modelparallellisme, hyperparameterconfiguraties, transfer learning en meer.

Voorheen wisten PyTorch Lightning-ontwikkelaars niet zeker hoe ze hun trainingscode naadloos konden migreren naar krachtige SageMaker GPU-clusters. Bovendien konden ze op geen enkele manier profiteren van de efficiëntieverbeteringen die door SageMaker Data Parallel werden geïntroduceerd.

Voor PyTorch Lightning zouden er over het algemeen weinig tot geen codewijzigingen moeten zijn om deze API's eenvoudig op SageMaker Training uit te voeren. In de voorbeeldnotitieboeken gebruiken we de DDPStrategie en DDP-plugin werkwijzen.

Er zijn drie stappen om PyTorch Lightning met SageMaker Data Parallel als geoptimaliseerde backend te gebruiken:

1. Gebruik een ondersteund AWS Deep Learning-container (DLC) als uw basisafbeelding, of maak optioneel uw eigen container en installeer de SageMaker Data Parallel-backend zelf. Zorg ervoor dat PyTorch Lightning is opgenomen in uw benodigde pakketten, zoals met a requirements.txt bestand.
2. Breng een paar kleine codewijzigingen aan in uw trainingsscript om de geoptimaliseerde backend mogelijk te maken. Waaronder:
   1. Importeer de SM DDP-bibliotheek:

      import smdistributed.dataparallel.torch.torch_smddp
      

   2. Stel de PyTorch Lightning-omgeving in voor 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. Als u een versie van PyTorch Lightning gebruikt die ouder is dan 1.5.10, moet u nog een paar stappen toevoegen.
      1. Voeg eerst de omgevingsvariabele toe:

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

      2. Ten tweede, zorg ervoor dat u gebruikt DDPPlugin, liever dan DDPStrategy. Als u een recentere versie gebruikt, kunt u deze eenvoudig instellen door de requirements.txt in de source_dir voor je werk, dan is dit niet nodig. Zie de volgende code:

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

   4. Definieer optioneel uw procesgroep-backend als "smddp" in de DDPSTrategy object. Als u echter PyTorch Lightning gebruikt met de PyTorch DDP-backend, die ook wordt ondersteund, verwijdert u deze `process_group_backend`-parameter. Zie de volgende code:

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

3. Zorg ervoor dat u een distributiemethode hebt genoteerd in de schatter, zoals: distribution={"smdistributed":{"dataparallel":{"enabled":True} als je de Herring-backend gebruikt, of distribution={"pytorchddp":{"enabled":True}.

• Voor een volledige lijst van geschikte parameters in de distribution parameter, zie onze documentatie hier.

Nu kunt u uw SageMaker-trainingstaak starten! U kunt uw trainingstaak starten via de Python SDK, Boto3, de SageMaker-console, de AWS-opdrachtregelinterface (AWS CLI), en talloze andere methoden. Vanuit een AWS-perspectief is dit: een enkele API-opdracht: create-training-job. Of u deze opdracht nu start vanaf uw lokale terminal, een AWS Lambda functie, een Amazon SageMaker Studio notebook, een KubeFlow-pijplijn of een andere computeromgeving is helemaal aan jou.

Houd er rekening mee dat de integratie tussen PyTorch Lightning en SageMaker Data Parallel momenteel alleen wordt ondersteund voor nieuwere versies van PyTorch, vanaf 1.11. Bovendien is deze release alleen beschikbaar in de AWS DLC's voor SageMaker vanaf PyTorch 1.12. Zorg ervoor dat u naar deze afbeelding wijst als uw basis. In us-east-1, dit adres is als volgt:

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

Vervolgens kunt u uw Docker-container uitbreiden door dit als uw basisimage te gebruiken, of u kunt dit als een variabele doorgeven aan de image_uri argument van de SageMaker-trainingschatter.

Als gevolg hiervan kunt u uw PyTorch Lightning-code uitvoeren op de geoptimaliseerde GPU's van SageMaker Training, met de beste prestaties die beschikbaar zijn op AWS.

Voer PyTorch Distributed Data Parallel uit op SageMaker

The biggest problem PyTorch gedistribueerde gegevens parallel (DDP) oplost is bedrieglijk eenvoudig: snelheid. Een goed gedistribueerd trainingskader moet stabiliteit, betrouwbaarheid en vooral uitstekende prestaties op grote schaal bieden. PyTorch DDP levert dit op door toortsontwikkelaars API's te bieden om hun modellen te repliceren over meerdere GPU-apparaten, in zowel single-node- als multi-node-instellingen. Het framework beheert vervolgens het sharden van verschillende objecten uit de trainingsgegevensset naar elke modelkopie, waarbij het gemiddelde van de gradiënten voor elk van de modelkopieën wordt genomen om ze bij elke stap te synchroniseren. Dit levert één model op bij de totale voltooiing van de volledige trainingsrun. Het volgende diagram illustreert dit proces.

PyTorch DDP is gebruikelijk in projecten die grote datasets gebruiken. De precieze grootte van elke dataset zal sterk variëren, maar een algemene richtlijn is om datasets, rekengroottes en modelgroottes in vergelijkbare verhoudingen te schalen. Ook wel genoemd schaalwetten, staat de optimale combinatie van deze drie ter discussie en zal variëren op basis van toepassingen. Bij AWS kunnen we, op basis van het werken met meerdere klanten, duidelijk voordelen zien van parallelle strategieën voor gegevens wanneer de totale gegevenssetgrootte ten minste enkele tientallen GB's is. Wanneer de datasets nog groter worden, is het implementeren van een soort parallelle datastrategie een cruciale techniek om het algehele experiment te versnellen en uw time-to-value te verbeteren.

Voorheen hadden klanten die PyTorch DDP gebruikten voor gedistribueerde training op locatie of in andere computeromgevingen, geen raamwerk om hun projecten eenvoudig naar SageMaker Training te migreren om te profiteren van krachtige GPU's met een naadloos besturingsvlak. Ze moesten met name ofwel hun parallelle gegevensframework migreren naar SMDDP, of de mogelijkheden van PyTorch DDP op SageMaker Training handmatig ontwikkelen en testen. Tegenwoordig biedt SageMaker Training graag een naadloze ervaring voor klanten die hun PyTorch DDP-code aan boord nemen.

Om dit effectief te gebruiken, hoeft u geen wijzigingen aan te brengen in uw trainingsscripts.

U kunt deze nieuwe parameter zien in de volgende code. In de distribution parameter, gewoon toevoegen pytorchddp en zet ingeschakeld als 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"}}
)

Deze nieuwe configuratie begint bij SageMaker Python SDK-versies 2.102.0 en PyTorch DLC's 1.11.

Voor PyTorch DDP-ontwikkelaars die bekend zijn met de populaire fakkeltocht framework, is het handig om te weten dat dit niet nodig is in de SageMaker-trainingsomgeving, die al een robuuste fouttolerantie biedt. Om het herschrijven van code te minimaliseren, kunt u echter een ander opstartscript gebruiken dat deze opdracht uitvoert als uw toegangspunt.

Nu kunnen PyTorch-ontwikkelaars hun scripts eenvoudig naar SageMaker verplaatsen, zodat hun scripts en containers naadloos in meerdere computeromgevingen kunnen worden uitgevoerd.

Dit bereidt hen voor om in de toekomst te profiteren van De gedistribueerde trainingsbibliotheken van SageMaker die AWS-geoptimaliseerde trainingstopologieën bieden om tot 40% versnellingsverbeteringen te leveren. Voor PyTorch-ontwikkelaars is dit een enkele regel code! Voor PyTorch DDP-code kunt u eenvoudig de backend instellen op: smddp in de initialisatie (zie Een PyTorch-trainingsscript wijzigen), zoals weergegeven in de volgende code:

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

Zoals we hierboven zagen, kun je ook de backend van DDPStrategy naar smddp bij gebruik van Lightning. Dit kan leiden tot maximaal 40% algehele versnellingen voor grote clusters! Zie voor meer informatie over gedistribueerde training op SageMaker: ons on-demand webinar, ondersteunend laptops, relevant documentatie en papieren.

Conclusie

In dit bericht hebben we twee nieuwe functies geïntroduceerd binnen de SageMaker Training-familie. Deze maken het voor PyTorch-ontwikkelaars veel gemakkelijker om hun bestaande code op SageMaker te gebruiken, zowel PyTorch DDP als PyTorch Lightning.

We hebben ook laten zien hoe Amazon Search SageMaker Training gebruikt voor het trainen van hun deep learning-modellen, en in het bijzonder PyTorch Lightning met de door SageMaker Data Parallel geoptimaliseerde collectieve bibliotheek als backend. Door over te stappen op gedistribueerde training kon Amazon Search 7.3x snellere treintijden behalen.

Over de auteurs

Emily Webber kwam bij AWS net na de lancering van SageMaker en probeert sindsdien de wereld erover te vertellen! Naast het bouwen van nieuwe ML-ervaringen voor klanten, mediteert Emily graag en bestudeert ze het Tibetaans boeddhisme.

Karan Dhiman is een Software Development Engineer bij AWS, gevestigd in Toronto, Canada. Hij is erg gepassioneerd door Machine Learning-ruimte en het bouwen van oplossingen voor het versnellen van gedistribueerde computerworkloads.

Vishwa Karia is een Software Development Engineer bij AWS Deep Engine. Haar interesses liggen op het snijvlak van machine learning en gedistribueerde systemen en ze is ook gepassioneerd over het versterken van vrouwen in technologie en AI.

Eiman Elnahrawy is een Principal Software Engineer bij Amazon Search en leidt de inspanningen op het gebied van versnelling, schaling en automatisering van Machine Learning. Haar expertise omvat meerdere gebieden, waaronder machine learning, gedistribueerde systemen en personalisatie.