Harjoittele jättimäisiä malleja lähes lineaarisella skaalauksella käyttämällä sirpaloitujen tietojen rinnakkaisuutta Amazon SageMakerissa

Tavoittelemalla ylivertaista tarkkuutta syväoppimismallit sellaisilla aloilla kuin luonnollisen kielen käsittely ja tietokonenäkö ovat kasvaneet merkittävästi viime vuosina, ja ne lasketaan usein kymmenistä satoihin miljardeihin parametreihin. Näiden jättimäisten mallien kouluttaminen on haastavaa ja vaatii monimutkaisia ​​jakelustrategioita. Datatieteilijät ja koneoppimisinsinöörit etsivät jatkuvasti parasta tapaa optimoida koulutuslaskentansa, mutta he kamppailevat kuitenkin ylimääräisten viestintäkulujen kanssa, jotka voivat kasvaa klusterin kokonaiskoon myötä.

Tästä syystä aloitimme äskettäin skova data rinnakkaisuus on Amazon Sage Maker, uusi muistia säästävä hajautettu harjoitustekniikka SageMaker mallin rinnakkaiskirjasto (SMP).. Jaettujen tietojen rinnakkaisuus on suunniteltu äärimmäisen mittakaavan malleihin ja käyttää Amazonia talon sisällä MiCS hupun alla oleva tekniikka, tieteellinen pyrkimys minimoida viestintämittakaava alentamalla parametrien keräämiseen ja gradienttisynkronointiin perustuvaa kallista tiedonsiirtoa. 30B-parametrin GPT-2-mallilla, jonka sekvenssipituus on 2048, tämä uusi ominaisuus saavutti 141 TFLOP:a, mikä on 39.7 % nopeampi verrattuna DeepSpeed ​​ZeRO-3:een. 10B GPT-2-mallissa, jonka sekvenssin pituus on 512, tämä uusi ominaisuus saavutti myös 564 näytettä sekunnissa, mikä on 13.9 % nopeampi verrattuna PyTorchin Fully Sharded Data Parallel (FSDP) -datan rinnalle. Muista, että jättimäisessä mallikoulutuksessa jokainen nopeutetun prosenttiosuus merkitsee säästöjä ja tiimisi tuottavuutta.

Tässä blogiviestissä tarkastellaan ensin lähemmin sirpaloidun datan rinnakkaisuuden keskeisiä erottajia ja sitä, milloin sitä käytetään. Sitten opit opettamaan 30 B parametrin GPT-2-mallia SageMakerissa helposti tämän uuden ominaisuuden avulla. Lopuksi verrataan suorituskykyä muihin avoimen lähdekoodin vaihtoehtoihin, jotka ylittävät DeepSpeed ​​ZeRO:n jopa 39.7 % 256 GPU:lla.

Kuinka sirpaloidun tiedon rinnakkaisuus toimii ja milloin sitä käytetään

Ennen kuin esittelemme sirpaloidun datan rinnakkaisuuden, katsotaanpa sen laajempaa tekniikkaperhettä. Viimeaikaiset hajautetut koulutuslähestymistavat suurille malleille ovat siirtyneet paradigmaan, jossa malliparametrit, gradientit ja optimointitilat jaetaan datan rinnakkaisten solmujen kesken. Toisin kuin Pipeline Parallelism, jolla on luontainen monimutkaisuus valita tasoja osioitaviksi eri laitteille varsinkin kun kehys ei tue automaattinen mallin jakaminenTämä paradigma säilyttää tyylikkäästi tietojen rinnakkaisuuden yksinkertaisuuden ja poistaa samalla tietojen rinnakkaisuuden rajoitteen, jossa mallin on mahduttava yhteen GPU:hun.

Olemassa olevissa tämän paradigman piirissä olevissa kehyksissä, erityisesti DeepSpeed ​​ZeRO-3:ssa ja PyTorchin FSDP:ssä FairScalesta ylävirtaan, mallitilat ovat hajallaan. kaikki GPU:t, strategia, joka alentaa muistin kulutusta jokaisessa grafiikkasuorittimessa suuren tiedonsiirron kustannuksella, joka kasvaa klusterin koon myötä ja saa siten skaalautuvuuden laskemaan merkittävästi mittakaavassa. Sitä vastoin sirpaloitujen tietojen rinnakkaisuus SMP-kirjastoosiomallissa on kohdassa a mittakaavatietoinen tavalla osioiden jokaisen mallin tilan replikan vain sisällä osajoukko GPU:sta.

