Jak opravit chybu „PyTorch CUDA Out of Memory“ během trénování modelu

Běh trénování v PyTorch může fungovat několik kroků a poté se zastavit s chybou torch.OutOfMemoryError nebo zprávou jako CUDA out of memory. Tried to allocate .... Bezprostřední příčina je jednoduchá: další alokace paměti CUDA se nevešla. Užitečná otázka zní proč se nevešla.

Během trénování může paměť GPU obsahovat parametry modelu, gradienty, stav optimalizátoru, vstupní tenzory, dočasné pracovní plochy a aktivace uložené pro zpětné šíření. PyTorch také používá alokátor s mezipamětí, takže paměť zobrazená jako „reserved“ (rezervovaná) není totožná s pamětí aktuálně obsazenou živými tenzory. Tento rozdíl je důležitý při rozhodování, zda zmenšit zátěž, nebo zkoumat fragmentaci alokátoru.

Tato příručka vychází z aktuální dokumentace PyTorch a používá názvy aktuálního API pro AMP. Konkrétně PyTorch nyní dokumentuje torch.amp.autocast("cuda") a torch.amp.GradScaler("cuda"); starší vstupní body torch.cuda.amp.* jsou zastaralé. Viz Dokumentace automatické smíšené přesnosti v PyTorch.

Rychlá triáž: s jakým typem OOM máte co do činění?

Krok 1: Změřte paměť před změnou receptury trénování

Začněte zaznamenáním velikosti dávky, rozměrů vstupu, přesnosti a bodu, kde dojde k selhání. Poté zkontrolujte paměť živých tenzorů i paměť rezervovanou alokátorem. PyTorch poskytuje funkce memory_allocated(), memory_reserved(), varianty pro špičkovou hodnotu a memory_summary(). Aktuální dokumentace správy paměti CUDA vysvětluje, že alokátor s mezipamětí udržuje znovu použitelné bloky, což je důvod, proč nevyužitá rezervovaná paměť může být stále zobrazena jako obsazená v nástrojích pro monitorování GPU. Správa paměti CUDA v PyTorch.

import torch

torch.cuda.reset_peak_memory_stats()

# Zde spusťte jeden reprezentativní trénovací 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))

Pokud jednoduchý souhrn nestačí, PyTorch může zachytit snímky alokátoru pro hlubší analýzu. Jeho nástroje pro práci s pamětí mohou zaznamenat historii alokací a vytvořit snímek, který lze prozkoumat pomocí vizualizéru paměti PyTorch. PyTorch upozorňuje, že tyto nástroje vidí paměť spravovanou alokátorem PyTorch; alokace provedené přímo jinými knihovnami CUDA se tam nemusí objevit. Průvodce PyTorch pochopením využití paměti CUDA.

Nepovažujte torch.cuda.empty_cache() za obecné řešení OOM

torch.cuda.empty_cache() uvolňuje nevyužité bloky v mezipaměti, aby je mohly použít jiné aplikace GPU. PyTorch výslovně uvádí, že neuvolňuje paměť obsazenou živými tenzory, a tedy nezvyšuje množství paměti GPU dostupné pro PyTorch pro tenzory, které jsou stále živé. Může být užitečné mezi samostatnými experimenty nebo po smazání velkých objektů, ale nenahrazuje zmenšení aktivní paměťové náročné stopy.

Krok 2: Nejprve zmenšete aktivní pracovní množinu

Nejspolehlivější první opravou je obvykle menší mikro-dávka: počet vzorků zpracovaných jedním dopředným/zpětným průchodem. Paměť pro aktivace obvykle roste s velikostí dávky, rozlišením obrázku, délkou sekvence a dalšími rozměry vstupu. Pokud se model trénuje při velikosti dávky 32, ale selhává při 64, zmenšení dávky není obcházením problému v negativním smyslu; je to přímé snížení špičkové paměťové náročné stopy.

