Ako opraviť chybu „PyTorch CUDA Out of Memory“ počas trénovania modelu

Trénovacie behy v PyTorch môžu fungovať niekoľko krokov a potom sa zastaviť s chybou torch.OutOfMemoryError alebo správou ako CUDA out of memory. Tried to allocate .... Bezprostredná príčina je jednoduchá: ďalšia alokácia CUDA sa nezmestila do pamäte. Užitočnou otázkou je prečo sa nezmestila.

Počas trénovania môže pamäť GPU obsahovať parametre modelu, gradienty, stav optimalizátora, vstupné tenzory, dočasné pracovné priestory a aktivácie uložené pre spätný chod. PyTorch tiež používa alokátor s cache, takže pamäť zobrazená ako „reserved“ (rezervovaná) nie je totožná s pamäťou aktuálne obsadenou živými tenzormi. Tento rozdiel je dôležitý pri rozhodovaní, či znížiť záťaž, alebo vyšetriť fragmentáciu alokátoru.

Táto príručka sa riadi aktuálnou dokumentáciou PyTorch a používa súčasné názvy API pre AMP. Konkrétne, PyTorch teraz dokumentuje torch.amp.autocast("cuda") a torch.amp.GradScaler("cuda"); staršie vstupné body torch.cuda.amp.* sú zastarané. Pozri Dokumentácia automatickej zmiešanej presnosti PyTorch.

Rýchla triáž: s akým typom OOM sa stretávate?

Krok 1: Merajte pamäť pred zmenou receptu na trénovanie

Začnite zaznamenaním veľkosti dávky, rozmerov vstupu, presnosti a bodu, kde nastane zlyhanie. Potom skontrolujte pamäť živých tenzorov aj pamäť rezervovanú alokátorom. PyTorch poskytuje funkcie memory_allocated(), memory_reserved(), ich vrcholové varianty a memory_summary(). Súčasná dokumentácia správy pamäte CUDA vysvetľuje, že alokátor s cache udržiava znovupoužiteľné bloky, čo je dôvod, prečo sa nevyužitá rezervovaná pamäť môže stále javiť ako použitá v monitorovacích nástrojoch GPU. Správa pamäte CUDA v PyTorch.

import torch

torch.cuda.reset_peak_memory_stats()

# Tu spustite jeden reprezentatívny tréningový krok.

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

Ak jednoduchý súhrn nestačí, PyTorch môže zachytiť snímky alokátoru pre hlbšiu analýzu. Jeho nástroje na správu pamäte môžu zaznamenávať históriu alokácií a vytvárať snímku, ktorú možno skúmať pomocou vizualizátora pamäte PyTorch. PyTorch uvádza, že tieto nástroje vidia pamäť spravovanú alokátorom PyTorch; alokácie vykonané priamo inými knižnicami CUDA sa tam nemusia objaviť. Príručka PyTorch na pochopenie používania pamäte CUDA.

Nepovažujte torch.cuda.empty_cache() za všeobecné riešenie OOM

torch.cuda.empty_cache() uvoľňuje nevyužité cacheované bloky, aby ich mohli používať iné aplikácie GPU. PyTorch výslovne uvádza, že neuvoľňuje pamäť obsadenú živými tenzormi, a preto nezvyšuje množstvo pamäte GPU dostupnej pre PyTorch pre tenzory, ktoré sú stále živé. Môže byť užitočné medzi samostatnými experimentmi alebo po odstránení veľkých objektov, ale nenahrádza znižovanie aktívnej pamäťovej stopy.

Krok 2: Najprv znížte aktívnu pracovnú množinu

Najspoľahlivejšie prvé riešenie je zvyčajne menšia mikro-dávka: počet vzoriek spracovaných jedným dopredným/spätným prechodom. Pamäť aktivácií zvyčajne rastie s veľkosťou dávky, rozlíšením obrázka, dĺžkou sekvencie a ďalšími vstupnými rozmermi. Ak sa model trénuje pri veľkosti dávky 32, ale zlyhá pri 64, zníženie dávky nie je obchádzaním v negatívnom zmysle; je to priame zníženie dopytu po vrcholovej pamäti.

