Kā novērst “PyTorch CUDA Out of Memory” kļūdu modeļa apmācības laikā

PyTorch apmācības process var darboties vairākus soļus, bet pēc tam apstāties ar kļūdu torch.OutOfMemoryError vai ziņojumu, piemēram, CUDA out of memory. Tried to allocate .... Tiešais iemesls ir vienkāršs: nākamajai CUDA atmiņas rezervācijai nebija pietiekami daudz vietas. Noderīgākais jautājums ir kāpēc tai nebija pietiekami daudz vietas.

Apmācības laikā GPU atmiņā var atrasties modeļa parametri, gradienti, optimizatora stāvoklis, ievades tenzori, pagaidu darbvietas un atpakaļejošajai izplatīšanai saglabātās aktivizācijas. PyTorch izmanto arī kešatmiņas atmiņas dalītāju, tāpēc atmiņa, kas tiek rādīta kā “rezervēta”, nav identiska atmiņai, ko pašlaik aizņem aktīvie tenzori. Šī atšķirība ir svarīga, nosakot, vai samazināt slodzi, vai pētīt dalītāja fragmentāciju.

Šis ceļvedis balstās uz pašreizējo PyTorch dokumentāciju un izmanto pašreizējos AMP API nosaukumus. Konkrēti, PyTorch tagad dokumentē torch.amp.autocast("cuda") un torch.amp.GradScaler("cuda"); vecākie torch.cuda.amp.* ieejas punkti ir novecojuši. Skatiet PyTorch automātiskās jauktās precizitātes dokumentāciju.

Ātra triāža: ar kāda veida OOM kļūdu jūs saskaraties?

1. solis: Izmēriet atmiņu, pirms maināt apmācības recepti

Sāciet, reģistrējot partijas izmēru, ievades dimensijas, precizitāti un brīdi, kad rodas kļūme. Pēc tam pārbaudiet gan aktīvo tenzoru atmiņu, gan dalītāja rezervēto atmiņu. PyTorch nodrošina piekļuvi funkcijām memory_allocated(), memory_reserved(), to maksimālajām vērtībām un memory_summary(). Pašreizējā CUDA atmiņas pārvaldības dokumentācija skaidro, ka kešatmiņas dalītājs saglabā atkārtoti lietojamus blokus, tāpēc neizmantotā rezervētā atmiņa GPU uzraudzības rīkos joprojām var izskatīties kā aizņemta. PyTorch CUDA atmiņas pārvaldība.

import torch

torch.cuda.reset_peak_memory_stats()

# Šeit palaidiet vienu reprezentatīvu apmācības soli.

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

Ja vienkāršs kopsavilkums nav pietiekams, PyTorch var uzņemt dalītāja momentuzņēmumus dziļākai analīzei. Tā atmiņas rīki var reģistrēt rezervāciju vēsturi un izveidot momentuzņēmumu, ko var pārbaudīt, izmantojot PyTorch atmiņas vizualizatoru. PyTorch norāda, ka šie rīki redz tikai to atmiņu, ko pārvalda PyTorch dalītājs; citas CUDA bibliotēkas veiktās rezervācijas tur var neparādīties. PyTorch ceļvedis CUDA atmiņas lietošanas izpratnei.

Nekonsiderējiet torch.cuda.empty_cache() par vispārīgu OOM risinājumu

torch.cuda.empty_cache() atbrīvo neizmantotos kešatmiņas blokus, lai citas GPU lietotnes varētu tos izmantot. PyTorch skaidri norāda, ka tas neatbrīvo atmiņu, ko aizņem aktīvie tenzori, un tāpēc nepalielina GPU atmiņas daudzumu, kas pieejams PyTorch tenzoriem, kuri joprojām ir dzīvi. Tas var būt noderīgi starp atsevišķiem eksperimentiem vai pēc lielu objektu dzēšanas, taču tas nav aizstājējs aktīvās atmiņas pēdas samazināšanai.

2. solis: Vispirms samaziniet aktīvo darba kopu

Visuzticamākais pirmais risinājums parasti ir mazāka mikro-partija: paraugu skaits, ko apstrādā viens uzlādējošais/atpakaļejošais posms. Aktivizāciju atmiņa parasti pieaug līdz ar partijas izmēru, attēla izšķirtspēju, virknes garumu un citām ievades dimensijām. Ja modelis apmācās ar partijas izmēru 32, bet neizdodas ar 64, partijas samazināšana nav negatīva nozīme risinājums; tā ir tieša maksimālās atmiņas pieprasījuma samazināšana.

Attēliem telpiskās izšķirtspējas vai apgriešanas izmēra samazināšana var radīt lielu atšķirību. Transformeru un citiem virkņu modeļiem virknes garuma samazināšana var būt vēl svarīgāka, jo daži starpposma tenzori strauji pieaug līdz ar virknes garumu. Precīza mērogošana ir atkarīga no arhitektūras, tāpēc mēriet, nevis pieņemiet.

