Träna gigantiska modeller med nästan linjär skalning med hjälp av fragmenterad dataparallellism på Amazon SageMaker

I strävan efter överlägsen noggrannhet har djupinlärningsmodeller inom områden som naturlig språkbehandling och datorseende vuxit avsevärt i storlek under de senaste åren, ofta räknade i tiotals till hundratals miljarder parametrar. Att träna dessa gigantiska modeller är utmanande och kräver komplexa distributionsstrategier. Dataforskare och maskininlärningsingenjörer letar ständigt efter det bästa sättet att optimera sin träningsberäkning, men de kämpar med kommunikationskostnaderna som kan öka i takt med den totala klustrets storlek.

Det är därför vi nyligen lanserade shård dataparallellism on Amazon SageMaker, en ny minnesbesparande distribuerad träningsteknik i SageMaker modell parallell (SMP) bibliotek. Delad dataparallellism är specialbyggd för extremskaliga modeller och använder Amazon internt MiCS teknologi under huven, ett vetenskapligt försök att minimera kommunikationsskalan genom att få ner dyra kommunikationskostnader med rötter i parameterinsamling och gradientsynkronisering. Med en 30B parameter GPT-2-modell med sekvenslängd 2048, uppnådde denna nya funktion 141 TFLOPs, en hastighet på 39.7 % jämfört med DeepSpeed ​​ZeRO-3. För en 10B GPT-2-modell med sekvenslängd 512 uppnådde denna nya funktion också 564 sampel per sekund, en hastighet på 13.9 % jämfört med PyTorchs Fully Sharded Data Parallel (FSDP). Kom ihåg att i gigantisk modellträning översätts varje procent av speedupen till sparade dollar och ökad produktivitet i ditt team.

I det här blogginlägget kommer vi först att titta närmare på de viktigaste skillnaderna mellan fragmenterad dataparallellism och när den ska användas. Sedan kommer du att lära dig hur du enkelt tränar en 30B parameter GPT-2-modell på SageMaker med denna nya funktion. Slutligen kommer vi att jämföra prestandan med andra alternativ med öppen källkod, särskilt överträffa DeepSpeed ​​ZeRO med upp till 39.7 % på 256 GPU:er.

Hur delad dataparallellism fungerar och när den ska användas

Innan vi introducerar fragmenterad dataparallellism, låt oss titta på dess bredare teknikfamilj. Nyligen distribuerade träningsmetoder för stora modeller har flyttats till ett paradigm där modellparametrar, gradienter och optimerartillstånd delas över dataparallella noder. Till skillnad från Pipeline Parallelism som har den medfödda komplexiteten att välja lager för att partitionera över enheter, särskilt när ditt ramverk inte stöder automatisk modelldelning, detta paradigm bevarar elegant enkelheten hos dataparallellism, samtidigt som den tar bort dataparallellismens begränsning där en modell måste passa in i en enda GPU.

I befintliga ramverk som faller under detta paradigm, särskilt DeepSpeed ​​ZeRO-3 och PyTorchs FSDP uppströms från FairScale, är modelltillstånd splittrade över alla GPU:er, en strategi som sänker minnesförbrukningen på varje GPU till priset av stora kommunikationskostnader som ökar med klusterstorleken och därför gör att skalbarheten minskar avsevärt i skalan. I motsats härtill tillstånd delad dataparallellism i SMP-bibliotekspartitionsmodellen i en skala medveten sätt genom att partitionera varje replik av modelltillstånd endast inom en delmängd av GPU:er.

Låt oss titta närmare på skala-medveten modellpartitionering i MiCS, kärntekniken bakom fragmenterad data parallell. Intuitionen bakom denna design är att partitionering av träningstillstånd över hela den dataparallella gruppen kanske inte krävs för att träna en modell med tiotals miljarder parametrar. Till exempel räcker 8 V100 GPU:er (32GB vardera) för att hålla modelltillståndsrepliken av en 10B-parametermodell som behöver cirka 200GB minne när man tränar med Adam optimizer med blandad precision. Genom att begränsa en komplett kopia av modelltillstånd i minsta delmängd av GPU:er kan vi effektivt minska omfattningen av kommunikationsoverhead jämfört med DeepSpeed ​​och PyTorch FSDP. Delad data parallell utnyttjar också andra tekniker i MiCS som hierarkisk kommunikation och 2-hop gradientsynkronisering. För mer information, kolla in Nära linjär skalning av gigantisk modellträning på AWS or MiCS: Närlinjär skalning för träning av gigantisk modell på Public Cloud.