Pre obrázky môže zníženie priestorového rozlíšenia alebo veľkosti orezania spôsobiť veľký rozdiel. Pre transformery a iné sekvenčné modely môže byť zníženie dĺžky sekvencie ešte dôležitejšie, pretože niektoré medziprodukty silne rastú s dĺžkou sekvencie. Presné škálovanie závisí od architektúry, preto merajte, namiesto toho, aby ste predpokladali.

Tiež sa uistite, že kód pre vyhodnocovanie nebuduje gradienty zbytočne. Návod na výkon PyTorch odporúča vypnúť výpočet gradientov pre validáciu alebo inferenciu, keď nie sú potrebné, pretože autograd inak ukladá medziprodukty. Typický vzor 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)

Počas trénovania používajte optimizer.zero_grad(set_to_none=True), pokiaľ váš algoritmus nespoľahlivo závisí na behaviorálnom rozdieli medzi nulovým gradientom a gradientom None. Dokumentácia optimalizátora PyTorch uvádza, že nastavenie gradientov na None má zvyčajne nižšiu pamäťovú stopu a môže mierne zlepšiť výkon. Dokumentácia zero_grad optimalizátora PyTorch.

Krok 3: Udržujte väčšiu efektívnu dávku pomocou AMP a akumulácie gradientov

Používajte Automatic Mixed Precision, ak to model podporuje

Automatic Mixed Precision (AMP) spúšťa oprávnené operácie v nižšej presnosti, zatiaľ čo operácie vyžadujúce väčší rozsah alebo presnosť ponecháva v vhodných typoch. PyTorch dokumentuje, že AMP môže zlepšiť výkon a znížiť pamäťovú stopu pre mnohé CUDA úlohy, ale nie je numericky vhodné pre každý model. Konkrétne, PyTorch varuje, že niektoré modely predtrénované v bfloat16 môžu pretekať v float16.

Aktuálny vzor trénovania CUDA AMP 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 odporúča spustiť dopredný prechod a stratu pod autocast, potom opustiť kontext autocast pred spätným chodom. Ak float16 spôsobuje nestabilitu, preskúmajte, či je bfloat16 podporovaný a vhodný pre váš hardvér a model, namiesto predpokladu, že všetky režimy zmiešanej presnosti sa správajú identicky.

Používajte akumuláciu gradientov, keď potrebujete väčšiu efektívnu dávku

Akumulácia gradientov spracováva niekoľko menších mikro-dávok pred aktualizáciou optimalizátora. Ak je mikro-dávka 4 a akumulujete 8 krokov, efektívna dávka pre jednu aktualizáciu optimalizátora je 32 vzoriek na worker, za predpokladu, že každá mikro-dávka má štyri vzorky a nastavenie data-parallel nezmení túto 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)

Pre produkčný kód tiež ošetrite posledné čiastočné okno akumulácie, keď počet dávok nie je deliteľný accum_steps. Ak používate distribuované trénovanie, správanie synchronizácie gradientov môže zmeniť kompromis medzi pamäťou a výkonom, preto sa riaďte pokynmi pre akumuláciu distribuovaného API, namiesto kopírovania slučky pre jednu GPU bez zmien.

Krok 4: Vymeňte výpočty za pamäť, potom vyšetrite zadržiavanie a fragmentáciu

Checkpointing aktivácií

Checkpointing aktivácií znižuje pamäť tým, že nenecháva vybrané dopredné aktivácie živé až do spätného chodu. Namiesto toho ich PyTorch prepočíta počas spätného chodu. Toto vymieňa dodatočný výpočet za nižšiu pamäťovú stopu aktivácií. Súčasná dokumentácia checkpointingu PyTorch odporúča explicitne odovzdať use_reentrant=False. Dokumentácia checkpointingu aktivácií PyTorch.

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)

Checkpointujte vrstvy s veľkými uloženými aktiváciami a prijateľnou cenou prepočtu. Nepredpokladajte, že checkpointovanie každej operácie je optimálne; môže výrazne spomaliť trénovanie.

Hľadajte tenzory, ktoré držia výpočtové grafy živé

Ak pamäť rastie každú iteráciu namiesto toho, aby vrcholila približne na rovnakej úrovni, skontrolujte odkazy v Pythone. Bežným vzorom je ukladanie tenzorov spojených s grafom do zoznamu:

# Rizikové, ak sa udržiava mnoho krokov:
loss_history.append(loss)

# Uložte namiesto toho číslo Pythonu:
loss_history.append(loss.item())

