Hvernig á að laga „PyTorch CUDA Out of Memory“ villur við þjálfun líkana

PyTorch þjálfunarkeyrsla getur virkað í nokkur skref og stöðvað síðan með torch.OutOfMemoryError eða skilaboðum eins og CUDA out of memory. Tried to allocate .... Bein orsök er einföld: næsta CUDA úthlutun passaði ekki. Nytsami spurningin er hvers vegna hún passaði ekki.

Í þjálfun getur GPU-minni haldið líkistikum, stuðlum, ástandi lágmarksfallara, inntaksþenjum, tímabundnum vinnusvæðum og virkjunum sem geymdar eru fyrir bakdreifingu. PyTorch notar einnig geymslu minnisstýringu, svo minni sem sýnt er sem „reserved“ er ekki eins og minni sem er núverandi undir lifandi þenjum. Þessi munur skiptir máli þegar ákveðið er hvort minnka eigi vinnuálag eða rannsaka sundrun minnisstýringar.

Þessi leiðbeining fylgir núverandi PyTorch skjölun og notar núverandi AMP API nöfn. Sérstaklega skráir PyTorch nú torch.amp.autocast("cuda") og torch.amp.GradScaler("cuda"); eldri torch.cuda.amp.* inngangspunktar eru úreltir. Sjá PyTorch Automatic Mixed Precision skjölun.

Hraðgreining: hvaða tegund af OOM er um að ræða?

Skref 1: Mæltu minni áður en þjálfunaruppskriftin er breytt

Byrjaðu á að skrá batch-stærð, inntaksvíddir, nákvæmni og þau atriði þar sem bilunin á sér stað. Skoðaðu síðan bæði minni undir lifandi þenjum og „reserved“ minni í minnisstýringunni. PyTorch birtir memory_allocated(), memory_reserved(), toppafbrigði og memory_summary(). Núverandi CUDA minnisstýringarskjöl útskýra að geymslu minnisstýringin haldi endurnýtanlegum blokkum, sem er ástæðan fyrir því að ónotað „reserved“ minni getur enn birst sem notað í GPU vöktunarverkfærum. PyTorch CUDA minnisstýring.

import torch

torch.cuda.reset_peak_memory_stats()

# Keyrðu eitt dæmigerð þjálfunarþref hér.

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

Ef einföld samantekt nægir ekki, getur PyTorch tekið minnisstýringarsnapskot fyrir dýpri greiningu. Minnisverkfæri þess geta skráð úthlutunarsögu og búið til snapskot sem hægt er að skoða með PyTorch minnismyndrænum. PyTorch tekur fram að þessi verkfæri sjái minni sem stýrt er af PyTorch minnisstýringunni; úthlutanir sem gerðar eru beint af öðrum CUDA bókasöfnum gætu ekki birst þar. PyTorch leiðbeiningar um skilning á CUDA minnisnotkun.

Ekki líta á torch.cuda.empty_cache() sem almenna OOM lausn

torch.cuda.empty_cache() losar ónotaðar geymslublokkir svo að önnur GPU forrit geti notað þær. PyTorch segir skýrt að það losi ekki minni sem er undir lifandi þenjum og eykur því ekki magn GPU minnis sem er í boði fyrir PyTorch fyrir þenjur sem eru enn lifandi. Það getur verið gagnlegt á milli aðskildra tilrauna eða eftir að stórum hlutum er eytt, en það er ekki staðgengill fyrir að minnka virka minnisfótspor.

Skref 2: Minnkaðu virku vinnusettina fyrst

Áreiðanlegasta fyrsta lagfæringin er yfirleitt minni smásæmi (micro-batch): fjöldi sýna sem eru unnar í einni fram-/bakdreifingu. Virkjunarminni vex oft með batch-stærð, myndupplausn, runulengd og öðrum inntaksvíddum. Ef líkanið þjálfast við batch-stærð 32 en bilast við 64, er minnkun á batch ekki vinnubrögð í neikvæðum skilningi; það er bein minnkun á hámarks minnisþörf.

Fyrir myndir getur lækkun á rúmfræðilegri upplausn eða skurðarstærð skipt miklu máli. Fyrir transformer og önnur runulíkön getur minnkun á runulengd verið enn mikilvægari vegna þess að sum millistig þenja vaxa sterkt með runulengd. Nákvæm kvarðun fer eftir arkitektúr, svo mældu frekar en að gera ráð fyrir.

