Tren gigantiske modeller med nesten lineær skalering ved å bruke sønderdelt dataparallellisme på Amazon SageMaker

I jakten på overlegen nøyaktighet har dyplæringsmodeller innen områder som naturlig språkbehandling og datasyn vokst betydelig i størrelse de siste årene, ofte talt i titalls til hundrevis av milliarder av parametere. Trening av disse gigantiske modellene er utfordrende og krever komplekse distribusjonsstrategier. Dataforskere og maskinlæringsingeniører leter stadig etter den beste måten å optimalisere treningsberegningen på, men de sliter med kommunikasjonskostnadene som kan øke sammen med den totale klyngestørrelsen.

Dette er grunnen til at vi nylig lanserte shardt dataparallellisme on Amazon SageMaker, en ny minnebesparende distribuert treningsteknikk i SageMaker modell parallell (SMP) bibliotek. Delt dataparallellisme er spesialbygd for ekstremskalamodeller og bruker Amazon internt MiCS teknologi under panseret, en vitenskapelig innsats for å minimere kommunikasjonsskalaen ved å redusere kostbare kommunikasjonskostnader forankret i parameterinnsamling og gradientsynkronisering. Med en 30B parameter GPT-2-modell med sekvenslengde 2048, oppnådde denne nye funksjonen 141 TFLOP-er, en hastighet på 39.7 % sammenlignet med DeepSpeed ​​ZeRO-3. For en 10B GPT-2-modell med sekvenslengde 512, oppnådde denne nye funksjonen også 564 prøver per sekund, en hastighet på 13.9 % sammenlignet med PyTorchs Fully Sharded Data Parallel (FSDP). Husk at i gigantisk modelltrening vil hver prosentandel av speedup oversettes til penger spart og produktivitet oppnådd i teamet ditt.

I dette blogginnlegget skal vi først se nærmere på de viktigste differensiatorene for deling av dataparallellisme og når den skal brukes. Deretter lærer du hvordan du enkelt trener en 30B parameter GPT-2-modell på SageMaker med denne nye funksjonen. Til slutt vil vi sammenligne ytelsen med andre åpen kildekode-alternativer, spesielt bedre enn DeepSpeed ​​ZeRO med opptil 39.7 % på 256 GPUer.

Hvordan delt dataparallellisme fungerer og når den skal brukes

Før vi introduserer delt dataparallellisme, la oss se på den bredere teknikkfamilien. Nylige distribuerte treningstilnærminger for store modeller har flyttet til et paradigme der modellparametere, gradienter og optimeringstilstander deles på tvers av dataparallelle noder. I motsetning til Pipeline Parallelism som har den medfødte kompleksiteten ved å velge lag for å partisjonere på tvers av enheter, spesielt når rammeverket ditt ikke støtter automatisert modelloppdeling, bevarer dette paradigmet elegant enkelheten til dataparallellisme, samtidig som den fjerner dataparallellismens begrensning der en modell må passe inn i en enkelt GPU.

I eksisterende rammeverk som faller inn under dette paradigmet, spesielt DeepSpeed ​​ZeRO-3 og PyTorchs FSDP oppstrøms fra FairScale, er modelltilstander spredt over alle GPUer, en strategi som senker minneforbruket på hver GPU på bekostning av store kommunikasjonskostnader som øker med klyngestørrelsen og derfor får skalerbarheten til å synke betydelig i skala. I motsetning til dette, er delt dataparallellisme i SMP-bibliotekpartisjonsmodellen i en skalabevisst måte ved å partisjonere hver replika av modelltilstander bare innenfor en delmengde av GPUer.

La oss se nærmere på skalabevisste modellpartisjonering i MiCS, kjerneteknologien bak sharded data parallel. Intuisjonen bak dette designet er at partisjoneringstreningstilstander på tvers av hele den dataparallelle gruppen kanskje ikke er nødvendig for å trene en modell med titalls milliarder av parametere. For eksempel er 8 V100 GPUer (32 GB hver) tilstrekkelig til å holde modellstatene replika av en 10B-parameter modell som trenger omtrent 200 GB minne når du trener med Adam optimizer ved bruk av blandet presisjon. Ved å begrense en fullstendig kopi av modelltilstander i minste delsett av GPUer, kan vi effektivt redusere omfanget av kommunikasjonskostnader sammenlignet med DeepSpeed ​​og PyTorch FSDP. Delt dataparallell utnytter også andre teknikker i MiCS som hierarkisk kommunikasjon og 2-hop Gradient Synchronization. For mer informasjon, sjekk ut Nær-lineær skalering av gigantisk modelltrening på AWS or MiCS: Nær-lineær skalering for opplæring av gigantisk modell på Public Cloud.

