Kaip išspręsti „PyTorch CUDA Out of Memory“ klaidą mokymo metu

PyTorch mokymo procesas gali veikti kelis žingsnius, o tada sustoti su torch.OutOfMemoryError arba pranešimu, pvz., CUDA out of memory. Tried to allocate .... Tiesioginė priežastis yra paprasta: kitam CUDA paskyrimui neužteko vietos. Naudingas klausimas yra kodėl neužteko vietos.

Mokymo metu GPU atmintyje gali būti saugomi modelio parametrai, gradientai, optimizatoriaus būsena, įvesties tenzoriai, laikinos darbo sritys ir atgaliniam skaičiavimui išsaugotos aktyvacijos. PyTorch taip pat naudoja talpyklą naudojantį skirstytuvą, todėl atmintis, rodoma kaip „rezervuota“, nėra identiška atminčiai, kurią šiuo metu užima aktyvūs tenzoriai. Šis skirtumas yra svarbus sprendžiant, ar mažinti darbo krūvį, ar tirti skirstytuvo fragmentaciją.

Šiame vadove remiamasi dabartine PyTorch dokumentacija ir naudojami dabartiniai AMP API pavadinimai. Visų pirma, PyTorch dabar dokumentuoja torch.amp.autocast("cuda") ir torch.amp.GradScaler("cuda"); senesni torch.cuda.amp.* įėjimo taškai yra pasenę. Žr. PyTorch automatinio mišrios tikslumo dokumentaciją.

Greitas triažas: su kokio tipo OOM susiduriate?

1 žingsnis: Išmatuokite atmintį prieš keisdami mokymo receptą

Pirmiausia užfiksuokite partijos dydį, įvesties matmenis, tikslumą ir tašką, kuriame įvyksta klaida. Tada patikrinkite tiek aktyvių tenzorių atmintį, tiek skirstytuvo rezervuotą atmintį. PyTorch pateikia memory_allocated(), memory_reserved(), pikinių reikšmių variantus ir memory_summary(). Dabartinė CUDA atminties valdymo dokumentacija paaiškina, kad talpyklą naudojantis skirstytuvas išlaiko pakartotinai naudojamus blokus, todėl nepanaudota rezervuota atmintis GPU stebėjimo įrankiuose vis tiek gali atrodyti kaip panaudota. PyTorch CUDA atminties valdymas.

import torch

torch.cuda.reset_peak_memory_stats()

# Čia paleiskite vieną reprezentatyvų mokymo žingsnį.

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

Jei paprastos santraukos nepakanka, PyTorch gali fiksuoti skirstytuvo momentines nuotraukas gilesnei analizei. Jo atminties įrankiai gali įrašyti paskyrimų istoriją ir sukurti nuotrauką, kurią galima peržiūrėti naudojant PyTorch atminties vizualizatorių. PyTorch pažymi, kad šie įrankiai mato tik PyTorch skirstytuvo valdomą atmintį; tiesiogiai kitų CUDA bibliotekų atlikti paskyrimai ten gali būti nematomi. PyTorch vadovas suprasti CUDA atminties naudojimą.

Nelaikykite torch.cuda.empty_cache() visuotiniu OOM sprendimu

torch.cuda.empty_cache() atlaisvina nepanaudotus talpyklos blokus, kad juos galėtų naudoti kitos GPU programos. PyTorch aiškiai nurodo, kad ji neatleidžia atminties, kurią užima aktyvūs tenzoriai, todėl nedidina PyTorch prieinamos GPU atminties tenzoriams, kurie vis dar yra aktyvūs. Ji gali būti naudinga tarp atskirų eksperimentų arba ištrynus didelius objektus, tačiau tai nėra pakaitalas aktyvios atminties pėdsako mažinimui.

2 žingsnis: Pirmiausia sumažinkite aktyvų darbo rinkinį

Patikimiausias pirmasis sprendimas dažniausiai yra mažesnė mikro-partija: pavyzdžių skaičius, apdorotas vieno tiesioginio/atgalinio skaičiavimo metu. Aktyvacijų atmintis dažniausiai didėja didėjant partijos dydžiui, vaizdo skiriamajai gebai, sekos ilgiui ir kitiems įvesties matmenims. Jei modelis mokosi su 32 dydžio partija, bet nepavyksta su 64, partijos mažinimas nėra neigiamas kompromisas; tai yra tiesioginis pikinės atminties paklausos mažinimas.

Vaizdams erdvinės skiriamosios gebos arba apkarpyto vaizdo dydžio mažinimas gali turėti didelę įtaką. Transformeriams ir kitiems sekų modeliams sekos ilgio mažinimas gali būti dar svarbesnis, nes kai kurie tarpiniai tenzoriai stipriai didėja didėjant sekos ilgiui. Tikslus mastelio keitimas priklauso nuo architektūros, todėl matuokite, o ne spėliokite.

