Træn gigantiske modeller med nær-lineær skalering ved hjælp af sharded data parallelism på Amazon SageMaker

I jagten på overlegen nøjagtighed er deep learning-modeller inden for områder som naturlig sprogbehandling og computersyn vokset betydeligt i størrelse i de sidste par år, ofte talt i ti til hundreder af milliarder af parametre. Træning af disse gigantiske modeller er udfordrende og kræver komplekse distributionsstrategier. Datavidenskabsmænd og maskinlæringsingeniører leder konstant efter den bedste måde at optimere deres træningsberegning på, men kæmper alligevel med kommunikationsomkostningerne, der kan stige sammen med den samlede klyngestørrelse.

Det er derfor, vi for nylig lancerede shårdt dataparallelisme on Amazon SageMaker, en ny hukommelsesbesparende distribueret træningsteknik i SageMaker model parallel (SMP) bibliotek. Delt dataparallelisme er specialbygget til modeller i ekstrem skala og bruger Amazon internt MiCS teknologi under motorhjelmen, en videnskabelig indsats for at minimere kommunikationsskalaen ved at nedbringe dyre kommunikationsomkostninger forankret i parameterindsamling og gradientsynkronisering. Med en 30B parameter GPT-2 model med sekvenslængde 2048 opnåede denne nye funktion 141 TFLOP'er, en hastighed på 39.7 % sammenlignet med DeepSpeed ​​ZeRO-3. For en 10B GPT-2-model med sekvenslængde 512 opnåede denne nye funktion også 564 prøver i sekundet, en hastighed på 13.9 % sammenlignet med PyTorchs Fully Sharded Data Parallel (FSDP). Husk, at i gigantisk modeltræning oversættes hver procentdel af speedup til sparede dollars og opnået produktivitet i dit team.

I dette blogindlæg vil vi først se nærmere på de vigtigste differentiatorer af sharded data parallelism, og hvornår det skal bruges. Derefter lærer du, hvordan du nemt træner en 30B parameter GPT-2 model på SageMaker med denne nye funktion. Til sidst vil vi sammenligne ydeevnen med andre open source-muligheder, der især overgår DeepSpeed ​​ZeRO med op til 39.7 % på 256 GPU'er.

Hvordan fragmenteret dataparallelisme virker, og hvornår man skal bruge det

Før vi introducerer sharded data parallelism, lad os se på dens bredere teknikfamilie. Nylige distribuerede træningsmetoder for store modeller er flyttet til et paradigme, hvor modelparametre, gradienter og optimeringstilstande deles på tværs af dataparallelle noder. I modsætning til Pipeline Parallelism, som har den medfødte kompleksitet at vælge lag til at partitionere på tværs af enheder, især når dit framework ikke understøtter automatiseret modelopdeling, bevarer dette paradigme elegant enkelheden ved dataparallelisme, samtidig med at dataparallelismens begrænsning, hvor en model skal passe ind i en enkelt GPU, fjernes.

I eksisterende rammer, der falder ind under dette paradigme, især DeepSpeed ​​ZeRO-3 og PyTorchs FSDP upstreamet fra FairScale, er modeltilstande spredt ud over alle GPU'er, en strategi, der sænker hukommelsesforbruget på hver GPU på bekostning af store kommunikationsomkostninger, som øges med klyngestørrelsen og derfor får skalerbarheden til at falde betydeligt i skala. I modsætning hertil står sharded data parallelism i SMP-bibliotekspartitionsmodellen i en skalabevidst måde ved kun at opdele hver replika af modeltilstande indenfor en delmængde af GPU'er.

Lad os se nærmere på skalabevidst modelopdeling i MiCS, kerneteknologien bag sharded data parallel. Intuitionen bag dette design er, at opdeling af træningstilstande på tværs af hele den dataparallelle gruppe muligvis ikke er påkrævet for at træne en model med titusinder af milliarder af parametre. For eksempel er 8 V100 GPU'er (32 GB hver) tilstrækkelige til at holde modeltilstande replika af en 10B-parameter model, som har brug for omkring 200 GB hukommelse, når du træner med Adam optimizer ved brug af blandet præcision. Ved at begrænse en komplet replika af modeltilstande i mindste undersæt af GPU'er, kan vi effektivt reducere omfanget af kommunikationsomkostninger sammenlignet med DeepSpeed ​​og PyTorch FSDP. Sharded data parallel udnytter også andre teknikker i MiCS såsom hierarkisk kommunikation og 2-hop Gradient Synchronization. For mere information, tjek Nær-lineær skalering af gigantisk model træning på AWS or MiCS: Nær-lineær skalering til træning af gigantisk model på Public Cloud.

Nu, hvordan ved du, hvornår du skal vælge sharded data parallelt frem for andre distribuerede træningsteknikker? Den generelle regel er, at hvis din model har mindre end 1 milliard parametre og kan passe ind i GPU-hukommelse, SageMaker data parallelbibliotek or SageMaker træningskompiler kan være tilstrækkeligt for dig. Hvis du har større sprog- eller computervisionsmodeller, er vores forslag at træne det med sharded data parallelism-teknikken kombineret med aktiveringskontrol , aktivering aflastning i SageMaker model parallelbibliotek først, før andre teknikker som f.eks tensor parallelisme eller rørledningsparallelisme.