Nå, hvordan vet du når du skal velge sharded data parallelt fremfor andre distribuerte treningsteknikker? Den generelle regelen er at hvis modellen din har mindre enn 1 milliard parametere og kan passe inn i GPU-minne, SageMaker data parallellbibliotek or SageMaker treningskompilator kan være tilstrekkelig for deg. Hvis du har større språk- eller datasynsmodeller, er vårt forslag å trene det med sharded data parallelism-teknikken kombinert med aktiveringssjekkpunkt og aktivering avlasting i SageMaker modell parallellbibliotek først, før andre teknikker som f.eks tensorparallellisme eller rørledningsparallellisme.

Bruk av sønderdelt dataparallellisme for å trene GPT-2 på Amazon SageMaker

La oss nå lære hvordan du trener en GPT-2-modell med delte data parallelt, med SMP som innkapsler kompleksiteten for deg. Dette komplett opplæringsnotatbok leder deg gjennom hele prosessen, fra databehandling, definering og innsending av opplæringsjobber, til overvåking av treningslogger. Det følgende er en kort oversikt som fremhever viktige trinn for bruk av denne funksjonen.

1. Kom i gang

Delt dataparallellisme er tilgjengelig i PyTorch v1.12.0+ og fungerer med både FP16 og BF16. Den enkleste måten å bruke SMP-biblioteket på er gjennom en forhåndsbygd AWS Deep Learning Container for PyTorch. Men hvis du ønsker å ta med din egen Docker-beholder, kan du henvise til Lag din egen Docker-beholder med SageMaker Distributed Model Parallel Library. For å komme i gang, følg Endre et PyTorch-treningsskript for å tilpasse SMPs' APIer i treningsskriptet ditt. I denne delen roper vi bare ut noen få hovedtrinn med kodebiter fra opplæringsskriptet som er klart til bruk train_gpt_simple.py. Du kan følge kommentarene i manuset og API-dokument for å lære mer om hvor SMP APIer brukes.

Først importerer og initialiser biblioteket ved å ringe smdistributed.modelparallel.torch.init() i begynnelsen av treningsmanuset:

import smdistributed.modelparallel.torch as smp

smp.init(smp_config)

For det andre, pakk inn modellen som skal partisjoneres med smdistributed.modelparallel.torch.DistributedModel og bruk det returnerte DistributedModel objekt fremover:

from transformers import AutoModelForCausalLM

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

Pakk inn optimizeren med smdistributed.modelparallel.torch.DistributedOptimizer for lagring og lasting av optimaliseringstilstander.

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},
        )

Sett forover- og bakoverlogikken i en trinnfunksjon og dekorer den med smdistributed.modelparallel.torch.step.  Enhver beregning definert i smp.step-decorated funksjonen utføres på en distribuert måte.

@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. Klargjør datasettet

Vi bruker openwebtext er datasettet vi bruker i dette eksemplet. Notatboken bruker skriptet data_prep_512.py for å laste ned og forhåndsbehandle datasettet. Du kan også trene med andre datasett ved å modifisere data_pipeline.py. Når du har å gjøre med store datasett og modeller, kan du få fart på treningsjobben ved å bruke data som er lagret i Amazon FSx for Luster, som gir et høyytelses filsystem som er integrert med Amazon enkel lagringstjeneste (S3). Vennligst se instruksjonene fra Konfigurer datainndatakanal for å bruke Amazon FSx for Luster for veiledning om innstilling av et FSx Luster-filsystem som datainndatakanal.

3. Start treningsjobbene

Dette trinnet forutsetter at du allerede har gjort det endret treningsskriptet ditt og utarbeidet datasettet som nevnt i de foregående avsnittene. Til aktiver delt dataparallellisme, bare still inn sharded_data_parallel_degree i PyTorch Estimator. I denne opplæringen setter vi sharded_data_parallel_degree=128 og instace_count=32 for p4d.24xlarge noder, noe som indikerer at modelltilstandene vil bli delt over 128 GPUer blant de totalt 256 GPUene. Basert på denne valgte verdien, vil SMP da automatisk sette dataparallellgraden til 2 (fordi 256/128=2), noe som betyr at vi har to replikaer for dataparallellisme. En generell regel for å velge en ideell verdi for sharded_data_parallel_degree er å legge til en node til i delingsgruppen for hver 3B med modellparametere. I denne opplæringen er modellstørrelsen vår 30B, så vi bør bruke minst 10 noder for skjæring. Og fordi 16 noder (128 GPUer) er den minste kraften på 2 over terskelen, satte vi sharded_data_parallel_degree=128.