Pārliecinieties arī, ka novērtēšanas kods nevajadzīgi neveido gradientus. PyTorch veiktspējas norādījumos ieteicams izslēgt gradientu aprēķināšanu validācijai vai secināšanai, kad gradienti nav nepieciešami, jo autograd citādi saglabā starpposma buferus. Tipisks paraugs ir:

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)

Apmācības laikā izmantojiet optimizer.zero_grad(set_to_none=True), ja vien jūsu algoritms nen paļaujas uz uzvedības atšķirību starp nulles gradientu un None gradientu. PyTorch optimizatora dokumentācija norāda, ka gradientu iestatīšana uz None parasti samazina atmiņas pēdu un var nedaudz uzlabot veiktspēju. PyTorch optimizatora zero_grad dokumentācija.

3. solis: Saglabājiet lielāku efektīvo partiju, izmantojot AMP un gradientu uzkrāšanu

Izmantojiet automātisko jaukto precizitāti, ja modelis to atbalsta

Automātiskā jauktā precizitāte (AMP) veic piemērotas operācijas zemākā precizitātē, saglabājot operācijas, kurām nepieciešams lielāks diapazons vai precizitāte, atbilstošos tipos. PyTorch dokumentē, ka AMP var uzlabot veiktspēju un samazināt atmiņas pēdu daudziem CUDA slodzes veidiem, taču tas nav skaitliski piemērots katram modelim. Konkrēti, PyTorch brīdina, ka daži modeļi, kas iepriekš apmācīti ar bfloat16, var piedzīvot pārpildi, izmantojot float16.

Pašreizējais CUDA AMP apmācības paraugs ir:

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 iesaka veikt uzlādējošo posmu un zaudējumu aprēķinu zem autocast, bet pirms atpakaļejošā posma atstāt autocast kontekstu. Ja float16 rada nestabilitāti, izpētiet, vai bfloat16 ir atbalstīts un piemērots jūsu aparatūrai un modelim, nevis pieņemiet, ka visi jauktas precizitātes režīmi uzvedas identiski.

Izmantojiet gradientu uzkrāšanu, ja nepieciešama lielāka efektīvā partija

Gradientu uzkrāšana apstrādā vairākas mazākas mikro-partijas, pirms atjaunina optimizatoru. Ja mikro-partija ir 4 un jūs uzkrājat 8 soļus, efektīvā partija vienam optimizatora atjauninājumam ir 32 paraugi uz vienu darbinieku, pieņemot, ka katrai mikro-partijai ir četri paraugi un datu paralēlā iestatīšana nemaina šo aritmētiku.

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)

Ražošanas kodā apstrādājiet arī pēdējo daļējo uzkrāšanas logu, ja partiju skaits nav dalāms ar accum_steps. Ja izmantojat sadalītu apmācību, gradientu sinhronizācijas uzvedība var mainīt atmiņas/veiktspējas kompromisu, tāpēc sekojiet sadalītās API uzkrāšanas norādījumiem, nevis kopējiet viena GPU ciklu nemainītu.

4. solis: Apmainiet aprēķinus pret atmiņu, pēc tam izpētiet saglabāšanu un fragmentāciju

Aktivizāciju kontrolpunkti

Aktivizāciju kontrolpunkti samazina atmiņu, nesaglabājot izvēlētās uzlādējošās aktivizācijas līdz atpakaļejošajam posmam. Tā vietā PyTorch tos pārrēķina atpakaļejošā posma laikā. Tas apmaina papildu aprēķinus pret zemāku aktivizāciju atmiņas pēdu. Pašreizējā PyTorch kontrolpunktu dokumentācija iesaka skaidri nodot use_reentrant=False. PyTorch aktivizāciju kontrolpunktu dokumentācija.

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)

Izveidojiet kontrolpunktus slāņiem ar lielām saglabātajām aktivizācijām un pieņemamu pārrēķina izmaksu. Nepieņemiet, ka kontrolpunktu izveide katrai operācijai ir optimāla; tas var būtiski palēnināt apmācību.

Meklējiet tenzorus, kas uztur aprēķinu grafikus dzīvus

Ja atmiņa pieaug katrā iterācijā, nevis sasniedz aptuveni vienādu līmeni, pārbaudiet Python atsauces. Bieži sastopams paraugs ir grafika savienotu tenzoru glabāšana sarakstā:

# Riskanti, ja saglabāts daudzos soļos:
loss_history.append(loss)

# Glabājiet Python skaitli tā vietā:
loss_history.append(loss.item())

