Tanítson gigantikus modelleket közel lineáris skálázással az Amazon SageMaker szaggatott adatpárhuzamának segítségével

A kiváló pontosságra törekedve a mélytanulási modellek olyan területeken, mint a természetes nyelvi feldolgozás és a számítógépes látás, jelentősen megnőttek az elmúlt néhány évben, gyakran tíz-százmilliárd paraméterben. E gigantikus modellek betanítása kihívást jelent, és összetett terjesztési stratégiákat igényel. Az adattudósok és a gépi tanulási mérnökök folyamatosan keresik a legjobb módszert képzési számításaik optimalizálására, de küszködnek a kommunikációs többletterheléssel, amely a klaszter teljes méretével együtt növekedhet.

Ezért indítottuk el a közelmúltban skemény adatpárhuzam on Amazon SageMaker, egy új memóriakímélő elosztott edzéstechnika a SageMaker modell párhuzamos (SMP) könyvtár. A megosztott adatok párhuzamosságát kifejezetten az extrém léptékű modellekhez tervezték, és házon belül használja az Amazont MiCS technológia a motorháztető alatt, tudományos erőfeszítés a kommunikációs lépték minimalizálására a paramétergyűjtésben és a gradiens szinkronizálásban gyökerező költséges kommunikációs költségek csökkentésével. A 30B paraméterű, 2 szekvenciahosszú GPT-2048 modellel ez az új funkció 141 TFLOP-t ért el, ami 39.7%-os sebesség a DeepSpeed ​​ZeRO-3-hoz képest. Az 10 szekvenciahosszúságú 2B GPT-512 modellnél ez az új funkció 564 mintát is elért másodpercenként, ami 13.9%-os sebesség a PyTorch Fully Sharded Data Parallel (FSDP) funkciójához képest. Ne feledje, hogy a gigantikus modellképzés során a gyorsulás minden százaléka megtakarított dollárt és megnövelt termelékenységet jelent csapatában.

Ebben a blogbejegyzésben először közelebbről megvizsgáljuk a szilánkos adatok párhuzamosságának fő megkülönböztető tényezőit, és azt, hogy mikor érdemes használni. Ezután megtanulhatja, hogyan taníthat meg könnyedén egy 30B paraméterű GPT-2 modellt a SageMakeren ezzel az új funkcióval. Végül összehasonlítjuk a teljesítményt más nyílt forráskódú opciókkal, amelyek a DeepSpeed ​​ZeRO teljesítményét akár 39.7%-kal is felülmúlják 256 GPU-n.

Hogyan működik a szilánkos adatok párhuzamossága, és mikor kell használni

Mielőtt bemutatnánk a szilánkos adatok párhuzamosságát, nézzük meg a szélesebb technikacsaládját. A nagy modellek legújabb elosztott betanítási megközelítései egy olyan paradigmába költöztek, ahol a modellparaméterek, a gradiensek és az optimalizáló állapotai az adatokkal párhuzamos csomópontok között vannak megosztva. Ellentétben a Pipeline Parallelism-mel, amelynek veleszületett összetettsége az eszközök közötti particionáláshoz szükséges rétegek kiválasztása, különösen akkor, ha a keretrendszer nem támogatja automatizált modellfelosztás, ez a paradigma elegánsan megőrzi az adatok párhuzamosságának egyszerűségét, miközben megszünteti az adatpárhuzamosság azon korlátját, hogy a modellnek egyetlen GPU-ba kell illeszkednie.

Az e paradigma alá tartozó meglévő keretrendszerekben, nevezetesen a DeepSpeed ​​ZeRO-3-ban és a PyTorch FSDP-ben, amely a FairScale-től felfelé áramlik, a modellállapotok szét vannak osztva. minden GPU-k, egy olyan stratégia, amely csökkenti az egyes GPU-k memóriafelhasználását nagy kommunikációs többletterheléssel, amely a fürt méretével növekszik, és ezért a méretezhetőség jelentősen csökken a skálán. Ezzel szemben a szilánkos adatok párhuzamossága az SMP könyvtárpartíciós modellben az a léptéktudatos úgy, hogy a modellállapotok minden replikáját csak belül particionálja egy részhalmaz a GPU-k közül.

