Kako odpraviti napako “PyTorch CUDA Out of Memory” med usposabljanjem modela

Usposabljanje modela v PyTorchu lahko deluje nekaj korakov, nato pa se ustavi z napako torch.OutOfMemoryError ali sporočilom, kot je CUDA out of memory. Tried to allocate .... Neposredni vzrok je preprost: naslednja dodelitev pomnilnika CUDA ni bila mogoča. Koristno vprašanje je zakaj ni bila mogoča.

Med usposabljanjem lahko pomnilnik GPU vsebuje parametre modela, gradienti, stanje optimizatorja, vhodni tenzorji, začasni delovni prostori in aktivacije, shranjene za nazajno propagacijo. PyTorch uporablja tudi predpomnilniški dodeljevalnik pomnilnika, zato pomnilnik, prikazan kot “rezerviran”, ni enak pomnilniku, ki ga trenutno zasedajo aktivni tenzorji. Ta razlika je pomembna pri odločanju, ali zmanjšati obremenitev ali preiskati fragmentacijo dodeljevalnika.

Ta vodnik sledi trenutni dokumentaciji PyTorch in uporablja trenutna imena API AMP. Zlasti PyTorch zdaj dokumentira torch.amp.autocast("cuda") in torch.amp.GradScaler("cuda"); starejše vstopne točke torch.cuda.amp.* so zastarele. Glejte dokumentacijo PyTorch Automatic Mixed Precision.

Hitra triaža: s kakšno vrsto OOM se soočate?

Korak 1: Izmerite pomnilnik, preden spremenite recept za usposabljanje

Začnite z beleženjem velikosti serije, vhodnih dimenzij, natančnosti in točke, kjer pride do napake. Nato preverite tako pomnilnik aktivnih tenzorjev kot pomnilnik, ki ga je rezerviral dodeljevalnik. PyTorch omogoča dostop do memory_allocated(), memory_reserved(), vrhunskih variant in memory_summary(). Trenutna dokumentacija o upravljanju pomnilnika CUDA pojasnjuje, da predpomnilniški dodeljevalnik ohranja ponovno uporabne bloke, zato se neuporabljen rezerviran pomnilnik lahko še vedno prikaže kot zaseden v orodjih za spremljanje GPU. Upravljanje pomnilnika CUDA v PyTorchu.

import torch

torch.cuda.reset_peak_memory_stats()

# Tukaj izvedite en reprezentativen korak usposabljanja.

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

Če preprost povzetek ni dovolj, lahko PyTorch zajame posnetke dodeljevalnika za poglobljeno analizo. Njegova orodja za pomnilnik lahko zabeležijo zgodovino dodelitev in ustvarijo posnetek, ki ga lahko pregledate z vizualizatorjem pomnilnika PyTorch. PyTorch opozarja, da ta orodja vidijo pomnilnik, ki ga upravlja dodeljevalnik PyTorch; dodelitve, ki jih neposredno izvajajo druge knjižnice CUDA, se morda ne bodo pojavile tam. Vodnik PyTorch za razumevanje uporabe pomnilnika CUDA.

Ne obravnavajte torch.cuda.empty_cache() kot splošno rešitev za OOM

torch.cuda.empty_cache() sprosti neuporabljene predpomnjene bloke, da jih lahko uporabijo druge aplikacije GPU. PyTorch izrecno navaja, da ne sprosti pomnilnika, ki ga zasedajo aktivni tenzorji, in zato ne poveča količine pomnilnika GPU, ki je na voljo PyTorchu za tenzorje, ki so še vedno aktivni. Koristen je med ločenimi eksperimenti ali po brisanju velikih objektov, vendar ni nadomestilo za zmanjšanje aktivne zasedenosti pomnilnika.

Korak 2: Najprej zmanjšajte aktivni delovni nabor

Najbolj zanesljiv prvi popravek je običajno manjša mikro-serija: število vzorcev, obdelanih v enem naprednem/nazajnem prehodu. Pomnilnik aktivacij običajno narašča z velikostjo serije, ločljivostjo slike, dolžino zaporedja in drugimi vhodnimi dimenzijami. Če se model usposablja pri velikosti serije 32, a odpove pri 64, zmanjšanje serije ni zaobid v negativnem smislu; gre za neposredno zmanjšanje povpraševanja po vrhunskem pomnilniku.

