Kuinka korjata "PyTorch CUDA Out of Memory" -virhe mallin koulutuksen aikana

PyTorch-koulutusajo voi toimia useita vaiheita ja pysähtyä sitten torch.OutOfMemoryError-virheeseen tai viestiin, kuten CUDA out of memory. Tried to allocate .... Välitön syy on yksinkertainen: seuraava CUDA-muistivaraus ei mahtunut. Hyödyllinen kysymys on miksi se ei mahtunut.

Koulutuksen aikana GPU-muisti voi sisältää malliparametrit, gradientit, optimointitilan, syöttötensorit, väliaikaiset työtilat ja taaksepäinpropagaatiota varten tallennetut aktivoinnit. PyTorch käyttää myös välimuistiallokaattoria, joten "varattu"-muistina näytetty muisti ei ole sama kuin elävien tensorien tällä hetkellä käyttämä muisti. Tämä ero on tärkeä päätettäessä, kannattaako työmäärää pienentää vai tutkia allokaattorin fragmentaatiota.

Tämä opas seuraa nykyistä PyTorch-dokumentaatiota ja käyttää nykyisiä AMP API -nimiä. Erityisesti PyTorch dokumentoi nyt torch.amp.autocast("cuda") ja torch.amp.GradScaler("cuda"); vanhemmat torch.cuda.amp.* -päätepisteet ovat vanhentuneita. Katso PyTorch Automatic Mixed Precision -dokumentaatio.

Nopea triage: minkä tyyppinen OOM-virhe on kyseessä?

Vaihe 1: Mittaa muisti ennen koulutusreseptin muuttamista

Aloita kirjaamalla batch-koko, syöttömitat, tarkkuus ja kohta, jossa virhe tapahtuu. Tarkista sitten sekä elävien tensorien muisti että allokaattorin varaama muisti. PyTorch tarjoaa memory_allocated(), memory_reserved(), huippuarvomuunnelmat ja memory_summary(). Nykyinen CUDA-muistinhallintadokumentaatio selittää, että välimuistiallokaattori säilyttää uudelleenkäytettävät lohkot, minkä vuoksi käyttämätön varattu muisti voi silti näyttää käytetyltä GPU-seurantaohjelmissa. PyTorch CUDA-muistinhallinta.

import torch

torch.cuda.reset_peak_memory_stats()

# Suorita yksi edustava koulutusvaihe tässä.

print("allocated GB:",
      torch.cuda.memory_allocated() / 1024**3)
print("reserved GB:",
      torch.cuda.memory_reserved() / 1024**3)
print("peak allocated GB:",
      torch.cuda.max_memory_allocated() / 1024**3)
print(torch.cuda.memory_summary(abbreviated=True))

Jos yksinkertainen yhteenveto ei riitä, PyTorch voi tallentaa allokaattorin tilannekuvia syvempää analyysiä varten. Sen muistityökalut voivat tallentaa varaushistorian ja tuottaa tilannekuvan, jota voidaan tutkia PyTorch-muistivisualisoijalla. PyTorch huomauttaa, että nämä työkalut näkevät PyTorch-allokaattorin hallinnoiman muistin; muiden CUDA-kirjastojen tekemät suorat varaukset eivät välttämättä näy siellä. PyTorch-opas CUDA-muistinkäytön ymmärtämiseen.

Älä pidä torch.cuda.empty_cache():a yleisenä OOM-korjauksena

torch.cuda.empty_cache() vapauttaa käyttämättömät välimuistilohkot, jotta muut GPU-sovellukset voivat käyttää niitä. PyTorch toteaa nimenomaisesti, ettei se vapauta elävien tensorien varaamaa muistia eikä siten lisää PyTorchin käytettävissä olevaa GPU-muistia yhä eläville tensorille. Se voi olla hyödyllinen erillisten kokeiden välillä tai suurten objektien poistamisen jälkeen, mutta se ei korvaa aktiivisen muistijalanjäljen pienentämistä.

Vaihe 2: Pienennä ensin aktiivinen työjoukko

Luotettavin ensimmäinen korjaus on yleensä pienempi mikrobatcher: yhden eteenpäin/taaksepäin-vaiheen käsittelemien näytteiden määrä. Aktivointimuisti kasvaa yleensä batch-koon, kuvan resoluution, sekvenssin pituuden ja muiden syöttömittojen mukana. Jos malli koulutetaan batch-kokoon 32 mutta epäonnistuu kokoon 64, batch-koon pienentäminen ei ole kielteisessä mielessä kiertotie; se on huippumuistitarpeen suoraa pienentämistä.

Kuvien kohdalla tilallisen resoluution tai rajauskoon pienentäminen voi tehdä suuren eron. Transformer-malleissa ja muissa sekvenssimalleissa sekvenssin pituuden pienentäminen voi olla vielä tärkeämpää, koska jotkin välitensorit kasvavat voimakkaasti sekvenssin pituuden myötä. Tarkka skaalaus riippuu arkkitehtuurista, joten mittaa älä oletta.