Nézzük meg közelebbről a méretarányos modellparticionálás a MiCS-ben, a párhuzamos adatok mögött meghúzódó alapvető technológia. A terv mögött meghúzódó megérzés az, hogy a tanítási állapotok particionálása a teljes adat-párhuzamos csoportban előfordulhat, hogy nem szükséges egy több tízmilliárd paraméterrel rendelkező modell betanításához. Például 8 V100 GPU (mindegyik 32 GB) elegendő egy 10B-s paraméterű modell modellállapotának replikájának megtartásához, amely körülbelül 200 GB memóriát igényel, ha vegyes pontosságú Adam optimalizálóval edz. A modellállapotok teljes replikájának korlátozásával a legkisebb A GPU-k részhalmaza hatékonyan csökkentheti a kommunikációs többletterhelést a DeepSpeedhez és a PyTorch FSDP-hez képest. A megosztott adatok párhuzamossága más technikákat is kihasznál a MiCS-ben, mint például a hierarchikus kommunikáció és a 2-hop gradiens szinkronizálás. További információért nézze meg A gigantikus modell-képzés közel lineáris skálázása AWS-en or MiCS: Közel-lineáris skálázás a gigantikus modell betanításához nyilvános felhőn.

Most honnan tudja, hogy mikor válassza a szilánkos adatokat párhuzamosan más elosztott képzési technikákkal szemben? Az általános szabály az, hogy ha a modell 1 milliárdnál kevesebb paraméterrel rendelkezik, és elfér a GPU memóriájában, SageMaker adatok párhuzamos könyvtár or SageMaker képzési fordító elég lehet neked. Ha nagyobb nyelvi vagy számítógépes látásmodelljei vannak, azt javasoljuk, hogy ezt a szilánkos adatok párhuzamosítási technikájával, kombinálva aktiválás ellenőrzőpont és a aktiválás tehermentesítése először a SageMaker modell párhuzamos könyvtárában, mielőtt más technikákat, például tenzor párhuzamosság vagy csővezeték párhuzamosság.

A szilánkos adatok párhuzamosságának használata a GPT-2 betanítására az Amazon SageMakeren

Most pedig tanuljuk meg, hogyan lehet párhuzamosan betanítani egy GPT-2-modellt szilánkos adatokkal, és az SMP beágyazza a bonyolultságot. Ez teljes oktatófüzet végigvezeti a teljes folyamaton, az adatfeldolgozástól a képzési munkák meghatározásán és benyújtásán át a képzési naplók figyeléséig. Az alábbiakban egy rövid áttekintés található, amely kiemeli a funkció használatának legfontosabb lépéseit.

1. Fogj neki

A megosztott adatok párhuzamossága a PyTorch v1.12.0+ verzióban érhető el, és az FP16-tal és a BF16-tal is működik. Az SMP-könyvtár használatának legegyszerűbb módja a PyTorch számára előre beépített AWS Deep Learning Container. Ha azonban saját Docker-konténerét szeretné vinni, akkor hivatkozhat erre Hozzon létre saját Docker-tárolót a SageMaker Distributed Model Parallel Library segítségével. A kezdéshez kövesse PyTorch képzési parancsfájl módosítása az SMP-k API-jának adaptálásához a képzési szkriptben. Ebben a részben csak néhány fő lépést hívunk meg a használatra kész képzési szkript kódrészleteivel train_gpt_simple.py. Követheti a megjegyzéseket a forgatókönyvben és API dokumentum hogy többet megtudjon az SMP API-k használatáról.

Először is importálja és inicializálja a könyvtárat hívással smdistributed.modelparallel.torch.init() a képzési szkript elején:

import smdistributed.modelparallel.torch as smp

smp.init(smp_config)