Rovnaký problém môže nastať, keď cacheujete výstupy modelu, mapy pozornosti, skryté stavy alebo validačné tenzory bez ich odpojenia alebo presunu mimo GPU. Odstráňte odkazy, ktoré už nepotrebujete, a použite detach() len vtedy, keď zámerne chcete tenzor odpojený od autogradu.

Ladenie alokátoru len po tom, čo štatistiky poukazujú na fragmentáciu

Súčasná dokumentácia PyTorch uprednostňuje premennú prostredia PYTORCH_ALLOC_CONF. Staršia PYTORCH_CUDA_ALLOC_CONF zostáva aliasom pre spätnú kompatibilitu. Tento detail pomenovania sa zmenil v aktuálnych dokumentoch, takže nové konfigurácie by mali používať uprednostňovaný názov. Premenné prostredia CUDA v PyTorch.

Dve možnosti alokátoru sú obzvlášť relevantné:

• expandable_segments:True je experimentálne a je navrhnuté tak, aby znižovalo nepoužiteľné kúsky pamäte, keď sa veľkosti alokácií menia, napríklad pri úlohách, kde sa veľkosti dávok alebo tenzorov líšia.
• max_split_size_mb môže znížiť fragmentáciu s natívnym alokátorom, ale PyTorch ho výslovne opisuje ako poslednú možnosť pre úlohy, ktoré zlyhávajú s OOM, zatiaľ čo ukazujú veľké množstvo neaktívnych rozdelených blokov. Môže tiež zhoršiť výkon a je ignorovaný backendom cudaMallocAsync.

# Príklad pre úlohu s rôznymi veľkosťami alokácií:
export PYTORCH_ALLOC_CONF=expandable_segments:True

Nekopírujte príznaky alokátoru z iného stroja bez kontroly memory_summary() alebo snímky. Skutočný problém s kapacitou – kde živé tenzory už vypĺňajú GPU – nebude vyriešený ladením fragmentácie.

Kedy jedna GPU stále nezmestí model

Ak jedna vzorka pri najmenšej praktickej vstupnej veľkosti stále spôsobuje OOM, problémom môže byť model a stav optimalizátora, nie dávka. V tom bode zvážte menšiu architektúru, parametre s nižšou presnosťou tam, kde je to numericky vhodné, stratégie CPU/offload alebo šardované distribuované trénovanie.

Fully Sharded Data Parallel (FSDP) v PyTorch môže šardovať parametre modelu medzi workerami data-parallel a jeho stratégia FULL_SHARD tiež šarduje gradienty a stavy optimalizátorov. To môže znížiť pamäť na GPU v porovnaní s plne replikovanou data parallelism, za cenu komunikácie a zložitejšieho správania trénovania. Dokumentácia PyTorch FSDP.

Praktické poradie operácií

Checklist: ako zistiť, že OOM je skutočne opravené

• Spustite niekoľko reprezentatívnych tréningových iterácií, nie len jeden úspešný dopredný prechod.
• Resetujte a zaznamenajte max_memory_allocated(), aby ste poznali nový vrchol.
• Potvrďte, že pamäť GPU dosahuje stabilný rozsah namiesto toho, aby rástla každú iteráciu.
• Validujte stratu a gradienty po povolení zmiešanej presnosti.
• Potvrďte, že akumulácia gradientov zachováva plán aktualizácie optimalizátora, ktorý ste zamýšľali.
• Spustite validačný prechod pod torch.no_grad(), keď nie sú vyžadované gradienty.
• Ak ste zmenili nastavenia alokátoru, porovnajte štatistiky pamäte a priepustnosť pred a po.
• Nepovažujte problém za vyriešený len preto, že nvidia-smi ukazuje menej rezervovanej pamäte po empty_cache(); samotná tréningová úloha sa musí dokončiť pri svojom normálnom vrchole.

OOM CUDA je najlepšie považovať za problém s rozpočtom pamäte, nie za jedinú chybu PyTorch. Merajte vrchol, najprv znížte živú pracovnú množinu, potom použite zmiešanú presnosť, akumuláciu a checkpointing ako zámerné kompromisy. Prejdite na ladenie alokátoru len vtedy, keď štatistiky alokátoru naznačujú fragmentáciu, a prejdite na šardovanie alebo iný model, keď sa model sám už pohodlne nezmestí na jednu GPU.