U obrázků může snížení prostorového rozlišení nebo velikosti ořezu způsobit velký rozdíl. U transformérů a jiných sekvenčních modelů může být snížení délky sekvence ještě důležitější, protože některé mezivýsledné tenzory silně rostou s délkou sekvence. Přesné škálování závisí na architektuře, takže měřte, místo abyste předpokládali.

Také se ujistěte, že kód pro vyhodnocování nepostaví gradienty zbytečně. Doporučení pro výkon v PyTorch doporučují vypnout výpočet gradientů pro validaci nebo inferenci, když gradienty nejsou potřeba, protože autograd jinak ukládá mezivýsledné buffery. 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)

Během trénování používejte optimizer.zero_grad(set_to_none=True), pokud váš algoritmus nezávisí na behaviorálním rozdílu mezi nulovým gradientem a gradientem None. Dokumentace optimalizátoru PyTorch uvádí, že nastavení gradientů na None má obecně nižší paměťovou náročnost a může mírně zlepšit výkon. Dokumentace zero_grad pro optimalizátor PyTorch.

Krok 3: Udržujte větší efektivní dávku pomocí AMP a akumulace gradientů

Použijte Automatickou smíšenou přesnost, pokud to model podporuje

Automatická smíšená přesnost (AMP) spouští vhodné operace v nižší přesnosti, zatímco operace vyžadující větší rozsah nebo přesnost ponechává v příslušných typech. PyTorch dokumentuje, že AMP může zlepšit výkon a snížit paměťovou náročnost pro mnoho úloh CUDA, ale není numericky vhodná pro každý model. Konkrétně PyTorch varuje, že některé modely předtrénované v bfloat16 mohou přetékat v float16.

Aktuální vzor trénování AMP pro CUDA 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 doporučuje spustit dopředný průchod a výpočet ztráty pod autocast, a poté opustit kontext autocast před zpětným šířením. Pokud float16 způsobuje nestabilitu, zkontrolujte, zda je podporováno a vhodné bfloat16 pro váš hardware a model, místo abyste předpokládali, že všechny režimy smíšené přesnosti se chovají identicky.

Použijte akumulaci gradientů, když potřebujete větší efektivní dávku

Akumulace gradientů zpracovává několik menších mikro-dávek před aktualizací optimalizátoru. Pokud je mikro-dávka 4 a vy akumulujete 8 kroků, efektivní dávka pro jednu aktualizaci optimalizátoru je 32 vzorků na worker, za předpokladu, že každá mikro-dávka má čtyři vzorky a nastavení datové paralelizace tuto aritmetiku nemění.

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)

Pro produkční kód také ošetřete poslední částečné okno akumulace, když není počet dávek dělitelný accum_steps. Pokud používáte distribuované trénování, chování synchronizace gradientů může změnit kompromis mezi pamětí a výkonem, takže dodržujte doporučení pro akumulaci v distribučním API, místo abyste kopírovali smyčku pro jednu GPU beze změny.

Krok 4: Vyměňte výpočetní výkon za paměť, poté zkoumejte zadržování a fragmentaci

Ukládání aktivací (Activation checkpointing)

Ukládání aktivací (checkpointing) snižuje paměť tím, že nedrží vybrané dopředné aktivace živé až do zpětného šíření. Místo toho je PyTorch přepočítá během zpětného šíření. To vyměňuje dodatečný výpočetní výkon za nižší paměťovou náročnost aktivací. Aktuální dokumentace checkpointingu v PyTorch doporučuje explicitně předat use_reentrant=False. Dokumentace ukládání aktivací v 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)

Ukládejte vrstvy s velkými uloženými aktivacemi a přijatelnou náročností přepočtu. Nepředpokládejte, že ukládání každé operace je optimální; může výrazně zpomalit trénování.

Hledejte tenzory, které udržují výpočetní grafy živé

Pokud paměť roste každou iterací, místo aby kulminovala přibližně na stejné úrovni, zkontrolujte odkazy v Pythonu. Běžným vzorem je ukládání tenzorů spojených s grafem do seznamu:

# Rizikové, pokud je udržováno mnoho kroků:
loss_history.append(loss)

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