Másodszor csomagolja be a particionálni kívánt modellt smdistributed.modelparallel.torch.DistributedModel és használja a visszaküldöttet DistributedModel tárgy a továbbiakban:

from transformers import AutoModelForCausalLM

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

Tekerje be az optimalizálót ezzel smdistributed.modelparallel.torch.DistributedOptimizer optimalizáló állapotok mentéséhez és betöltéséhez.

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

Tedd lépésfüggvénybe az előre és hátra logikát, és díszítsd vele smdistributed.modelparallel.torch.step.  Bármely számítás, amelyet a smp.step-decorated funkció elosztott módon kerül végrehajtásra.

@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. Készítse elő az adatkészletet

Az általunk használt nyílt webszöveg az az adatkészlet, amelyet ebben a példában használunk. A notebook a szkriptet használja data_prep_512.py az adatkészlet letöltéséhez és előfeldolgozásához. Más adatkészletekkel is edzhet módosítással data_pipeline.py. Ha nagy adatkészlettel és modellel foglalkozik, felgyorsíthatja a betanítási munkát a tárolt adatok használatával Amazon FSx Lusterhez, amely natívan integrált nagy teljesítményű fájlrendszert biztosít Amazon egyszerű tárolási szolgáltatás (S3). Kérjük, olvassa el az utasításokat Állítsa be az adatbeviteli csatornát az Amazon FSx for Luster használatához útmutatásért az FSx Luster fájlrendszer adatbeviteli csatornaként való beállításához.

3. Indítsa el a képzési munkákat

Ez a lépés feltételezi, hogy már megtette módosította az edzési szkriptet és elkészítette az adatkészletet az előző szakaszokban leírtak szerint. Nak nek engedélyezze a szilánkos adatok párhuzamosságát, egyszerűen állítsa be a sharded_data_parallel_degree a PyTorch Becslő. Ebben az oktatóanyagban beállítjuk sharded_data_parallel_degree=128 és a instace_count=32 p4d.24xlarge csomópontokhoz, ami azt jelzi, hogy a modellállapotok 128 GPU között lesznek felosztva az összesen 256 GPU között. A kiválasztott érték alapján az SMP ezután automatikusan 2-re állítja az adatok párhuzamos fokát (mert 256/128=2), ami azt jelenti, hogy két másolatunk lesz az adatpárhuzamosságról. Általános szabály az ideális érték kiválasztásához sharded_data_parallel_degree Az a lényeg, hogy minden 3B modellparaméterenként egy további csomópontot adjunk a megosztási csoporthoz. Ebben az oktatóanyagban a modellünk mérete 30B, tehát legalább 10 csomópontot kell használnunk a felosztáshoz. És mivel a 16 csomópont (128 GPU) a küszöb feletti 2 legkisebb teljesítménye, beállítjuk sharded_data_parallel_degree=128.

Az ellenőrzőpontokhoz egy ellenőrzőpont-kezelő segédprogramot is biztosítunk sharded_data_parallel_checkpoint.py , beleértve a teljes rekonstrukciós segédprogramot state_dict haladó felhasználási esetekre. Végül elindíthatunk egy elosztott képzési feladatot a fit() meghívásával a Becslőben.

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. Figyelje a képzési munkákat

Hozzáférhet az edzésnaplókhoz, és nyomon követheti a GPU és a memória kihasználtságát amazonfelhőóra. Ügyeljen arra, hogy nézze meg az „algo-1” naplóit mert ez az a fő csomópont, amelynek kimeneti adatfolyama tartalmazza a betanítási munkanaplókat az összes példányból.

Benchmarking teljesítmény