Varmista myös, että evaluointikoodi ei rakenna gradientteja tarpeettomasti. PyTorchin suorituskykyohjeet suosittelevat gradienttilaskennan poistamista käytöstä validointia tai päättelyä varten, kun gradientteja ei tarvita, koska autograd tallentaa muuten välipuskureita. Tyypillinen kuvio on:

model.eval()
with torch.no_grad():
    for x, y in val_loader:
        x = x.cuda(non_blocking=True)
        y = y.cuda(non_blocking=True)
        pred = model(x)

Koulutuksen aikana käytä optimizer.zero_grad(set_to_none=True), ellei algoritmissasi luoteta nollagradientin ja None-gradientin väliseen käyttäytymiserotukseen. PyTorchin optimointiohjeistus toteaa, että gradienttien asettaminen arvoon None on yleensä muistitehokkaampaa ja voi parantaa suorituskykyä hieman. PyTorch optimizer zero_grad -dokumentaatio.

Vaihe 3: Säilytä suurempi tehollinen batch AMP:n ja gradienttien kertymisen avulla

Käytä Automatic Mixed Precisionia, kun malli tukee sitä

Automatic Mixed Precision (AMP) suorittaa kelvolliset operaatiot matalammalla tarkkuudella pitäen operaatiot, jotka vaativat suuremman alueen tai tarkkuuden, sopivissa tyypeissä. PyTorch dokumentoi, että AMP voi parantaa suorituskykyä ja pienentää muistijalanjälkeä monissa CUDA-työmäärissä, mutta se ei ole numeerisesti sopiva jokaiselle mallille. Erityisesti PyTorch varoittaa, että jotkin bfloat16:lla esikoulutetut mallit voivat ylittää float16:n rajat.

Nykyinen CUDA AMP -koulutuskuvio on:

scaler = torch.amp.GradScaler("cuda")

for inputs, targets in train_loader:
    inputs = inputs.cuda(non_blocking=True)
    targets = targets.cuda(non_blocking=True)
    optimizer.zero_grad(set_to_none=True)

    with torch.amp.autocast("cuda", dtype=torch.float16):
        outputs = model(inputs)
        loss = loss_fn(outputs, targets)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

PyTorch suosittelee eteenpäinpropagaation ja häviön suorittamista autocastin alla, jättäen autocast-kontekstin ennen taaksepäinpropagaatiota. Jos float16 aiheuttaa epästabiilisuutta, tutki, tukeeko laitteistosi ja mallisi bfloat16:ta ja onko se sopiva, äläkä oleta kaikkien sekatarkkuustilojen käyttäytyvän identtisesti.

Käytä gradienttien kertymistä, kun tarvitset suuremman tehollisen batchin

Gradienttien kertyminen käsittelee useita pienempiä mikrobatcheja ennen optimointiohjelman päivitystä. Jos mikrobatcher on 4 ja kerrytät 8 vaihetta, yhden optimointipäivityksen tehollinen batch on 32 näytettä per työntekijä, olettaen että jokaisessa mikrobatcherissa on neljä näytettä ja data-parallel-asetus ei muuta tätä laskentaa.

accum_steps = 8
optimizer.zero_grad(set_to_none=True)

for step, (inputs, targets) in enumerate(train_loader):
    inputs = inputs.cuda(non_blocking=True)
    targets = targets.cuda(non_blocking=True)

    with torch.amp.autocast("cuda", dtype=torch.float16):
        outputs = model(inputs)
        loss = loss_fn(outputs, targets) / accum_steps

    scaler.scale(loss).backward()

    if (step + 1) % accum_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad(set_to_none=True)

Tuotantokoodissa käsittele myös lopullinen osittainen kertymisikkuna, kun batchien määrä ei ole jaollinen luvulla accum_steps. Jos käytät hajautettua koulutusta, gradienttien synkronointikäyttäytyminen voi muuttaa muistin/suorituskyvyn kompromissia, joten noudata hajautetun API:n kertymisohjeita äläkä kopioi yksittäisen GPU:n silmukkaa sellaisenaan.

Vaihe 4: Vaihtoehto laskenta muistille, tutki sitten säilytystä ja fragmentaatiota

Aktivointien checkpointing

Aktivointien checkpointing vähentää muistia pitämättä valittuja eteenpäin-aktivointeja elossa taaksepäinpropagaatioon asti. Sen sijaan PyTorch laskee ne uudelleen taaksepäinpropagaation aikana. Tämä vaihtaa lisälaskennan pienempään aktivointimuistijalanjälkeen. Nykyinen PyTorch-checkpoint-dokumentaatio suosittelee nimenomaisesti use_reentrant=False:n välittämistä. PyTorch aktivointien checkpointing -dokumentaatio.

from torch.utils.checkpoint import checkpoint

def forward(self, x):
    x = checkpoint(self.block1, x, use_reentrant=False)
    x = checkpoint(self.block2, x, use_reentrant=False)
    return self.head(x)

Tee checkpoint kerroksille, joilla on suuret tallennetut aktivoinnit ja hyväksyttävä uudelleenlaskennan kustannus. Älä oleta, että checkpointing jokaiselle operaatiolle on optimaalista; se voi hidastaa koulutusta merkittävästi.