Katsotaanpa tarkemmin mittakaavatietoinen mallin osiointi MiCS:ssä, rinnakkaisen sirpaloidun tiedon takana oleva ydinteknologia. Tämän suunnittelun taustalla oleva intuitio on, että koulutustilojen osiointia koko data-rinnakkaisryhmässä ei välttämättä vaadita mallin kouluttamiseen, jossa on kymmeniä miljardeja parametreja. Esimerkiksi 8 V100 GPU:ta (kukin 32 Gt) riittää pitämään mallitilat kopio 10B-parametrisesta mallista, joka tarvitsee noin 200 Gt muistia, kun harjoittelet Adam Optimizerilla sekatarkkuudella. Rajoittamalla mallin tilojen täydellinen kopio pienin Grafiikkasuorittimien osajoukko, voimme tehokkaasti pienentää tietoliikenteen laajuutta verrattuna DeepSpeed- ja PyTorch FSDP:hen. Jaettu data rinnakkain hyödyntää myös muita MiCS-tekniikoita, kuten hierarkkista viestintää ja 2-hop gradientsynkronointia. Katso lisätietoja AWS:n jättimäisen mallikoulutuksen lähes lineaarinen skaalaus or MiCS: lähes lineaarinen skaalaus jättimäisen mallin kouluttamiseen julkisessa pilvessä.

Mistä tiedät, milloin sinun tulee valita sirpaloitu data rinnakkain muihin hajautettuihin harjoitustekniikoihin verrattuna? Yleissääntönä on, että jos mallissasi on alle miljardi parametria ja se mahtuu GPU-muistiin, SageMaker-tietojen rinnakkaiskirjasto or SageMaker koulutuksen kääntäjä voi riittää sinulle. Jos sinulla on suurempia kieli- tai tietokonenäkömalleja, ehdotuksemme on kouluttaa se sirpaloidun datan rinnakkaistekniikalla yhdistettynä aktivoinnin tarkistuspiste ja aktivoinnin purkaminen SageMaker-mallin rinnakkaiskirjastossa ensin, ennen muita tekniikoita, kuten tensorin rinnakkaisuus tai putkilinjan yhdensuuntaisuus.

Sirpaloidun datan rinnakkaisuuden käyttäminen GPT-2:n harjoittamiseen Amazon SageMakerissa

Opitaan nyt kouluttamaan GPT-2-mallia, jossa on sirpaloituja tietoja rinnakkain, ja SMP kapseloi monimutkaisuuden puolestasi. Tämä täydellinen opetusohjelmamuistikirja opastaa sinut koko prosessin läpi tietojenkäsittelystä, koulutustöiden määrittämisestä ja lähettämisestä koulutuslokien seurantaan. Seuraavassa on lyhyt yleiskatsaus tämän ominaisuuden käytön tärkeimmistä vaiheista.

1. Aloittaa

Jaettujen tietojen rinnakkaisuus on saatavilla PyTorch v1.12.0+:ssa ja toimii sekä FP16:n että BF16:n kanssa. Helpoin tapa käyttää SMP-kirjastoa on PyTorchin valmiiksi rakennetun AWS Deep Learning Containerin kautta. Jos kuitenkin haluat tuoda oman Docker-konttisi, voit viitata Luo oma Docker-säiliö SageMaker Distributed Model Parallel Library -kirjaston avulla. Aloita seuraamalla Muokkaa PyTorch-harjoituskomentosarjaa mukauttaaksesi SMP-sovellusliittymiä harjoitusskriptiisi. Tässä osiossa mainitaan vain muutama päävaihe koodinpätkillä käyttövalmiin harjoitusskriptistä train_gpt_simple.py. Voit seurata käsikirjoituksen kommentteja ja API-asiakirja saadaksesi lisätietoja siitä, missä SMP-sovellusliittymiä käytetään.

Ensin tuo ja alusta kirjasto soittamalla smdistributed.modelparallel.torch.init() harjoitusskriptin alussa:

import smdistributed.modelparallel.torch as smp

smp.init(smp_config)