Tā pati problēma var rasties, kešojot modeļa izvades, uzmanības kartes, slēptos stāvokļus vai validācijas tenzorus, neatvienojot tos vai nepārvietojot no GPU. Dzēsiet atsauces, kuras vairs nevajag, un izmantojiet detach() tikai tad, ja apzināti vēlaties tenzoru, kas ir atvienots no autograd.

Pielāgojiet dalītāju tikai tad, kad statistika norāda uz fragmentāciju

Pašreizējā PyTorch dokumentācija dod priekšroku vides mainīgajam PYTORCH_ALLOC_CONF. Vecākais PYTORCH_CUDA_ALLOC_CONF paliek kā alias atpakaļejošai saderībai. Šī nosaukuma detaļa mainījās pašreizējā dokumentācijā, tāpēc jaunām konfigurācijām jāizmanto vēlamais nosaukums. PyTorch CUDA vides mainīgie.

Divas dalītāja opcijas ir īpaši nozīmīgas:

• expandable_segments:True ir eksperimentāls un paredzēts, lai samazinātu nelietojamus atmiņas gabaliņus, mainoties rezervāciju izmēriem, piemēram, slodzēm, kuru partiju vai tenzoru izmēri mainās.
• max_split_size_mb var samazināt fragmentāciju ar dabisko dalītāju, bet PyTorch to skaidri raksturo kā pēdējo līdzekli slodzēm, kas neizdodas ar OOM, vienlaikus rādot lielu daudzumu neaktīvu sadalītu bloku. Tas var arī pasliktināt veiktspēju, un cudaMallocAsync aizmugurē tas tiek ignorēts.

# Piemērs slodzei ar mainīgiem rezervāciju izmēriem:
export PYTORCH_ALLOC_CONF=expandable_segments:True

Nekopējiet dalītāja karodziņus no citas mašīnas, nepārbaudot memory_summary() vai momentuzņēmumu. Patiesu kapacitātes problēmu—kur aktīvie tenzori jau aizpilda GPU—nevar atrisināt ar fragmentācijas pielāgošanu.

Kad viens GPU joprojām nevar ietilpināt modeli

Ja viena parauga apstrāde ar mazāko praktisko ievades izmēru joprojām izraisa OOM, problēma var būt modelis un optimizatora stāvoklis, nevis partija. Šajā brīdī apsveriet mazāku arhitektūru, zemākas precizitātes parametrus, kur tas ir skaitliski piemērots, CPU/offload stratēģijas vai sadalītu sadalītu apmācību.

PyTorch Fully Sharded Data Parallel (FSDP) var sadalīt modeļa parametrus starp datu paralēliem darbiniekiem, un tā FULL_SHARD stratēģija sadala arī gradientus un optimizatora stāvokļus. Tas var samazināt atmiņu uz vienu GPU, salīdzinot ar pilnībā replicētu datu paralēlismu, par cenu komunikācijas un sarežģītākas apmācības uzvedības. PyTorch FSDP dokumentācija.

Praktiska darbību secība

Kontrolsaraksts: kā zināt, ka OOM kļūda ir patiešām novērsta

• Palaidiet vairākus reprezentatīvus apmācības iterācijas, nevis tikai vienu veiksmīgu uzlādējošo posmu.
• Atiestatiet un reģistrējiet max_memory_allocated(), lai zinātu jauno pīķi.
• Pārliecinieties, ka GPU atmiņa sasniedz stabilu diapazonu, nevis pieaug katrā iterācijā.
• Validējiet zaudējumus un gradientus pēc jauktas precizitātes iespējošanas.
• Pārliecinieties, ka gradientu uzkrāšana saglabā jūsu paredzēto optimizatora atjaunināšanas grafiku.
• Palaidiet validācijas posmu zem torch.no_grad(), kad gradienti nav nepieciešami.
• Ja mainījāt dalītāja iestatījumus, salīdziniet atmiņas statistiku un caurlaidspēju pirms un pēc.
• Nenosauciet problēmu atrisinātu tikai tāpēc, ka nvidia-smi rāda mazāk rezervētas atmiņas pēc empty_cache(); pašai apmācības slodzei ir jāpabeidz ar savu parasto pīķi.

CUDA OOM kļūda ir vislabāk jāuztver kā atmiņas budžeta problēma, nevis kā viena PyTorch kļūda. Izmēriet pīķi, vispirms samaziniet dzīvo darba kopu, pēc tam izmantojiet jaukto precizitāti, uzkrāšanu un kontrolpunktus kā apzinātus kompromisus. Pārejiet uz dalītāja pielāgošanu tikai tad, kad dalītāja statistika norāda uz fragmentāciju, un pārejiet uz sadalīšanu vai citu modeli, kad pats modelis vairs ērti neietilpst vienā GPU.