Kako riješiti grešku “PyTorch CUDA Out of Memory” tijekom treniranja modela

Treniranje u PyTorchu može raditi nekoliko koraka, a zatim se zaustaviti s greškom torch.OutOfMemoryError ili porukom poput CUDA out of memory. Tried to allocate .... Trenutni uzrok je jednostavan: sljedeća CUDA alokacija nije mogla biti smještena. Korisno pitanje je zašto nije mogla biti smještena.

Tijekom treniranja, GPU memorija može sadržavati parametre modela, gradijente, stanje optimizatora, ulazne tenzore, privremene radne prostore i aktivacije spremljene za backward pass. PyTorch također koristi caching alocator, pa memorija prikazana kao “reserved” (rezervirana) nije identična memoriji koju trenutno zauzimaju živi tenzori. Ta razlika je važna pri odlučivanju treba li smanjiti opterećenje ili istražiti fragmentaciju alocirača.

Ovaj vodič prati trenutnu PyTorch dokumentaciju i koristi trenutne nazive AMP API-ja. Posebno, PyTorch sada dokumentira torch.amp.autocast("cuda") i torch.amp.GradScaler("cuda"); starije torch.cuda.amp.* ulazne točke su zastarjele. Pogledajte PyTorch Automatic Mixed Precision dokumentaciju.

Brza dijagnostika: s kojom vrstom OOM-a imate posla?

Korak 1: Mjerite memoriju prije mijenjanja recepta za treniranje

Počnite bilježenjem veličine batcha, dimenzija ulaza, preciznosti i točke u kojoj se događa kvar. Zatim provjerite memoriju živih tenzora i memoriju rezerviranu od strane alocirača. PyTorch izlaže memory_allocated(), memory_reserved(), vršne varijante i memory_summary(). Trenutna dokumentacija o upravljanju CUDA memorijom objašnjava da caching alocator čuva ponovno iskoristive blokove, zbog čega neiskorištena rezervirana memorija i dalje može izgledati kao zauzeta u alatima za praćenje GPU-a. PyTorch CUDA upravljanje memorijom.

import torch

torch.cuda.reset_peak_memory_stats()

# Ovdje pokrenite jedan reprezentativni korak treniranja.

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

Ako jednostavan sažetak nije dovoljan, PyTorch može snimiti snimke alocirača za dublju analizu. Njegovi alati za memoriju mogu zabilježiti povijest alokacija i proizvesti snimku koja se može pregledati pomoću PyTorch vizualizatora memorije. PyTorch napominje da ovi alati vide memoriju kojom upravlja PyTorch alocator; alokacije koje izravno vrše druge CUDA biblioteke možda se neće pojaviti tamo. PyTorch vodič za razumijevanje korištenja CUDA memorije.

Nemojte tretirati torch.cuda.empty_cache() kao opće rješenje za OOM

torch.cuda.empty_cache() oslobađa neiskorištene predmemorirane blokove kako bi ih druge GPU aplikacije mogle koristiti. PyTorch izričito navodi da ne oslobađa memoriju koju zauzimaju živi tenzori i stoga ne povećava količinu GPU memorije dostupnu PyTorchu za tenzore koji su još uvijek živi. Može biti koristan između odvojenih eksperimenata ili nakon brisanja velikih objekata, ali nije zamjena za smanjenje aktivnog otiska memorije.

Korak 2: Prvo smanjite aktivni radni skup

Najpouzdanije prvo rješenje obično je manji mikro-batch: broj uzoraka obrađenih u jednom forward/backward prolazu. Memorija aktivacija obično raste s veličinom batcha, rezolucijom slike, duljinom niza i drugim dimenzijama ulaza. Ako se model trenira s batch veličinom 32, ali ne uspijeva s 64, smanjenje batcha nije zaobilazno rješenje u negativnom smislu; to je izravno smanjenje potražnje za vršnom memorijom.

Kod slika, smanjenje prostorne rezolucije ili veličine cropa može napraviti veliku razliku. Za transformere i druge modele sekvenci, smanjenje duljine sekvence može biti još važnije jer neki međutenzori jako rastu s duljinom sekvence. Točno skaliranje ovisi o arhitekturi, pa mjerite umjesto da pretpostavljate.