Passaðu einnig að matkóði sé ekki að byggja stuðla án þess að þörf sé á. PyTorch afköstaleiðbeiningar mæla með því að slökkva á stuðlareikningi fyrir staðfestingu eða ályktun þegar stuðlar eru ekki nauðsynlegir, vegna þess að autograd geymir annars millibuffer. Dæmigerð mynstur er:

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)

Í þjálfun, notaðu optimizer.zero_grad(set_to_none=True) nema reikniritið þitt treysti á hegðunarmuninn á núllstuðli og None stuðli. PyTorch lágmarksfallara skjöl segja að setja stuðla í None hafi yfirleitt lægra minnisfótspor og geti bætt afköst lítillega. PyTorch lágmarksfara zero_grad skjölun.

Skref 3: Haltu stærri virkri batch-stærð með AMP og stuðlasöfnun

Notaðu Automatic Mixed Precision þegar líkanið styður það

Automatic Mixed Precision (AMP) keyrir hæfar aðgerðir í lægri nákvæmni en heldur aðgerðum sem þurfa meira svið eða nákvæmni í viðeigandi tegundum. PyTorch skráir að AMP geti bætt afköst og minnkað minnisfótspor fyrir mörg CUDA vinnuálag, en það er ekki tölulega hentugt fyrir öll líkön. Sérstaklega varar PyTorch við að sum líkön sem eru fyrirþjálfuð í bfloat16 geti ofrunnið í float16.

Núverandi CUDA AMP þjálfunarmynstur er:

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 mælir með því að keyra framdreifingu og tap undir autocast, en yfirgefa autocast samhengið áður en bakdreifing hefst. Ef float16 veldur óstöðugleika, rannsakaðu hvort bfloat16 sé stutt og viðeigandi fyrir vélbúnaðinn og líkanið þitt, frekar en að gera ráð fyrir að allir blandaðir nákvæmnishættir hegði sér eins.

Notaðu stuðlasöfnun þegar þú þarft á stærri virkri batch-stærð að halda

Stuðlasöfnun vinnur nokkur minni smásæmi áður en lágmarksfallarinn er uppfærður. Ef smásæmið er 4 og þú safnar 8 þrefum, er virk batch-stærð fyrir eina lágmarksfarauppfærslu 32 sýnir á hverja vinnuþráð, að gefnu að hvert smásæmi hafi fjórar sýnir og gagnasamsíða uppsetningin breyti ekki þeirri útreikningi.

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)

Fyrir framleiðslukóða, meðhöndlaðu einnig síðasta hluta safnunar glugga þegar fjöldi batcha er ekki deilanlegur með accum_steps. Ef þú notar dreifða þjálfun, getur hegðun stuðlasamstillingar breytt minnis/afkasta viðskiptum, svo fylgdu leiðbeiningum dreifðs API um söfnun frekar en að afrita ein-GPU lykkju óbreytta.

Skref 4: Skiptu reikniaflinu fyrir minni, rannsakaðu síðan varðveislu og sundrun

Virkjunarpunktar (Activation checkpointing)

Virkjunarpunktar minnka minni með því að halda ekki valdar framdreifingarvirkjanir lifandi þar til bakdreifing. Í staðinn endurreiknar PyTorch þær í bakdreifingu. Þetta skiptir auknu reikniafli fyrir lægra virkjunarminnisfótspor. Núverandi PyTorch checkpoint skjöl mæla með því að gefa use_reentrant=False skýrt. PyTorch virkjunarpunkta skjölun.

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)

Veldu lög með virkjunarpunktum sem hafa stórar geymdar virkjanir og ásættanlegan endurreikningskostnað. Ekki gera ráð fyrir að virkjunarpunktar fyrir hverja aðgerð séu bestir; það getur hægt verulega á þjálfun.

Leitaðu að þenjum sem halda reikniritagröfum lifandi

Ef minni vex í hverri endurtekningu frekar en að ná toppi við u.þ.b. sama stig, skoðaðu Python tilvísanir. Algengt mynstur er að geyma graf-tengdar þenjur í lista:

# Hættulegt ef haldið er í mörg þref:
loss_history.append(loss)

# Geymdu Python tölu í staðinn:
loss_history.append(loss.item())

