Hogyan javítsd a „PyTorch CUDA Out of Memory” hibát modell tanítás közben

Egy PyTorch tanítási futás működhet néhány lépésig, majd megállhat a torch.OutOfMemoryError hibaüzenettel vagy egy olyan üzenettel, mint a CUDA out of memory. Tried to allocate .... A közvetlen ok egyszerű: a következő CUDA allokáció nem fért el. A hasznos kérdés az, hogy miért nem fért el.

Tanítás közben a GPU memória tárolhatja a modell paramétereit, a gradienseket, az optimalizáló állapotát, a bemeneti tenzorokat, az ideiglenes munkaterületeket és a backpropagációhoz mentett aktivációkat. A PyTorch cache-elt allokatort is használ, így a „lelőfoglalt” (reserved) memóriaként megjelenő érték nem azonos az élő tenzorok által jelenleg elfoglalt memóriával. Ez a különbség fontos, amikor eldöntjük, hogy csökkentsük-e a munkaterhelést, vagy vizsgáljuk meg az allokátor fragmentációját.

Ez az útmutató a jelenlegi PyTorch dokumentációt követi, és a jelenlegi AMP API neveket használja. Különösen fontos, hogy a PyTorch most már a torch.amp.autocast("cuda") és a torch.amp.GradScaler("cuda") függvényeket dokumentálja; a régebbi torch.cuda.amp.* belépési pontok elavultak. Lásd: PyTorch Automatic Mixed Precision dokumentáció.

Gyors triázs: milyen típusú OOM hibával állsz szemben?

1. lépés: Mérd meg a memóriát, mielőtt megváltoztatnád a tanítási receptet

Kezdd a batch méret, a bemeneti dimenziók, a pontosság és a hiba bekövetkezési pontjának rögzítésével. Ezután vizsgáld meg az élő tenzorok memóriáját és az allokátor által lefoglalt memóriát is. A PyTorch elérhetővé teszi a memory_allocated(), memory_reserved(), a csúcsérték-változatokat és a memory_summary() függvényt. A jelenlegi CUDA memóriakezelési dokumentáció elmagyarázza, hogy a cache-elt allokátor újrafelhasználható blokkokat tart fenn, ezért a nem használt lefoglalt memória még mindig használtként jelenhet meg a GPU figyelőeszközökben. PyTorch CUDA memóriakezelés.

import torch

torch.cuda.reset_peak_memory_stats()

# Futtass itt egy reprezentatív tanítási lépést.

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

Ha egy egyszerű összefoglaló nem elég, a PyTorch képes allokációs pillanatképeket rögzíteni mélyebb elemzéshez. Memóriaeszközei rögzíthetik az allokációs előzményeket, és olyan pillanatképet készíthetnek, amely megvizsgálható a PyTorch memória vizualizálóval. A PyTorch megjegyzi, hogy ezek az eszközök a PyTorch allokátor által kezelt memóriát látják; más CUDA könyvtárak által közvetlenül végzett allokációk nem feltétlenül jelennek meg ott. PyTorch útmutató a CUDA memóriahasználat megértéséhez.

Ne kezeld a torch.cuda.empty_cache() függvényt általános OOM megoldásként

A torch.cuda.empty_cache() felszabadítja a használaton kívüli cache-elt blokkokat, hogy más GPU alkalmazások használhassák őket. A PyTorch kifejezetten kimondja, hogy ez nem szabadítja fel az élő tenzorok által elfoglalt memóriát, és ezért nem növeli a PyTorch számára elérhető GPU memória mennyiségét még élő tenzorok esetén. Hasznos lehet különálló kísérletek között vagy nagy objektumok törlése után, de nem helyettesíti az aktív memória-lábnyom csökkentését.

2. lépés: Csökkentsd először az aktív munkakészletet

A legmegbízhatóbb első javítás általában egy kisebb micro-batch: az egy forward/backward pass során feldolgozott minták száma. Az aktivációs memória általában nő a batch mérettel, a felbontással, a szekvencia hosszával és egyéb bemeneti dimenziókkal. Ha a modell 32-es batch mérettel tanítható, de 64-nél elbukik, a batch csökkentése nem negatív értelemben vett kerülő megoldás; ez a csúcsmemória-igény közvetlen csökkentése.

Képek esetén a térbeli felbontás vagy a crop méret csökkentése nagy különbséget jelenthet. Transformer modellek és egyéb szekvencia modellek esetén a szekvencia hosszának csökkentése még fontosabb lehet, mivel egyes közbenső tenzorok erősen nőnek a szekvencia hosszával. A pontos skálázás az architektúrától függ, ezért mérj, ne feltételezz.