Također provjerite ne gradi kôd za evaluaciju nepotrebno. PyTorchove smjernice za performanse preporučuju onemogućavanje izračuna gradijenata za validaciju ili inferenciju kada gradijenti nisu potrebni, jer autograd inače sprema međuspremničke buffere. Tipičan obrazac je:

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)

Tijekom treniranja koristite optimizer.zero_grad(set_to_none=True) osim ako vaš algoritam ne ovisi o razlici u ponašanju između nultog gradijenta i None gradijenta. PyTorchova dokumentacija optimizatora navodi da postavljanje gradijenata na None općenito ima manji memorijski otisak i može umjereno poboljšati performanse. PyTorch optimizer zero_grad dokumentacija.

Korak 3: Zadržite veću efektivnu batch veličinu uz AMP i akumulaciju gradijenata

Koristite Automatic Mixed Precision kada model to podržava

Automatic Mixed Precision (AMP) pokreće kvalificirane operacije u nižoj preciznosti, dok operacije koje zahtijevaju veći raspon ili preciznost zadržava u odgovarajućim tipovima. PyTorch dokumentira da AMP može poboljšati performanse i smanjiti memorijski otisak za mnoga CUDA opterećenja, ali nije numerički prikladan za svaki model. Posebno, PyTorch upozorava da neki modeli pred-trenirani u bfloat16 mogu doživjeti overflow u float16.

Trenutni CUDA AMP obrazac za treniranje je:

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 preporučuje pokretanje forward passa i gubitka pod autocastom, a zatim napuštanje autocast konteksta prije backwarda. Ako float16 uzrokuje nestabilnost, istražite je li bfloat16 podržan i prikladan za vaš hardver i model, umjesto da pretpostavljate da se svi mješoviti precizijski načini ponašaju identično.

Koristite akumulaciju gradijenata kada vam je potrebna veća efektivna batch veličina

Akumulacija gradijenata obrađuje nekoliko manjih mikro-batcheva prije ažuriranja optimizatora. Ako je mikro-batch 4, a akumulirate 8 koraka, efektivni batch za jedno ažuriranje optimizatora je 32 uzorka po radniku, pod pretpostavkom da svaki mikro-batch ima četiri uzorka i da postavka data-parallel ne mijenja tu aritmetiku.

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)

Za produkcijski kôd, također obradite konačni djelomični prozor akumulacije kada broj batcheva nije djeljiv s accum_steps. Ako koristite distribuirano treniranje, ponašanje sinkronizacije gradijenata može promijeniti omjer memorije i performansi, pa slijedite smjernice za akumulaciju distribuiranog API-ja umjesto da kopirate petlju za jedan GPU nepromijenjenu.

Korak 4: Razmijenite računanje za memoriju, zatim istražite zadržavanje i fragmentaciju

Checkpointing aktivacija

Checkpointing aktivacija smanjuje memoriju tako što ne čuva odabrane forward aktivacije živima do backwarda. Umjesto toga, PyTorch ih ponovno izračunava tijekom backwarda. To razmjenjuje dodatno računanje za niži memorijski otisak aktivacija. Trenutna PyTorch checkpoint dokumentacija preporučuje izričito prosljeđivanje use_reentrant=False. PyTorch checkpointing aktivacija 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)

Checkpointirajte slojeve s velikim spremljenim aktivacijama i prihvatljivom cijenom ponovnog izračuna. Nemojte pretpostavljati da je checkpointiranje svake operacije optimalno; to može značajno usporiti treniranje.

Potražite tenzore koji održavaju računske grafove živima

Ako memorija raste pri svakoj iteraciji umjesto da doseže vrh na otprilike istoj razini, provjerite Python reference. Uobičajeni obrazac je spremanje tenzora povezanih s grafom u listu:

# Rizično ako se zadržava mnogo koraka:
loss_history.append(loss)

# Spremite Python broj umjesto toga:
loss_history.append(loss.item())