Brug af sharded data parallelism til at træne GPT-2 på Amazon SageMaker

Lad os nu lære, hvordan man træner en GPT-2-model med opdelte data parallelt, hvor SMP indkapsler kompleksiteten for dig. Dette komplet tutorial notesbog fører dig gennem hele processen, fra databehandling, definering og indsendelse af træningsjob til overvågning af træningslogfiler. Det følgende er en kort oversigt, der fremhæver de vigtigste trin til brug af denne funktion.

1. Kom igang

Delt dataparallelisme er tilgængelig i PyTorch v1.12.0+ og fungerer med både FP16 og BF16. Den nemmeste måde at bruge SMP-biblioteket på er gennem en forudbygget AWS Deep Learning Container til PyTorch. Men hvis du ønsker at medbringe din egen Docker container, kan du henvise til Opret din egen Docker-container med SageMaker Distributed Model Parallel Library. Følg med for at komme i gang Rediger et PyTorch-træningsscript at tilpasse SMP'ers API'er i dit træningsscript. I dette afsnit kalder vi kun nogle få hovedtrin frem med kodestykker fra træningsscriptet, der er klar til brug train_gpt_simple.py. Du kan følge kommentarerne i manuskriptet og API-dokument for at lære mere om, hvor SMP API'er bruges.

Først skal du importere og initialisere biblioteket ved at ringe smdistributed.modelparallel.torch.init() i begyndelsen af ​​træningsscriptet:

import smdistributed.modelparallel.torch as smp

smp.init(smp_config)

For det andet skal du pakke den model, der skal opdeles med smdistributed.modelparallel.torch.DistributedModel og brug det returnerede DistributedModel objekt fremadrettet:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_config(model_config)
model = smp.DistributedModel(model, trace_device="gpu", backward_passes_per_step=args.gradient_accumulation)

Pak optimizeren ind med smdistributed.modelparallel.torch.DistributedOptimizer til at gemme og indlæse optimeringstilstande.

from torch import optim

optimizer = optim.Adam(
    param_groups, betas=(args.beta1, args.beta2), lr=args.lr, weight_decay=args.weight_decay
)

optimizer = smp.DistributedOptimizer(
        optimizer, 
        static_loss_scale=None, 
        dynamic_loss_scale=True,
        dynamic_loss_args={"scale_window": 1000, "min_scale": 1, "delayed_shift": 2},
        )

Sæt frem og tilbage logikken i en trinfunktion og dekorer den med smdistributed.modelparallel.torch.step.  Enhver beregning defineret inde i smp.step-decorated funktion udføres på en distribueret måde.

@smp.step
def train_step(model, optimizer, input_ids, attention_mask, args):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    model.backward(loss)

    return loss

@smp.step
def test_step(model, input_ids, attention_mask):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    
    return loss

2. Forbered datasættet

Vi bruger openwebtext er det datasæt, vi bruger i dette eksempel. Notesbogen bruger scriptet data_prep_512.py at downloade og forbehandle datasættet. Du kan også træne med andre datasæt ved at ændre data_pipeline.py. Når du har med stort datasæt og model at gøre, kan du fremskynde træningsjobbet ved at bruge data gemt i Amazon FSx til Luster, som giver et højtydende filsystem, der er integreret med Amazon Simple Storage Service (S3). Se venligst instruktionerne fra Konfigurer datainputkanal til at bruge Amazon FSx til Luster til vejledning om indstilling af et FSx Luster-filsystem som datainputkanal.

3. Start træningsjobene