Ügyelj arra is, hogy az értékelési kód ne építsen feleslegesen gradienseket. A PyTorch teljesítményi útmutatója javasolja a gradiensszámítás letiltását validáció vagy inferencia során, amikor a gradiensekre nincs szükség, mert az autograd egyébként közbenső puffermentéseket végez. Egy tipikus minta:

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)

Tanítás közben használj optimizer.zero_grad(set_to_none=True) parancsot, hacsak az algoritmusod nem támaszkodik a zérus gradiens és a None gradiens viselkedésbeli különbségére. A PyTorch optimalizáló dokumentációja kimondja, hogy a gradiensek None-ra állítása általában kisebb memória-lábnyommal jár, és mérsékelten javíthatja a teljesítményt. PyTorch optimalizáló zero_grad dokumentáció.

3. lépés: Tarts fenn nagyobb effektív batch méretet AMP és gradiens-akumuláció segítségével

Használj Automatic Mixed Precision-t, ha a modell támogatja

Az Automatic Mixed Precision (AMP) az alkalmas műveleteket alacsonyabb pontossággal futtatja, miközben a nagyobb tartományt vagy pontosságot igénylő műveleteket megfelelő típusokban tartja. A PyTorch dokumentációja szerint az AMP javíthatja a teljesítményt és csökkentheti a memória-lábnyomot sok CUDA munkaterhelés esetén, de nem numerikusan alkalmas minden modellhez. Különösen figyelmeztet a PyTorch, hogy egyes bfloat16-ban előtanított modellek túlcsordulhatnak float16-ban.

Egy jelenlegi CUDA AMP tanítási minta:

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

A PyTorch azt javasolja, hogy a forward pass-t és a loss számítást az autocast kontextusban futtasd, majd hagyd el az autocast kontextust a backward előtt. Ha a float16 instabilitást okoz, vizsgálja meg, hogy a bfloat16 támogatott-e és megfelelő-e a hardveredhez és modelledhez, ahelyett, hogy feltételeznéd, hogy minden vegyes pontosságú mód azonosan viselkedik.

Használj gradiens-akumulációt, ha nagyobb effektív batch méretre van szükséged

A gradiens-akumuláció több kisebb micro-batch-et dolgoz fel az optimalizáló frissítése előtt. Ha a micro-batch mérete 4, és 8 lépést akkumulálsz, az effektív batch méret egy optimalizáló frissítéshez 32 minta munkafolyamat-onként, feltételezve, hogy minden micro-batch négy mintát tartalmaz, és az adat-párhuzamos beállítás nem változtat ezen a számításon.

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)

Éles kódban kezelj egy végső részleges akkumulációs ablakot is, amikor a batchek száma nem osztható accum_steps-szel. Ha elosztott tanítást használsz, a gradiens-szinkronizáció viselkedése megváltoztathatja a memória/teljesítmény kompromisszumot, ezért kövesd az elosztott API akkumulációs útmutatását, ahelyett, hogy egy GPU-s hurkot változtatás nélkül másolnál.

4. lépés: Cseréld a számítást memóriára, majd vizsgálj meg retention-t és fragmentációt

Aktiváció-checkpointing

Az aktiváció-checkpointing csökkenti a memóriát azáltal, hogy nem tartja életben a kiválasztott forward aktivációkat a backward-ig. Ehelyett a PyTorch újraszámítja őket a backward során. Ez további számítást cserél el alacsonyabb aktivációs memória-lábnyomért. A jelenlegi PyTorch checkpoint dokumentációja kifejezetten javasolja a use_reentrant=False átadását. PyTorch aktiváció-checkpointing dokumentáció.

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)

Checkpointold a rétegeket nagy mentett aktivációkkal és elfogadható újraszámítási költséggel. Ne feltételezd, hogy minden művelet checkpointolása optimális; ez jelentősen lelassíthatja a tanítást.

Keress olyan tenzorokat, amelyek életben tartják a számításgráfokat

Ha a memória minden iterációval nő, ahelyett, hogy nagyjából azonos szinten érne el csúcsot, vizsgálj meg Python hivatkozásokat. Egy gyakori minta a gráfhoz kapcsolódó tenzorok listában tárolása:

# Kockázatos, ha sok lépésig megmarad:
loss_history.append(loss)

# Tárolj helyette egy Python számot:
loss_history.append(loss.item())