Taip pat įsitikinkite, kad vertinimo kodas be reikalo nekuria gradientų. PyTorch našumo gairės rekomenduoja išjungti gradientų skaičiavimą validacijai ar išvadai, kai gradientai nereikalingi, nes autogradas kitaip išsaugo tarpinius buferius. Tipiškas rašymo būdas yra:

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)

Mokymo metu naudokite optimizer.zero_grad(set_to_none=True), nebent jūsų algoritmas remiasi elgsenos skirtumu tarp nulinio gradiento ir None gradiento. PyTorch optimizatoriaus dokumentacija nurodo, kad gradientų nustatymas į None paprastai turi mažesnį atminties pėdsaką ir gali šiek tiek pagerinti našumą. PyTorch optimizatoriaus zero_grad dokumentacija.

3 žingsnis: Išlaikykite didesnę efektyvią partiją naudodami AMP ir gradientų kaupimą

Naudokite automatinį mišrų tikslumą, kai tai palaiko modelis

Automatinis mišrus tikslumas (AMP) vykdo tinkamas operacijas mažesniu tikslumu, o operacijas, kurioms reikia didesnio diapazono ar tikslumo, palieka tinkamuose tipuose. PyTorch dokumentuoja, kad AMP gali pagerinti našumą ir sumažinti atminties pėdsaką daugeliui CUDA darbo krūvių, tačiau jis nėra skaitmeniškai tinkamas kiekvienam modeliui. Visų pirma, PyTorch įspėja, kad kai kurie modeliai, iš anksto apmokyti naudojant bfloat16, gali patirti perpildymą naudojant float16.

Dabartinis CUDA AMP mokymo rašymo būdas yra:

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 rekomenduoja tiesioginį skaičiavimą ir nuostolį vykdyti autocast kontekste, o prieš atgalinį skaičiavimą palikti autocast kontekstą. Jei float16 sukelia nestabilumą, ištirkite, ar bfloat16 yra palaikomas ir tinkamas jūsų aparatūrai bei modeliui, vietoj to, kad manytumėte, jog visi mišraus tikslumo režimai veikia vienodai.

Naudokite gradientų kaupimą, kai reikia didesnės efektyvios partijos

Gradientų kaupimas apdoroja kelias mažesnes mikro-partijas prieš atnaujinant optimizatorių. Jei mikro-partija yra 4, o jūs kaupiate 8 žingsnius, efektyvi partija vienam optimizatoriaus atnaujinimui yra 32 pavyzdžiai vienam darbuotojui, darant prielaidą, kad kiekviena mikro-partija turi keturis pavyzdžius ir duomenų lygiagretumo sąranka nekeičia šios aritmetikos.

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)

Produkcijos kodui taip pat apdorokite galutinį dalinį kaupimo langą, kai partijų skaičius nesidalija iš accum_steps. Jei naudojate paskirstytąjį mokymąsi, gradientų sinchronizavimo elgsena gali pakeisti atminties/našumo kompromisą, todėl laikykitės paskirstytos API kaupimo gairių, o ne kopijuokite vieno GPU ciklą nepakeistą.

4 žingsnis: Keiskite skaičiavimą į atmintį, tada tirkite išlaikymą ir fragmentaciją

Aktyvacijų įrašymas (Checkpointing)

Aktyvacijų įrašymas sumažina atmintį neišlaikant pasirinktų tiesioginių aktyvacijų gyvų iki atgalinio skaičiavimo. Vietoj to, PyTorch jas perskaičiuoja atgalinio skaičiavimo metu. Tai keičia papildomą skaičiavimą į mažesnį aktyvacijų atminties pėdsaką. Dabartinė PyTorch checkpoint dokumentacija rekomenduoja aiškiai perduoti use_reentrant=False. PyTorch aktyvacijų įrašymo dokumentacija.

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)

Įrašykite sluoksnius, turinčius dideles išsaugotas aktyvacijas ir priimtiną perskaičiavimo kaštą. Nelaikykite, kad kiekvienos operacijos įrašymas yra optimalus; tai gali žymiai sulėtinti mokymą.

Ieškokite tenzorių, kurie išlaiko skaičiavimo grafus gyvus

Jei atmintis didėja kiekvieną iteraciją, o ne pasiekia maždaug tą patį lygį, patikrinkite Python nuorodas. Dažnas rašymo būdas yra saugoti su grafu susijusius tenzorius sąraše:

# Rizikinga, jei išlaikoma daug žingsnių:
loss_history.append(loss)

# Vietoj to saugokite Python skaičių:
loss_history.append(loss.item())