Stejný problém může nastat, když mezipaměťujete výstupy modelu, mapy pozornosti, skryté stavy nebo validační tenzory bez jejich odpojení (detach) nebo přesunu mimo GPU. Smažte odkazy, které již nepotřebujete, a použijte detach() pouze tehdy, když záměrně chcete tenzor odpojit od autogradu.

Laděte alokátor pouze tehdy, když statistiky ukazují na fragmentaci

Aktuální dokumentace PyTorch upřednostňuje proměnnou prostředí PYTORCH_ALLOC_CONF. Starší PYTORCH_CUDA_ALLOC_CONF zůstává aliasem pro zpětnou kompatibilitu. Tento detail pojmenování se v aktuální dokumentaci změnil, takže nové konfigurace by měly používat preferovaný název. Proměnné prostředí CUDA v PyTorch.

Dvě možnosti alokátoru jsou obzvláště relevantní:

• expandable_segments:True je experimentální a je navrženo tak, aby snižovalo nepoužitelné útržky paměti, když se velikosti alokací mění, například u úloh, jejichž velikosti dávek nebo tenzorů kolísají.
• max_split_size_mb může snížit fragmentaci s nativním alokátorem, ale PyTorch to výslovně popisuje jako poslední možnost pro úlohy, které selhávají s OOM, zatímco vykazují velké množství neaktivních rozdělených bloků. Může také poškodit výkon a je ignorováno backendem cudaMallocAsync.

# Příklad pro úlohu s měnícími se velikostmi alokací:
export PYTORCH_ALLOC_CONF=expandable_segments:True

Nekopírujte příznaky alokátoru z jiného stroje bez kontroly memory_summary() nebo snímku. Skutečný problém s kapacitou—kde živé tenzory již zaplňují GPU—nebude vyřešen laděním fragmentace.

Když jedna GPU stále nedokáže pojmout model

Pokud jeden vzorek při nejmenší praktické vstupní velikosti stále způsobuje OOM, problém může být v modelu a stavu optimalizátoru, nikoli v dávce. V takovém případě zvažte menší architekturu, parametry s nižší přesností, kde je to numericky vhodné, strategie CPU/offload nebo šardované distribuované trénování.

Fully Sharded Data Parallel (FSDP) v PyTorch může šardovat parametry modelu mezi workerů datové paralelizace a jeho strategie FULL_SHARD šarduje také gradienty a stavy optimalizátoru. To může snížit paměť na GPU ve srovnání s plně replikovanou datovou paralelizací, za cenu komunikace a složitějšího chování trénování. Dokumentace FSDP v PyTorch.

Praktické pořadí operací

Kontrolní seznam: jak poznat, že je OOM skutečně opraveno

• Spusťte několik reprezentativních trénovacích iterací, nejen jeden úspěšný dopředný průchod.
• Resetujte a zaznamenejte max_memory_allocated(), abyste znali novou špičku.
• Ověřte, že paměť GPU dosahuje stabilního rozsahu, místo aby rostla každou iterací.
• Ověřte ztrátu a gradienty po povolení smíšené přesnosti.
• Ověřte, že akumulace gradientů zachovává plán aktualizací optimalizátoru, který jste zamýšleli.
• Spusťte validační průchod pod torch.no_grad(), když gradienty nejsou vyžadovány.
• Pokud jste změnili nastavení alokátoru, porovnejte statistiky paměti a propustnost před a po.
• Nepovažujte problém za vyřešený pouze proto, že nvidia-smi ukazuje méně rezervované paměti po empty_cache(); samotná trénovací úloha musí dokončit svůj obvyklý špičkový stav.

OOM CUDA je nejlepší považovat za problém s rozpočtem paměti, nikoli za jednu chybu v PyTorch. Změřte špičku, nejprve zmenšete živou pracovní množinu, poté použijte smíšenou přesnost, akumulaci a ukládání aktivací jako záměrné kompromisy. Přejděte na ladění alokátoru pouze tehdy, když statistiky alokátoru naznačují fragmentaci, a přejděte na šardování nebo jiný model, když se model sám již pohodlně nevejde na jednu GPU.