Nu, hur vet du när du ska välja fragmenterad data parallellt framför andra distribuerade träningstekniker? Den allmänna regeln är att om din modell har mindre än 1 miljard parametrar och kan passa in i GPU-minnet, SageMaker data parallellbibliotek or SageMaker träningskompilator kan vara tillräckligt för dig. Om du har större språk- eller datorseendemodeller är vårt förslag att träna upp det med tekniken för fragmenterad dataparallellism i kombination med aktiveringskontroll och aktivering avlastning i SageMakers modell parallellbibliotek först, innan andra tekniker som t.ex tensorparallellism eller pipelineparallellism.

Använder delad dataparallellism för att träna GPT-2 på Amazon SageMaker

Låt oss nu lära oss hur man tränar en GPT-2-modell med delad data parallellt, med SMP som kapslar in komplexiteten åt dig. Detta komplett handledningsanteckningsbok leder dig genom hela processen, från databearbetning, definition och inlämning av utbildningsjobb, till övervakning av träningsloggar. Vad som följer är en kort översikt som belyser viktiga steg för att använda den här funktionen.

1. Komma igång

Delad dataparallellism är tillgänglig i PyTorch v1.12.0+ och fungerar med både FP16 och BF16. Det enklaste sättet att använda SMP-biblioteket är genom en förbyggd AWS Deep Learning Container för PyTorch. Men om du vill ta med din egen Docker-container kan du hänvisa till Skapa din egen Docker-behållare med SageMaker Distributed Model Parallel Library. För att komma igång, följ Ändra ett PyTorch-träningsskript för att anpassa SMPs API:er i ditt träningsskript. I det här avsnittet ropar vi bara ut några få huvudsteg med kodavsnitt från det färdiga träningsskriptet train_gpt_simple.py. Du kan följa kommentarerna i manuset och API-dokument för att lära dig mer om var SMP API:er används.

Importera och initiera först biblioteket genom att ringa smdistributed.modelparallel.torch.init() i början av träningsmanuset:

import smdistributed.modelparallel.torch as smp

smp.init(smp_config)

För det andra, slå in modellen som ska partitioneras med smdistributed.modelparallel.torch.DistributedModel och använd det returnerade DistributedModel objekt framöver:

from transformers import AutoModelForCausalLM

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

Slå in optimeraren med smdistributed.modelparallel.torch.DistributedOptimizer för att spara och ladda optimeringslägen.

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

Lägg framåt och bakåt logiken i en stegfunktion och dekorera den med smdistributed.modelparallel.torch.step.  Alla beräkningar definierade i smp.step-decorated funktionen utförs på ett distribuerat sätt.

@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. Förbered datauppsättningen

Vi använder openwebtext är den datauppsättning vi använder i det här exemplet. Anteckningsboken använder skriptet data_prep_512.py för att ladda ner och förbehandla datamängden. Du kan också träna med andra datauppsättningar genom att modifiera data_pipeline.py. När du har att göra med stor datamängd och modell kan du påskynda träningsjobbet genom att använda data lagrad i Amazon FSx för Luster, som ger ett högpresterande filsystem integrerat med Amazon enkel lagringstjänst (S3). Se instruktionerna från Konfigurera datainmatningskanal för att använda Amazon FSx för Luster för vägledning om att ställa in ett FSx Luster-filsystem som datainmatningskanal.

3. Starta träningsjobben

Detta steg förutsätter att du redan har gjort det ändrat ditt träningsskript och förberedde datamängden som nämnts i de föregående avsnitten. Till aktivera fragmenterad dataparallellism, ställ helt enkelt in sharded_data_parallel_degree i PyTorch Estimator. I den här handledningen ställer vi in sharded_data_parallel_degree=128 och instace_count=32 för p4d.24xlarge noder, vilket indikerar att modelltillstånden kommer att delas över 128 GPU:er bland de totalt 256 GPU:erna. Baserat på detta valda värde ställer SMP sedan automatiskt in dataparallellgraden till 2 (eftersom 256/128=2), vilket betyder att vi har två repliker för dataparallellism. En allmän regel för att välja ett idealiskt värde för sharded_data_parallel_degree är att lägga till ytterligare en nod till delningsgruppen för varje 3B av modellparametrar. I den här handledningen är vår modellstorlek 30B, så vi bör använda minst 10 noder för skärning. Och eftersom 16 noder (128 GPU:er) är den minsta power-of-2 över tröskeln, ställer vi in sharded_data_parallel_degree=128.