Etsi tensorit, jotka pitävät laskentagraafit elossa

Jos muisti kasvaa jokaisella iteraatiolla sen sijaan, että se huipentuisi suunnilleen samaan tasoon, tutki Python-viittauksia. Yleinen kuvio on graafiin kytkettyjen tensorien tallentaminen listaan:

# Riskialtista, jos säilytetään monien vaiheiden ajan:
loss_history.append(loss)

# Tallenna Python-luku sen sijaan:
loss_history.append(loss.item())

Samasta ongelmasta voi olla kyse, kun välimuistissa olevat mallitulosteet, attention-kartat, piilotilat tai validointitensorit eivät ole irrotettu tai siirretty pois GPU:sta. Poista viittaukset, joita et enää tarvitse, ja käytä detach():a vain, kun haluat tietoisesti tensorin, joka on irrotettu autogradista.

Säädä allokaattoria vain, kun tilastot osoittavat fragmentaatiota

Nykyinen PyTorch-dokumentaatio suosii ympäristömuuttujaa PYTORCH_ALLOC_CONF. Vanhempi PYTORCH_CUDA_ALLOC_CONF on edelleen alias taaksepäin yhteensopivuuden vuoksi. Tämä nimeämistieto muuttui nykyisissä dokumenteissa, joten uusien konfiguraatioiden tulisi käyttää suositeltua nimeä. PyTorch CUDA-ympäristömuuttujat.

Kaksi allokaattorivaihtoehtoa ovat erityisen relevantteja:

• expandable_segments:True on kokeellinen ja suunniteltu vähentämään käyttökelpoisia muistisiruja, kun varauskoot muuttuvat, kuten työmäärissä, joiden batch- tai tensorikoot vaihtelevat.
• max_split_size_mb voi vähentää fragmentaatiota natiivilla allokaattorilla, mutta PyTorch kuvaa sen nimenomaisesti viimeisenä keinona työmäärille, jotka epäonnistuvat OOM-virheeseen näyttäen suuren määrän epäaktiivisia split-lohkoja. Se voi myös heikentää suorituskykyä, ja cudaMallocAsync-backend ohittaa sen.

# Esimerkki työmäärälle, jossa varauskoot vaihtelevat:
export PYTORCH_ALLOC_CONF=expandable_segments:True

Älä kopioi allokaattorilippuja toiselta koneelta tarkistamatta memory_summary():a tai tilannekuvaa. Todellista kapasiteettiongelmaa – jossa elävät tensorit täyttävät jo GPU:n – ei ratkaista fragmentaation säädöllä.

Kun yksi GPU ei silti mahduta mallia

Jos yksittäinen näyte pienimmällä käytännöllisellä syöttökoolla aiheuttaa edelleen OOM-virheen, ongelma voi olla malli ja optimointitila eikä batch. Tässä vaiheessa harkitse pienempää arkkitehtuuria, matalamman tarkkuuden parametreja numeerisesti sopivissa tapauksissa, CPU/offload-strategioita tai sharded hajautettua koulutusta.

PyTorchin Fully Sharded Data Parallel (FSDP) voi sharded malliparametrit data-parallel-työntekijöiden kesken, ja sen FULL_SHARD-strategia sharded myös gradientit ja optimointitilat. Tämä voi vähentää per-GPU-muistia verrattuna täysin replikoituun data-parallelismiin, viestinnän ja monimutkaisemman koulutuskäyttäytymisen kustannuksella. PyTorch FSDP -dokumentaatio.

Käytännön toimintajärjestys

Tarkistuslista: kuinka tietää, että OOM on todella korjattu

• Suorita useita edustavia koulutusiteraatioita, älä vain yhtä onnistunutta eteenpäinvaihetta.
• Nollaa ja tallenna max_memory_allocated(), jotta tiedät uuden huipun.
• Varmista, että GPU-muisti saavuttaa vakaan alueen sen sijaan, että se kasvaisi jokaisella iteraatiolla.
• Validoi häviö ja gradientit sekatarkkuuden käyttöönoton jälkeen.
• Varmista, että gradienttien kertyminen säilyttää suunnitellun optimointipäivitysaikataulun.
• Suorita validointivaihe torch.no_grad():n alla, kun gradientteja ei vaadita.
• Jos muuttit allokaattoriasetuksia, vertaa muistitilastoja ja läpimenoa ennen ja jälkeen.
• Älä kutsu ongelmaa ratkaistuksi pelkästään siksi, että nvidia-smi näyttää vähemmän varattua muistia empty_cache():n jälkeen; itse koulutustyömäärän on täytyttävä suorittaa normaali huippunsa.

CUDA OOM on käsiteltävä parhaiten muistibudongelmana, ei yhtenä PyTorch-bugina. Mittaa huippu, pienennä elävää työjoukkoa ensin, käytä sitten sekatarkkuutta, kertymistä ja checkpointingia tietoisina kompromisseina. Siirry allokaattorin säätöön vain, kun allokaattorin tilastot osoittavat fragmentaatiota, ja siirry shardingiin tai toiseen malliin, kun malli itsessään ei enää mahdu mukavasti yhdelle GPU:lle.