Isti problem se može dogoditi kada predmemorirate izlaze modela, attention mape, skrivena stanja ili tenzore validacije bez odvajanja ili premještanja s GPU-a. Obrišite reference koje vam više nisu potrebne i koristite detach() samo kada namjerno želite tenzor odvojen od autograda.

Prilagođavajte alocator samo nakon što statistika ukazuje na fragmentaciju

Trenutna PyTorch dokumentacija preferira varijablu okruženja PYTORCH_ALLOC_CONF. Starija PYTORCH_CUDA_ALLOC_CONF ostaje alias za povratnu kompatibilnost. Ova pojedinost o imenovanju promijenjena je u trenutnoj dokumentaciji, pa nove konfiguracije trebaju koristiti preferirani naziv. PyTorch CUDA varijable okruženja.

Dvije opcije alocirača su posebno relevantne:

• expandable_segments:True je eksperimentalna i dizajnirana je za smanjenje neiskoristivih komadića memorije kada se veličine alokacija mijenjaju, kao kod opterećenja čije se veličine batcha ili tenzora mijenjaju.
• max_split_size_mb može smanjiti fragmentaciju s nativnim alociračem, ali PyTorch ga izričito opisuje kao posljednje utočište za opterećenja koja ne uspijevaju s OOM-om dok pokazuju veliku količinu neaktivnih split blokova. Također može naškoditi performansama i ignorira ga cudaMallocAsync backend.

# Primjer za opterećenje s promjenjivim veličinama alokacija:
export PYTORCH_ALLOC_CONF=expandable_segments:True

Nemojte kopirati zastavice alocirača s drugog računala bez provjere memory_summary() ili snimke. Pravi problem kapaciteta—gdje živi tenzori već ispunjavaju GPU—neće biti riješen podešavanjem fragmentacije.

Kada jedan GPU i dalje ne može smjestiti model

Ako jedan uzorak pri najmanjoj praktičnoj ulaznoj veličini i dalje izaziva OOM, problem može biti model i stanje optimizatora, a ne batch. U tom trenutku razmislite o manjoj arhitekturi, parametrima niže preciznosti gdje je numerički prikladno, CPU/offload strategijama ili shardiranom distribuiranom treniranju.

PyTorchov Fully Sharded Data Parallel (FSDP) može shardirati parametre modela preko data-parallel radnika, a njegova FULL_SHARD strategija također shardira gradijente i stanja optimizatora. To može smanjiti memoriju po GPU-u u usporedbi s potpuno repliciranom data paralelizmom, uz cijenu komunikacije i složenijeg ponašanja treniranja. PyTorch FSDP dokumentacija.

Praktičan redoslijed operacija

Popis za provjeru: kako znati da je OOM stvarno riješen

• Pokrenite nekoliko reprezentativnih iteracija treniranja, a ne samo jedan uspješan forward pass.
• Resetirajte i zabilježite max_memory_allocated() kako biste znali novi vrh.
• Potvrdite da GPU memorija doseže stabilan raspon umjesto da raste pri svakoj iteraciji.
• Validirajte gubitak i gradijente nakon omogućavanja mješovite preciznosti.
• Potvrdite da akumulacija gradijenata čuva raspored ažuriranja optimizatora koji ste namjeravali.
• Pokrenite validacijski prolaz pod torch.no_grad() kada gradijenti nisu potrebni.
• Ako ste promijenili postavke alocirača, usporedite statistiku memorije i propusnost prije i poslije.
• Nemojte smatrati problem riješenim samo zato što nvidia-smi pokazuje manje rezervirane memorije nakon empty_cache(); samo opterećenje treniranja mora se dovršiti pri svom normalnom vrhu.

CUDA OOM je najbolje tretirati kao problem memorijskog proračuna, a ne kao jedan PyTorch bug. Izmjerite vrh, prvo smanjite živi radni skup, zatim koristite mješovitu preciznost, akumulaciju i checkpointing kao namjerne kompromise. Pređite na podešavanje alocirača samo kada statistika alocirača ukazuje na fragmentaciju, a na shardiranje ili drugi model kada model sam više ne staje udobno na jedan GPU.