För checkpointing tillhandahåller vi också en uppsättning checkpointingverktyg i sharded_data_parallel_checkpoint.py , inklusive ett verktyg för att rekonstruera hela state_dict för avancerade användningsfall. Slutligen kan vi lansera ett distribuerat träningsjobb genom att anropa fit() på Estimatorn.

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. Övervaka träningsjobben

Du kan komma åt träningsloggarna och spåra GPU och minnesanvändning på amazoncloudwatch. Se till att titta på loggarna för "algo-1" eftersom det är huvudnoden vars utgångsström har träningsjobbloggarna från alla instanser.

Benchmarking prestanda

Vi benchmarkerade fragmenterad dataparallellism i SMP-biblioteket på både 16 och 32 p4d.24xlarge noder för sekvenslängd 512 respektive 2048. GPT30-modellen med 2B-parameter är konfigurerad att använda en dold bredd på 7168, 48 lager och 64 huvuden. Du kan använda exakt samma konfiguration där sekvenslängden är 2048 genom att ställa in model_config = "gpt2-30b" i handledningsanteckningsboken. Med den här inställningen uppnådde SMP 73.52 sampel per sekund, en hastighet på 39.7 % jämfört med DeepSpeed ​​ZeRO-3. Om din tokenstorlek är 500 miljarder betyder denna snabbhet nästan 367 timmars besparingar på p4d.24xlarge noder, vilket motsvarar mer än $12,000 XNUMX budget sparad per träning! Följande tabell sammanfattar våra benchmarkresultat.

1/ För varje modellkonfiguration testade vi olika funktioner, stadier och konfigurationer i DeepSpeed ​​ZeRO och valde den som ger bäst genomströmning som DeepSpeeds baslinje. Riktmärket kördes på Amazon Elastic Compute Cloud (Amazon EC2). 2/ Dessa resultat förlitar sig på förbättrade kommunikationskollektiv optimerade för AWS som kommer att göras tillgängliga snart. 3/ Tid att träna projiceras från hastighet baserat på antalet bearbetade tokens.

Sammanfattningsvis observerade vi konsekvent högre genomströmning med fragmenterad dataparallellism i SMP jämfört med DeepSpeed ​​över en rad modeller och konfigurationer. Denna nya funktion visade också en bättre minneseffektivitet jämfört med DeepSpeed, vilket gör att SMP kan passa en större batchstorlek och minska nivån av gradientackumulering som krävs för att passa en viss global batchstorlek.

Slutsats

I det här inlägget introducerade vi en ny distribuerad träningsteknik – fragmenterad dataparallellism – och hur den påskyndar gigantisk modellträning med nära linjär skalning på Amazon SageMaker. Vi gick också igenom hur man tränar en GPT-2-modell med den nya tekniken efter detta fullständigt exempel. Du kan följa Amazon SageMaker-exempel GitHub-repo för att spåra alla SageMaker modell parallella exempel eller delta i vår nästa distribuerade utbildningsverkstäder. För att lära dig mer om fragmenterad dataparallellism, se dokumentation.

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.

Kan Karakus är Senior Applied Scientist på AWS, som optimerar storskalig distribuerad djupinlärning på AWS. Hans forskningsintressen omfattar djupinlärning, distribuerad optimering, distribuerade system och informationsteori. Utanför jobbet tycker han om att cykla, resa, läsa och lära sig.

Rahul Huilgol är senior mjukvaruingenjör på AWS. Han arbetar med distribuerade djupinlärningssystem, för att göra det enkelt och prestanda att träna stora djupinlärningsmodeller i molnet. På fritiden gillar han fotografering, cykling och trädgårdsarbete.

Suhit Kodgule är en mjukvaruutvecklingsingenjör med AWS Artificiell Intelligens-grupp som arbetar med ramar för djupinlärning. På fritiden tycker han om att vandra, resa och laga mat.

Erin Ho är produktchef för AWS Deep Learning. Hon arbetar med produkter som gör det enklare för kunder att träna djupinlärningsmodeller på AWS. För skoj utanför jobbet tycker hon om att vandra och åka skidor.