Pri slikah lahko znižanje prostorske ločljivosti ali velikosti obreza naredi veliko razliko. Pri transformatorjih in drugih modelih zaporedij je zmanjšanje dolžine zaporedja lahko še pomembnejše, ker nekatere vmesne tenzorje močno naraščajo z dolžino zaporedja. Točno skaliranje je odvisno od arhitekture, zato merite in ne predpostavljajte.

Preverite tudi, ali koda za vrednotenje nepotrebno gradi gradiente. Smernice za zmogljivost PyTorch priporočajo onemogočanje izračuna gradientov za validacijo ali sklepanje, ko gradienti niso potrebni, ker autograd sicer shrani vmesne pomnilniške medpomnilnike. Tipičen vzorec 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)

Med usposabljanjem uporabite optimizer.zero_grad(set_to_none=True), razen če vaš algoritem temelji na vedenjski razliki med ničelnim gradientom in gradientom None. Dokumentacija optimizatorja PyTorch navaja, da nastavitev gradientov na None običajno zahteva manj pomnilnika in lahko zmerno izboljša zmogljivost. Dokumentacija PyTorch optimizatorja zero_grad.

Korak 3: Ohranite večjo efektivno serijo z AMP in akumulacijo gradientov

Uporabite samodejno mešano natančnost, ko jo model podpira

Samodejna mešana natančnost (AMP) izvaja ustrezne operacije v nižji natančnosti, hkrati pa ohranja operacije, ki potrebujejo večji obseg ali natančnost, v ustreznih tipih. PyTorch dokumentira, da lahko AMP izboljša zmogljivost in zmanjša zasedenost pomnilnika za mnoga obremenitve CUDA, vendar ni numerično primerna za vsak model. Zlasti PyTorch opozarja, da lahko nekateri modeli, predusposobljeni v bfloat16, pretečejo v float16.

Trenutni vzorec usposabljanja AMP za 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 priporoča izvajanje naprednega prehoda in izgube pod autocast, nato pa zapustite kontekst autocast pred nazajno propagacijo. Če float16 povzroča nestabilnost, preverite, ali je bfloat16 podprt in ustrezen za vašo strojno opremo in model, namesto da predpostavljate, da se vsi načini mešane natančnosti obnašajo enako.

Uporabite akumulacijo gradientov, ko potrebujete večjo efektivno serijo

Akumulacija gradientov obdela več manjših mikro-serij pred posodobitvijo optimizatorja. Če je mikro-serija 4 in akumulirate 8 korakov, je efektivna serija za eno posodobitev optimizatorja 32 vzorcev na delavca, ob predpostavki, da ima vsaka mikro-serija štiri vzorce in postavitev podatkovne paralelizacije ne spremeni te aritmetike.

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 produkcijsko kodo obravnavajte tudi končno delno okno akumulacije, ko število serij ni deljivo s accum_steps. Če uporabljate distribuirano usposabljanje, se lahko vedenje sinhronizacije gradientov spremeni, kar vpliva na razmerje med pomnilnikom in zmogljivostjo, zato sledite smernicam za akumulacijo distribuiranega API in ne kopirajte zanke za en GPU nespremenjeno.

Korak 4: Zamenjajte računanje za pomnilnik, nato preiskujte ohranjanje in fragmentacijo

Shranjevanje aktivacij v kontrolne točke

Shranjevanje aktivacij v kontrolne točke zmanjša pomnilnik tako, da ne ohranja izbranih aktivacij naprednega prehoda do nazajne propagacije. Namesto tega jih PyTorch ponovno izračuna med nazajno propagacijo. To zamenja dodatno računanje za nižjo zasedenost pomnilnika aktivacij. Trenutna dokumentacija PyTorch za kontrolne točke priporoča izrecno podajanje use_reentrant=False. Dokumentacija PyTorch za shranjevanje aktivacij v kontrolne točke.

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)

Shranjujte sloje v kontrolne točke z velikimi shranjenimi aktivacijami in sprejemljivimi stroški ponovnega računanja. Ne predpostavljajte, da je shranjevanje vsake operacije v kontrolne točke optimalno; lahko bistveno upočasni usposabljanje.

Iščite tenzorje, ki ohranjajo računske grafe aktivne

Če pomnilnik narašča pri vsaki iteraciji namesto da bi dosegel vrh na približno enaki ravni, preverite reference Python. Pogost vzorec je shranjevanje tenzorjev, povezanih z grafom, v seznam:

# Tvegano, če se ohranja za veliko korakov:
loss_history.append(loss)

# Shranite namesto tega število Python:
loss_history.append(loss.item())