Toiseksi, kääri osioitava malli smdistributed.modelparallel.torch.DistributedModel ja käytä palautettua DistributedModel kohde eteenpäin:

from transformers import AutoModelForCausalLM

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

Kääri optimoinnin kanssa smdistributed.modelparallel.torch.DistributedOptimizer optimointitilojen tallentamiseen ja lataamiseen.

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

Laita eteenpäin ja taaksepäin logiikka askelfunktioon ja koristele se smdistributed.modelparallel.torch.step.  Mikä tahansa sisällä määritetty laskenta smp.step-decorated toiminto suoritetaan hajautetusti.

@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. Valmistele tietojoukko

Käytämme avoin verkkoteksti on tässä esimerkissä käyttämämme tietojoukko. Muistikirja käyttää käsikirjoitusta data_prep_512.py ladataksesi ja esikäsitelläksesi tietojoukon. Voit myös harjoitella muiden tietojoukkojen kanssa muokkaamalla data_pipeline.py. Kun käsittelet suurta tietojoukkoa ja mallia, voit nopeuttaa harjoitustyötä käyttämällä tallennettuja tietoja Amazon FSx Lusterille, joka tarjoaa korkean suorituskyvyn tiedostojärjestelmän, joka on integroitu natiivisti Amazonin yksinkertainen tallennuspalvelu (S3). Katso ohjeet osoitteesta Määritä tiedonsyöttökanava käyttämään Amazon FSx:ää Lusterille saadaksesi ohjeita FSx Luster -tiedostojärjestelmän asettamisesta tiedonsyöttökanavaksi.

3. Aloita koulutustyöt

Tämä vaihe olettaa, että olet jo tehnyt muokkasi harjoitusohjelmaasi ja valmisti tietojoukon edellisissä osissa mainitulla tavalla. Vastaanottaja Ota sirpaloitujen tietojen rinnakkaisuus käyttöön, aseta vain sharded_data_parallel_degree vuonna PyTorch estimaattori. Tässä opetusohjelmassa asetamme sharded_data_parallel_degree=128 ja instace_count=32 p4d.24xlarge-solmuille, mikä osoittaa, että mallin tilat jaetaan 128 GPU:lle 256 GPU:n joukosta. Tämän valitun arvon perusteella SMP asettaa sitten automaattisesti tietojen rinnakkaisasteen arvoon 2 (koska 256/128=2), mikä tarkoittaa, että meillä on kaksi kopiota tietojen rinnakkaisuudesta. Yleinen sääntö ihanteellisen arvon valitsemiseksi sharded_data_parallel_degree on lisätä yksi solmu lisää jakamisryhmään jokaista 3B malliparametria kohden. Tässä opetusohjelmassa mallimme koko on 30 B, joten meidän tulisi käyttää vähintään 10 solmua sirpalointiin. Ja koska 16 solmua (128 GPU:ta) on pienin teho 2:sta kynnyksen yläpuolella, asetimme sharded_data_parallel_degree=128.

Tarkastuspisteitä varten tarjoamme myös joukon tarkistuspisteen apuohjelmia sharded_data_parallel_checkpoint.py , mukaan lukien apuohjelma, jolla rekonstruoidaan koko state_dict edistyneisiin käyttötapauksiin. Lopuksi voimme käynnistää hajautetun harjoittelun kutsumalla fit() estimaattorissa.

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. Valvo koulutustöitä

Voit käyttää harjoituslokeja ja seurata grafiikkasuorittimen ja muistin käyttöä amazonin pilvikello. Muista katsoa "algo-1":n lokit koska se on pääsolmu, jonka lähtövirrassa on koulutustyölokit kaikista esiintymistä.

Suorituskyky benchmarking

Vertailimme sirpaloitujen tietojen rinnakkaisuutta SMP-kirjastossa sekä 16:ssa että 32:ssa p4d.24xlarge-solmussa sekvenssin pituuksilla 512 ja 2048. 30B-parametrin GPT2-malli on määritetty käyttämään piilotettua leveyttä 7168, 48 kerrosta ja 64 päätä. Voit ottaa käyttöön täsmälleen saman kokoonpanon, jossa sekvenssin pituus on 2048 asettamalla model_config = "gpt2-30b" opetusohjelman muistikirjassa. Tällä asetuksella SMP saavutti 73.52 näytettä sekunnissa, mikä on 39.7 % nopeampi verrattuna DeepSpeed ​​ZeRO-3:een. Jos tunnuksesi koko on 500 miljardia, tämä nopeus tarkoittaa lähes 367 tunnin säästöä p4d.24xlarge-solmuissa, mikä vastaa yli 12,000 XNUMX dollarin budjettisäästöä harjoittelua kohden! Seuraavassa taulukossa on yhteenveto vertailutuloksistamme.