Ta pati problema gali kilti, kai talpykloje saugote modelio išvestis, dėmesio žemėlapius, paslėptas būsenas ar validacijos tenzorius neatjungdami jų arba neperkeldami iš GPU. Ištrinkite nuorodas, kurių jums nebereikia, ir naudokite detach() tik tada, kai sąmoningai norite tenzoriaus, atjungto nuo autogrado.

Derinkite skirstytuvą tik tada, kai statistika rodo fragmentaciją

Dabartinė PyTorch dokumentacija teikia pirmenybę aplinkos kintamajam PYTORCH_ALLOC_CONF. Senesnis PYTORCH_CUDA_ALLOC_CONF lieka aliasu atgaliniam suderinamumui. Ši pavadinimo detalė pasikeitė dabartinėje dokumentacijoje, todėl naujose konfigūracijose reikėtų naudoti pirmenybės teikiamą pavadinimą. PyTorch CUDA aplinkos kintamieji.

Ypač aktualūs du skirstytuvo parametrai:

• expandable_segments:True yra eksperimentinis ir sukurtas sumažinti nepanaudojamus atminties gabalus, kai keičiasi paskyrimų dydžiai, pvz., darbo krūviuose, kurių partijų ar tenzorių dydžiai kinta.
• max_split_size_mb gali sumažinti fragmentaciją naudojant natyvų skirstytuvą, tačiau PyTorch aiškiai apibūdina jį kaip paskutinę priemonę darbo krūviams, kurie patiria OOM klaidas, rodydami didelį neaktyvių padalintų blokų kiekį. Tai taip pat gali pabloginti našumą ir yra ignoruojamas cudaMallocAsync posistemio.

# Pavyzdys darbo krūviui su kintančiais paskyrimų dydžiais:
export PYTORCH_ALLOC_CONF=expandable_segments:True

Nekopijuokite skirstytuvo vėliavų iš kito kompiuterio nepatikrinę memory_summary() arba momentinės nuotraukos. Tikra talpos problema – kai aktyvūs tenzoriai jau užpildo GPU – nebus išspręsta fragmentacijos derinimu.

Kai vienas GPU vis tiek netelpa modelio

Jei vienas pavyzdys mažiausiu praktiniu įvesties dydžiu vis tiek sukelia OOM, problema gali būti modelis ir optimizatoriaus būsena, o ne partija. Tuo atveju apsvarstykite mažesnę architektūrą, mažesnio tikslumo parametrus, kur tai skaitmeniškai tinkama, CPU/iškrovimo strategijas arba padalintą paskirstytąjį mokymąsi.

PyTorch Fully Sharded Data Parallel (FSDP) gali padalinti modelio parametrus tarp duomenų lygiagretumo darbuotojų, o jo FULL_SHARD strategija taip pat padalina gradientus ir optimizatoriaus būsenas. Tai gali sumažinti atmintį vienam GPU, palyginti su visiškai dubliuotu duomenų lygiagretumu, tačiau tai kainuoja komunikaciją ir sudėtingesnę mokymo elgseną. PyTorch FSDP dokumentacija.

Praktinė veiksmų seka

Kontrolinis sąrašas: kaip žinoti, kad OOM tikrai išspręstas

• Paleiskite kelis reprezentatyvius mokymo iteracijas, o ne tik vieną sėkmingą tiesioginį skaičiavimą.
• Atstatykite ir užfiksuokite max_memory_allocated(), kad žinotumėte naują piką.
• Patvirtinkite, kad GPU atmintis pasiekia stabilų diapazoną, o ne didėja kiekvieną iteraciją.
• Patvirtinkite nuostolius ir gradientus įjungus mišrų tikslumą.
• Patvirtinkite, kad gradientų kaupimas išlaiko planuotą optimizatoriaus atnaujinimo tvarkaraštį.
• Paleiskite validacijos žingsnį naudojant torch.no_grad(), kai gradientai nereikalingi.
• Jei keitėte skirstytuvo nustatymus, palyginkite atminties statistiką ir našumą prieš ir po.
• Nelaikykite problemos išspręsta vien todėl, kad nvidia-smi rodo mažiau rezervuotos atminties po empty_cache(); pats mokymo darbo krūvis turi užbaigti savo įprastą piką.

CUDA OOM geriausia traktuoti kaip atminties biudžeto problemą, o ne kaip vieną PyTorch klaidą. Išmatuokite piką, pirmiausia sumažinkite gyvą darbo rinkinį, tada naudokite mišrų tikslumą, kaupimą ir įrašymą kaip sąmoningus kompromisus. Pereikite prie skirstytuvo derinimo tik tada, kai skirstytuvo statistika rodo fragmentaciją, ir pereikite prie padalinimo ar kito modelio, kai pats modelis nebetelpa viename GPU.