Ista težava se lahko pojavi, ko predpomnite izhode modela, zemljevide pozornosti, skrita stanja ali tenzorje validacije brez ločitve ali premika z GPU. Izbrišite reference, ki jih ne potrebujete več, in uporabite detach() le, ko namerno želite tenzor, ločen od autograda.

Prilagodite dodeljevalnik šele, ko statistika kaže na fragmentacijo

Trenutna dokumentacija PyTorch daje prednost spremenljivki okolja PYTORCH_ALLOC_CONF. Starejša PYTORCH_CUDA_ALLOC_CONF ostaja vzdevek za nazaj združljivost. Ta podrobnost poimenovanja se je spremenila v trenutnih dokumentih, zato bi morale nove konfiguracije uporabljati prednostno ime. Spremenljivke okolja CUDA v PyTorchu.

Dve možnosti dodeljevalnika sta še posebej pomembni:

• expandable_segments:True je eksperimentalna in je zasnovana za zmanjšanje neuporabnih rezin pomnilnika, ko se velikosti dodelitev spreminjajo, kot so obremenitve, katerih velikosti serij ali tenzorjev variirajo.
• max_split_size_mb lahko zmanjša fragmentacijo z izvirnim dodeljevalnikom, vendar jo PyTorch izrecno opisuje kot zadnjo možnost za obremenitve, ki odpovejo z OOM, hkrati pa kažejo veliko količino neaktivnih razcepljenih blokov. Lahko tudi škoduje zmogljivosti in jo zaledje cudaMallocAsync ignorira.

# Primer za obremenitev z različnimi velikostmi dodelitev:
export PYTORCH_ALLOC_CONF=expandable_segments:True

Ne kopirajte zastavic dodeljevalnika z drugega računalnika, ne da bi preverili memory_summary() ali posnetek. Pravi problem s kapaciteto – kjer aktivni tenzorji že zapolnijo GPU – ne bo rešen s prilagajanjem fragmentacije.

Ko en GPU še vedno ne more sprejeti modela

Če en vzorec pri najmanjši praktični vhodni velikosti še vedno povzroči OOM, je težava morda v modelu in stanju optimizatorja, ne v seriji. V tem trenutku razmislite o manjši arhitekturi, parametrih z nižjo natančnostjo, kjer je to numerično ustrezno, strategijah CPU/offload ali distribuiranem usposabljanju z razdelitvijo.

PyTorch Fully Sharded Data Parallel (FSDP) lahko razdeli parametre modela med delavce podatkovne paralelizacije, njegova strategija FULL_SHARD pa razdeli tudi gradienti in stanja optimizatorja. To lahko zmanjša pomnilnik na GPU v primerjavi s popolnoma replicirano podatkovno paralelizacijo, vendar na račun komunikacije in bolj zapletenega vedenja usposabljanja. Dokumentacija PyTorch FSDP.

Praktični vrstni red ukrepov

Seznam preverjanja: kako vedeti, da je OOM dejansko odpravljen

• Izvedite več reprezentativnih iteracij usposabljanja, ne le enega uspešnega naprednega prehoda.
• Ponastavite in zabeležite max_memory_allocated(), da poznate nov vrh.
• Potrdite, da pomnilnik GPU doseže stabilno območje namesto da bi naraščal pri vsaki iteraciji.
• Validirajte izgubo in gradienti po omogočanju mešane natančnosti.
• Potrdite, da akumulacija gradientov ohranja načrt posodobitev optimizatorja, ki ste ga želeli.
• Izvedite prehod validacije pod torch.no_grad(), ko gradienti niso zahtevani.
• Če ste spremenili nastavitve dodeljevalnika, primerjajte statistiko pomnilnika in prepustnost pred in po spremembi.
• Ne označite težave za rešeno samo zato, ker nvidia-smi pokaže manj rezerviranega pomnilnika po empty_cache(); samo delovno obremenitev usposabljanja mora dokončati pri svojem običajnem vrhu.

Napako CUDA OOM je najbolje obravnavati kot problem proračuna pomnilnika, ne kot posamezno napako PyTorch. Izmerite vrh, najprej zmanjšajte aktivni delovni nabor, nato uporabite mešano natančnost, akumulacijo in shranjevanje v kontrolne točke kot namjerne kompromise. Preidite na prilagajanje dodeljevalnika le, ko statistika dodeljevalnika kaže na fragmentacijo, in preidite na razdelitev ali drug model, ko se model sam ne prilega več udobno na en GPU.