Sama vandamálið getur komið upp þegar þú geymir líkanið úttak, athygliskort, falin ástand eða staðfestingarþenjur án þess að aðskilja þær eða færa þær af GPU. Eyddu tilvísunum sem þú þarft ekki lengur, og notaðu detach() aðeins þegar þú vilt meðvitað að þenja sé aðskilin frá autograd.

Stilltu minnisstýringuna aðeins eftir að tölfræði bendir til sundrunar

Núverandi PyTorch skjöl kjósa umhverfisbreytuna PYTORCH_ALLOC_CONF. Eldri PYTORCH_CUDA_ALLOC_CONF er enn alias fyrir bakflæðisvirkni. Þessi nafngreiningarbreyting átti sér stað í núverandi skjölum, svo nýjar stillingar ættu að nota kjörna nafnið. PyTorch CUDA umhverfisbreytur.

Tveir minnisstýringarvalkostir eru sérstaklega viðeigandi:

• expandable_segments:True er tilraunakennt og hannað til að minnka ónothæf minnisbrot þegar úthlutunarstærðir breytast, svo sem vinnuálag þar sem batch eða þenjustærðir breytast.
• max_split_size_mb getur minnkað sundrun með innbyggðri minnisstýringu, en PyTorch lýsir því skýrt sem lokaraðgerð fyrir vinnuálag sem bilast með OOM á meðan sýnt er mikið magn af óvirkum klofningsblokkum. Það getur einnig skaðað afköst og er hunsað af cudaMallocAsync bakendanum.

# Dæmi fyrir vinnuálag með breytilegum úthlutunarstærðum:
export PYTORCH_ALLOC_CONF=expandable_segments:True

Ekki afrita minnisstýringarflagg frá annarri tölvu án þess að skoða memory_summary() eða snapskot. Sannur afkastavandamál—þar sem lifandi þenjur fylla nú þegar GPU—verður ekki leyst með sundrunarstillingum.

Þegar eitt GPU getur samt ekki haldið líkaninu

Ef ein sýn við minnstu mögulegu inntaksstærð veldur samt OOM, gæti vandamálið verið líkanið og ástand lágmarksfallarans frekar en batch-stærðin. Á því stigi, íhugaðu minni arkitektúr, lægri nákvæmnisþenjur þar sem tölulega viðeigandi, CPU/afhleðslu aðferðir, eða skipta dreifða þjálfun.

PyTorch Fully Sharded Data Parallel (FSDP) getur skipt líkistikum milli gagnasamsíða vinnuþráða, og FULL_SHARD stefnan skiptir einnig stuðlum og ástandi lágmarksfallara. Þetta getur minnkað minni á hverju GPU samanborið við fullt endurtekna gagnasamsíða, á kostnað samskipta og flóknari þjálfunarhegðunar. PyTorch FSDP skjölun.

Hagnýt röð aðgerða

Checklist: hvernig á að vita að OOM sé í raun lagað

• Keyrðu nokkur dæmigerð þjálfunarþref, ekki bara eina vel heppnaða framdreifingu.
• Endurstilltu og skráðu max_memory_allocated() svo þú vitir nýja toppinn.
• Staðfestu að GPU minni nái stöðugu bili frekar en að aukast í hverri endurtekningu.
• Staðfestu tap og stuðla eftir að blönduð nákvæmni er virkjuð.
• Staðfestu að stuðlasöfnun varðveiti lágmarksfarauppfærslutímaáætlunina sem þú vildir.
• Keyrðu staðfestingarþref undir torch.no_grad() þegar stuðlar eru ekki krafist.
• Ef þú breyttir minnisstýringarstillingum, berðu saman minnistölfræði og gegnumstreymi áður og eftir.
• Ekki kalla vandamálið leyst aðeins vegna þess að nvidia-smi sýni minna „reserved“ minni eftir empty_cache(); þjálfunarvinnuálagið sjálft verður að ljúka við eðlilegan topp sinn.

CUDA OOM er best að meðhöndla sem minnisfjárhæðarvandamál, ekki sem eina PyTorch villu. Mældu toppinn, minnkaðu virku vinnusettina fyrst, notaðu síðan blandaða nákvæmni, söfnun og virkjunarpunkta sem meðvitaðar viðskiptir. Færðu þig aðeins yfir í minnisstýringarstillingar þegar minnisstýringartölfræði bendir til sundrunar, og færðu þig yfir í skiptingu eða annað líkan þegar líkanið sjálft passar ekki lengur þægilega á einu GPU.