Összehasonlítottuk a szilánkos adatok párhuzamosságát az SMP könyvtárban mind a 16, mind a 32 p4d.24xlarge csomóponton 512 és 2048 szekvenciahosszúság esetén. A 30B paraméterű GPT2 modell úgy van beállítva, hogy 7168 rejtett szélességet, 48 réteget és 64 fejet használjon. Pontosan ugyanazt a konfigurációt alkalmazhatja, ahol a sorozat hossza 2048 beállítással model_config = "gpt2-30b" az oktatófüzetben. Ezzel a beállítással az SMP 73.52 mintát ért el másodpercenként, ami 39.7%-os sebesség a DeepSpeed ​​ZeRO-3-hoz képest. Ha a token mérete 500 milliárd, ez a gyorsulás közel 367 órás megtakarítást jelent a p4d.24xlarge csomópontokon, ami több mint 12,000 XNUMX dolláros költségmegtakarításnak felel meg képzésenként! Az alábbi táblázat összefoglalja benchmark eredményeinket.

1/ Minden modellkonfigurációhoz különböző funkciókat, fokozatokat és konfigurációkat teszteltünk a DeepSpeed ​​ZeRO-ban, és a legjobb átvitelt biztosító DeepSpeed ​​alapvonalat választottuk. A benchmark futott Amazon rugalmas számítási felhő (Amazon EC2). 2/ Ezek az eredmények az AWS-re optimalizált továbbfejlesztett kommunikációs csoportokon alapulnak, amelyek hamarosan elérhetőek lesznek. 3/ Az edzésig eltelt időt a sebesség alapján vetítjük a feldolgozott tokenek száma alapján.

Összefoglalva, következetesen nagyobb átviteli sebességet figyeltünk meg a szilánkos adatok párhuzamosságával az SMP-ben, mint a DeepSpeed-hez képest, számos modellben és konfigurációban. Ez az új funkció a DeepSpeed-hez képest jobb memóriahatékonyságot is mutatott, lehetővé téve az SMP számára, hogy nagyobb kötegmérethez illeszkedjen, és csökkentse az egy adott globális kötegmérethez szükséges gradiens felhalmozódás szintjét.

Következtetés

Ebben a bejegyzésben bemutattunk egy új elosztott képzési technikát – a szilánkos adatok párhuzamosságát – és azt, hogy ez hogyan gyorsítja fel a gigantikus modellképzést az Amazon SageMaker közel lineáris skálázásával. Azt is végigjártuk, hogyan taníthatunk GPT-2 modellt az új technikával ezt követően teljes példa. Követheti a Amazon SageMaker példák GitHub repo nyomon követni az összes SageMaker modell párhuzamos példáját, vagy részt venni a következőn kiosztott oktatóműhelyeket. Ha többet szeretne megtudni a szilánkos adatok párhuzamosságáról, kérjük, tekintse meg a dokumentáció.

A szerzőkről

Emily Webber közvetlenül a SageMaker elindítása után csatlakozott az AWS-hez, és azóta is próbálja elmondani a világnak! Azon kívül, hogy új ML-élményeket építhet az ügyfelek számára, Emily élvezi a meditációt és a tibeti buddhizmus tanulmányozását.

Can Karakus az AWS vezető alkalmazott tudósa, aki az AWS-en optimalizálja a nagyszabású elosztott mély tanulást. Kutatási területe a mély tanulás, az elosztott optimalizálás, az elosztott rendszerek és az információelmélet. Munkán kívül szeret kerékpározni, utazni, olvasni és tanulni.

Rahul Huilgol az AWS vezető szoftvermérnöke. Elosztott mély tanulási rendszereken dolgozik, hogy megkönnyítse és hatékony legyen a nagy mélységű tanulási modellek felhőben betanítása. Szabadidejében szeret fényképezni, kerékpározni és kertészkedni.

Suhit Kodgule egy szoftverfejlesztő mérnök az AWS mesterséges intelligencia csoporttal, amely mély tanulási keretrendszereken dolgozik. Szabadidejében szeret túrázni, utazni és főzni.

Erin Ho az AWS Deep Learning termékmenedzsere. Olyan termékeken dolgozik, amelyek megkönnyítik az ügyfelek számára a mély tanulási modellek AWS-en való betanítását. Külső munkája miatt szeret túrázni és síelni.