For sjekkpunkter tilbyr vi også et sett med sjekkpunktverktøy sharded_data_parallel_checkpoint.py , inkludert et verktøy for å rekonstruere hele state_dict for avansert bruk. Til slutt kan vi lansere en distribuert treningsjobb ved å ringe fit() på Estimator.

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åk treningsjobbene

Du kan få tilgang til treningsloggene og spore GPU og minneutnyttelse på Amazon CloudWatch. Sørg for å se på loggene til "algo-1" fordi det er hovednoden hvis utdatastrøm har treningsjobbloggene fra alle forekomster.

Benchmarking ytelse

Vi benchmarkerte delt dataparallellisme i SMP-biblioteket på både 16 og 32 p4d.24xlarge noder for henholdsvis sekvenslengde 512 og 2048. 30B-parameter GPT2-modellen er konfigurert til å bruke en skjult bredde på 7168, 48 lag og 64 hoder. Du kan bruke nøyaktig samme konfigurasjon der sekvenslengden er 2048 ved å stille inn model_config = "gpt2-30b" i opplæringsnotatboken. Med denne innstillingen oppnådde SMP 73.52 prøver per sekund, en hastighet på 39.7 % sammenlignet med DeepSpeed ​​ZeRO-3. Hvis tokenstørrelsen din er 500 milliarder, betyr denne hastigheten nesten 367 timer med besparelser på p4d.24xlarge noder, tilsvarende mer enn $12,000 XNUMX budsjett spart per trening! Tabellen nedenfor oppsummerer referanseresultatene våre.

1/ For hver modellkonfigurasjon testet vi forskjellige funksjoner, stadier og konfigurasjoner i DeepSpeed ​​ZeRO og valgte den som gir best gjennomstrømning som DeepSpeed-grunnlinje. Benchmark ble kjørt på Amazon Elastic Compute Cloud (Amazon EC2). 2/ Disse resultatene er avhengige av forbedrede kommunikasjonskollektiver optimalisert for AWS som vil bli gjort tilgjengelig snart. 3/ Tid til å trene projiseres fra hastighet basert på antall behandlede tokens.

Oppsummert observerte vi konsekvent høyere gjennomstrømning med sønderdelt dataparallellisme i SMP sammenlignet med DeepSpeed ​​på tvers av en rekke modeller og konfigurasjoner. Denne nye funksjonen demonstrerte også en bedre minneeffektivitet sammenlignet med DeepSpeed, noe som gjorde det mulig for SMP å passe til en større batchstørrelse og redusere nivået av gradientakkumulering som kreves for å passe til en bestemt global batchstørrelse.

konklusjonen

I dette innlegget introduserte vi en ny distribuert treningsteknikk – sønderdelt dataparallellisme – og hvordan den fremskynder gigantisk modelltrening med nesten lineær skalering på Amazon SageMaker. Vi gikk også gjennom hvordan man trener en GPT-2-modell med den nye teknikken etter dette komplett eksempel. Du kan følge Amazon SageMaker Eksempler GitHub repo for å spore alle SageMaker modell parallelle eksempler eller delta på vårt neste delte ut opplæringsverksteder. For å lære mer om deling av dataparallellisme, se dokumentasjon.

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.

Kan Karakus er Senior Applied Scientist ved AWS, og optimaliserer distribuert dyp læring i stor skala på AWS. Hans forskningsinteresser dekker dyp læring, distribuert optimalisering, distribuerte systemer og informasjonsteori. Utenom jobben liker han å sykle, reise, lese og lære.

Rahul Huilgol er senior programvareingeniør hos AWS. Han jobber med distribuerte dyplæringssystemer, for å gjøre det enkelt og effektivt å trene opp store dyplæringsmodeller i skyen. På fritiden liker han fotografering, sykling og hagearbeid.

Suhit Kodgule er en programvareutviklingsingeniør med AWS Artificial Intelligence-gruppe som jobber med rammeverk for dyp læring. På fritiden liker han å gå tur, reise og lage mat.

Erin Ho er produktsjef for AWS Deep Learning. Hun jobber med produkter som gjør det enklere for kunder å trene dyplæringsmodeller på AWS. For moro utenfor jobben liker hun å gå tur og gå på ski.