Ugyanez a probléma előfordulhat, ha modell kimeneteket, figyelem térképeket, rejtett állapotokat vagy validációs tenzorokat cache-elsz anélkül, hogy leválasztanád (detach) vagy lemozgatnád a GPU-ról. Töröld a már nem szükséges hivatkozásokat, és csak akkor használd a detach()-et, ha szándékosan egy autograd-tól leválasztott tenzort szeretnél.

Állítsd be az allokatort csak akkor, ha a statisztikák fragmentációra utalnak

A jelenlegi PyTorch dokumentáció az PYTORCH_ALLOC_CONF környezeti változót részesíti előnyben. A régebbi PYTORCH_CUDA_ALLOC_CONF továbbra is alias a visszafelé kompatibilitás érdekében. Ez az elnevezési részlet megváltozott a jelenlegi dokumentációban, ezért az új konfigurációknak a preferált nevet kell használniuk. PyTorch CUDA környezeti változók.

Két allokációs opció különösen releváns:

• Az expandable_segments:True kísérleti jellegű, és arra tervezték, hogy csökkentse a használhatatlan memória-szeleteket, amikor az allokációs méretek változnak, például olyan munkaterhelések esetén, ahol a batch vagy tenzor méretek változnak.
• A max_split_size_mb csökkentheti a fragmentációt a natív allokatort használva, de a PyTorch kifejezetten végső megoldásként (last resort) írja le olyan munkaterhelések esetén, amelyek OOM hibával buknak el, miközben nagy mennyiségű inaktív split blokkot mutatnak. Ronthatja a teljesítményt is, és a cudaMallocAsync backend figyelmen kívül hagyja.

# Példa változó allokációs méretű munkaterheléshez:
export PYTORCH_ALLOC_CONF=expandable_segments:True

Ne másolj allokációs jelzőket egy másik gépről anélkül, hogy ellenőriznéd a memory_summary() kimenetét vagy egy pillanatképet. Egy valódi kapacitási probléma – ahol az élő tenzorok már kitöltik a GPU-t – nem oldható meg fragmentáció hangolással.

Amikor egyetlen GPU még mindig nem fér el a modell

Ha egyetlen minta a legkisebb gyakorlati bemeneti méret mellett is OOM hibát okoz, a probléma lehet a modell és az optimalizáló állapota, nem a batch. Ezen a ponton fontolj meg egy kisebb architektúrát, alacsonyabb pontosságú paramétereket, ahol numerikusan megfelelő, CPU/offload stratégiákat vagy sharded elosztott tanítást.

A PyTorch Fully Sharded Data Parallel (FSDP) képes a modell paramétereit shardingolni az adat-párhuzamos munkafolyamatok között, és a FULL_SHARD stratégia a gradienseket és az optimalizáló állapotokat is shardingolja. Ez csökkentheti a GPU-onkénti memóriát a teljesen replikált adat-párhuzamossághoz képest, a kommunikáció és a bonyolultabb tanítási viselkedés árán. PyTorch FSDP dokumentáció.

Gyakorlati műveleti sorrend

Ellenőrzőlista: honnan tudod, hogy az OOM tényleg megjavult

• Futtass le több reprezentatív tanítási iterációt, ne csak egy sikeres forward pass-t.
• Állítsd vissza és rögzítsd a max_memory_allocated() értéket, hogy tudd az új csúcsértéket.
• Erősítsd meg, hogy a GPU memória stabil tartományba kerül, ahelyett, hogy minden iterációval nőne.
• Validáld a loss-t és a gradienseket a vegyes pontosság engedélyezése után.
• Erősítsd meg, hogy a gradiens-akumuláció megőrzi a tervezett optimalizáló-frissítési ütemtervet.
• Futtass egy validációs pass-t torch.no_grad() alatt, amikor a gradiensekre nincs szükség.
• Ha megváltoztattad az allokátor beállításait, hasonlítsd össze a memória statisztikákat és az áteresztőképességet előtte és utána.
• Ne tekintsd megoldottnak a problémát csak azért, mert az nvidia-smi kevesebb lefoglalt memóriát mutat az empty_cache() után; a tanítási munkaterhelésnek magának kell befejeződnie a normál csúcsértékén.

A CUDA OOM hibát legjobban memória-költségvetési problémaként kell kezelni, nem egyetlen PyTorch hibaként. Mérd meg a csúcsértéket, csökkentsd először az élő munkakészletet, majd használd a vegyes pontosságot, az akkumulációt és a checkpointingot tudatos kompromisszumokként. Csak akkor térj át az allokátor hangolására, ha az allokátor statisztikái fragmentációra utalnak, és csak akkor térj át shardingra vagy egy másik modellre, ha a modell maga már nem fér el kényelmesen egyetlen GPU-n.