1/ Testasimme DeepSpeed ​​ZeRO:ssa erilaisia ​​ominaisuuksia, vaiheita ja kokoonpanoja jokaiselle mallikonfiguraatiolle ja valitsimme DeepSpeed-perusviivaksi sen, joka tarjoaa parhaan suorituskyvyn. Vertailuarvoa ajettiin Amazonin elastinen laskentapilvi (Amazon EC2). 2/ Nämä tulokset perustuvat parannettuihin viestintäryhmiin, jotka on optimoitu AWS:lle ja jotka tulevat saataville pian. 3/ Harjoitukseen kuluva aika ennustetaan nopeudesta käsiteltyjen merkkien määrän perusteella.

Yhteenvetona voidaan todeta, että havaitsimme jatkuvasti korkeamman suorituskyvyn sirpaloidun tiedon rinnakkaisuuden SMP:ssä verrattuna DeepSpeediin useissa malleissa ja kokoonpanoissa. Tämä uusi ominaisuus osoitti myös paremman muistin tehokkuuden kuin DeepSpeed, mikä mahdollisti SMP:n sovittamaan suuremman eräkoon ja vähentämään gradientin kertymisen tasoa, joka tarvitaan sovittamaan tiettyyn globaaliin eräkokoon.

Yhteenveto

Tässä viestissä esittelimme uuden hajautetun harjoitustekniikan – sirpaloidun datan rinnakkaisuuden – ja kuinka se nopeuttaa jättimäistä malliharjoittelua lähes lineaarisella skaalalla Amazon SageMakerissa. Kävimme myös läpi kuinka koulutetaan GPT-2-malli uudella tekniikalla tämän jälkeen täydellinen esimerkki. Voit seurata Amazon SageMaker -esimerkkejä GitHub-reposta seurata kaikkia SageMaker-mallien rinnakkaisia ​​esimerkkejä tai osallistua seuraavaan jaettu koulutustyöpajoja. Lisätietoja sirpaloidun tiedon rinnakkaisuudesta on kohdassa dokumentointi.

Tietoja kirjoittajista

Emily Webber liittyi AWS:ään heti SageMakerin julkaisun jälkeen ja on yrittänyt kertoa siitä maailmalle siitä lähtien! Asiakkaille uusien ML-elämysten rakentamisen lisäksi Emily nauttii meditoinnista ja Tiibetin buddhalaisuuden opiskelusta.

Voiko Karakus on AWS:n vanhempi soveltuva tutkija, joka optimoi laajamittaisen hajautetun syväoppimisen AWS:ssä. Hänen tutkimusalueensa ovat syväoppiminen, hajautettu optimointi, hajautetut järjestelmät ja tietoteoria. Työn ulkopuolella hän nauttii pyöräilystä, matkustamisesta, lukemisesta ja oppimisesta.

Rahul Huilgol on vanhempi ohjelmistosuunnittelija AWS:ssä. Hän työskentelee hajautettujen syväoppimisjärjestelmien parissa, jotta suurten syväoppimismallien opettaminen pilvessä olisi helppoa ja tehokasta. Vapaa-ajallaan hän harrastaa valokuvaamista, pyöräilyä ja puutarhanhoitoa.

Suhit Kodgule on ohjelmistokehitysinsinööri, jolla on AWS-keinoälyryhmä, joka työskentelee syvän oppimiskehyksen parissa. Vapaa-ajallaan hän harrastaa patikointia, matkustamista ja ruoanlaittoa.

Erin Ho on AWS Deep Learningin tuotepäällikkö. Hän työskentelee tuotteiden parissa, joiden avulla asiakkaiden on helpompi kouluttaa syväoppimismalleja AWS:ssä. Ulkopuolisen työn iloksi hän harrastaa patikointia ja hiihtoa.