Dette trin forudsætter, at du allerede har ændret dit træningsscript og forberedte datasættet som nævnt i de foregående afsnit. Til aktivere sharded data parallelism, skal du blot indstille sharded_data_parallel_degree i PyTorch Estimator. I denne tutorial sætter vi sharded_data_parallel_degree=128 , instace_count=32 for p4d.24xlarge noder, hvilket indikerer, at modeltilstandene vil blive fordelt over 128 GPU'er blandt de i alt 256 GPU'er. Baseret på denne valgte værdi, sætter SMP automatisk dataparallelgraden til 2 (fordi 256/128=2), hvilket betyder, at vi har to replikaer til dataparallelisme. En generel regel for at vælge en ideel værdi for sharded_data_parallel_degree er at tilføje en node mere til delingsgruppen for hver 3B af modelparametre. I denne tutorial er vores modelstørrelse 30B, så vi bør bruge mindst 10 noder til sharding. Og fordi 16 noder (128 GPU'er) er den mindste power-of-2 over tærsklen, sætter vi sharded_data_parallel_degree=128.

Til checkpointing leverer vi også et sæt checkpointing-værktøjer i sharded_data_parallel_checkpoint.py , herunder et værktøj til at rekonstruere det fulde state_dict til avancerede brugssager. Endelig kan vi starte et distribueret træningsjob ved at kalde fit() på Estimatoren.

smp_estimator = PyTorch(
    entry_point="train_gpt_simple.py",
    instance_type="ml.p4d.24xlarge",
    source_dir=os.getcwd(),
    volume_size=500,
    instance_count=32,
    distribution={
        "mpi": {
            "enabled": True,
            "processes_per_host": processes_per_host,
            "custom_mpi_options": mpioptions,
        },
        "smdistributed": {
            "modelparallel": {
                "enabled": True,
                "parameters": {
                    "ddp": True,
                    "skip_tracing": True,
                    "delayed_parameter_initialization": True,
                    "offload_activations": True,
                    "activation_loading_horizon": 4,
                    # To enable sharded data parallelism.
                    # Here we shard model states across 128 GPUs. 
                    "sharded_data_parallel_degree": 128, 
                    "fp16": False,
                    "bf16": True,
                    # This is to disable pipeline parallelism.
                    "partitions": 1,
                },
            }
        },
    },
    framework_version="1.12",
    py_version="py38",
    hyperparameters=hyperparameters,
    checkpoint_s3_uri=checkpoint_s3_uri if not use_fsx else None,
    checkpoint_local_path=hyperparameters["checkpoint-dir"] if use_fsx else None,
    ...
)

smp_estimator.fit(inputs=data_channels)

4. Overvåg træningsjob

Du kan få adgang til træningsloggene og spore GPU og hukommelsesudnyttelse på amazoncloudwatch. Sørg for at se på logfilerne for "algo-1" fordi det er hovedknudepunktet, hvis outputstrøm har træningsjobloggene fra alle instanser.

Benchmarking ydeevne

Vi benchmarkerede sharded data-parallelisme i SMP-biblioteket på både 16 og 32 p4d.24xlarge noder for henholdsvis sekvenslængde 512 og 2048. 30B-parameter GPT2-modellen er konfigureret til at bruge en skjult bredde på 7168, 48 lag og 64 hoveder. Du kan anvende nøjagtig den samme konfiguration, hvor sekvenslængden er 2048 ved at indstille model_config = "gpt2-30b" i vejledningsnotesbogen. Med denne indstilling opnåede SMP 73.52 prøver i sekundet, en hastighed på 39.7 % sammenlignet med DeepSpeed ​​ZeRO-3. Hvis din tokenstørrelse er 500 milliarder, betyder denne hastighed næsten 367 timers besparelser på p4d.24xlarge noder, hvilket svarer til mere end $12,000 budget, der spares pr. træning! Følgende tabel opsummerer vores benchmarkresultater.

1/ For hver modelkonfiguration testede vi forskellige funktioner, stadier og konfigurationer i DeepSpeed ​​ZeRO og valgte den, der giver den bedste gennemstrømning som DeepSpeed-baseline. Benchmark blev kørt på Amazon Elastic Compute Cloud (Amazon EC2). 2/ Disse resultater er afhængige af forbedrede kommunikationskollektiver optimeret til AWS, som snart vil blive gjort tilgængelige. 3/ Tid til at træne er projekteret ud fra hastighed baseret på antallet af behandlede tokens.

Sammenfattende observerede vi konsekvent højere gennemløb med sønderdelt dataparallelisme i SMP sammenlignet med DeepSpeed ​​på tværs af en række modeller og konfigurationer. Denne nye funktion demonstrerede også en bedre hukommelseseffektivitet sammenlignet med DeepSpeed, hvilket gør det muligt for SMP at passe til en større batchstørrelse og reducere niveauet af gradientakkumulering, der kræves for at passe til en bestemt global batchstørrelse.

Konklusion

I dette indlæg introducerede vi en ny distribueret træningsteknik - sharded data parallelism - og hvordan den fremskynder gigantisk modeltræning med næsten lineær skalering på Amazon SageMaker. Vi gik også igennem, hvordan man træner en GPT-2-model med den nye teknik efter dette komplet eksempel. Du kan følge Amazon SageMaker Eksempler GitHub repo at spore alle SageMaker-modeller parallelle eksempler eller deltage i vores næste uddelte træningsværksteder. For at lære mere om sharded data parallelism, se venligst dokumentation.

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.

Kan Karakus er Senior Applied Scientist hos AWS, der optimerer distribueret dyb læring i stor skala på AWS. Hans forskningsinteresser dækker over deep learning, distribueret optimering, distribuerede systemer og informationsteori. Uden for arbejdet nyder han at cykle, rejse, læse og lære.

Rahul Huilgol er senior softwareingeniør hos AWS. Han arbejder på distribuerede deep learning-systemer, for at gøre det nemt og effektivt at træne store deep learning-modeller i skyen. I sin fritid nyder han fotografering, cykling og havearbejde.

Suhit Kodgule er en softwareudviklingsingeniør med AWS Artificial Intelligence-gruppe, der arbejder med deep learning-rammer. I sin fritid nyder han at vandre, rejse og lave mad.

Erin Ho er produktchef for AWS Deep Learning. Hun arbejder på produkter, der gør det nemmere for kunderne at træne deep learning-modeller på AWS. For sjov udenfor arbejde nyder hun at vandre og stå på ski.