Så åtgärdar du “PyTorch CUDA Out of Memory” under modellträning

En PyTorch-träningskörning kan fungera i flera steg och sedan avbrytas med torch.OutOfMemoryError eller ett meddelande som CUDA out of memory. Tried to allocate .... Den omedelbara orsaken är enkel: nästa CUDA-allokering fick inte plats. Den användbara frågan är varför den inte fick plats.

Under träning kan GPU-minnet innehålla modellparametrar, gradienter, optimerartillstånd, indatatensorer, tillfälliga arbetsområden och aktiveringar som sparats för bakåtpropagering. PyTorch använder också en cachande allokatör, så minne som visas som “reserverat” är inte identiskt med minne som för närvarande upptas av levande tensorer. Denna distinktion är viktig när man avgör om man ska minska arbetsbelastningen eller undersöka allokatorsfragmentering.

Denna guide följer den aktuella PyTorch-dokumentationen och använder de nuvarande AMP-API-namnen. I synnerhet dokumenterar PyTorch nu torch.amp.autocast("cuda") och torch.amp.GradScaler("cuda"); de äldre torch.cuda.amp.*-ingångspunkterna är föråldrade. Se PyTorch Automatic Mixed Precision-dokumentation.

Snabb triage: vilken typ av OOM hanterar du?

Steg 1: Mät minnet innan du ändrar träningsreceptet

Börja med att registrera batchstorlek, indatadimensioner, precision och punkten där felet inträffar. Inspektera sedan både levande tensor-minne och allokatorreserverat minne. PyTorch exponerar memory_allocated(), memory_reserved(), toppvarianter och memory_summary(). Den aktuella CUDA-minneshanteringsdokumentationen förklarar att den cachande allokatören behåller återanvändbara block, vilket är varför oanvänt reserverat minne fortfarande kan visas som använt i GPU-övervakningsverktyg. PyTorch CUDA-minneshantering.

import torch

torch.cuda.reset_peak_memory_stats()

# Kör ett representativt träningssteg här.

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

Om en enkel sammanfattning inte räcker kan PyTorch fånga allokatorsögonblicksbilder för djupare analys. Dess minnesverktyg kan registrera allokeringshistorik och producera en ögonblicksbild som kan inspekteras med PyTorchs minnesvisualiserare. PyTorch noterar att dessa verktyg ser minne som hanteras av PyTorch-allokatören; allokeringar som görs direkt av andra CUDA-bibliotek kan inte synas där. PyTorch-guide för att förstå CUDA-minnesanvändning.

Behandla inte torch.cuda.empty_cache() som en allmän OOM-lösning

torch.cuda.empty_cache() frigör oanvända cachade block så att andra GPU-applikationer kan använda dem. PyTorch anger uttryckligen att det inte frigör minne som upptas av levande tensorer och ökar därför inte mängden GPU-minne som är tillgängligt för PyTorch för tensorer som fortfarande är levande. Det kan vara användbart mellan separata experiment eller efter att stora objekt raderats, men det är ingen ersättning för att minska den aktiva minnesavtrycket.

Steg 2: Minska den aktiva arbetsuppsättningen först

Den mest pålitliga första lösningen är vanligtvis en mindre mikrobatch: antalet prover som behandlas av ett framåt-/bakåtsteg. Aktiveringsminnet växer vanligtvis med batchstorlek, bildupplösning, sekvenslängd och andra indatadimensioner. Om modellen tränas vid batchstorlek 32 men misslyckas vid 64, är att minska batchen inte en workaround i negativ mening; det är en direkt minskning av peak-minnesbehovet.

För bilder kan sänkt spatial upplösning eller beskärd storlek göra en stor skillnad. För transformers och andra sekvensmodeller kan det vara ännu viktigare att minska sekvenslängden eftersom vissa mellanliggande tensorer växer starkt med sekvenslängden. Den exakta skalningen beror på arkitekturen, så mät istället för att anta.

Säkerställ också att utvärderingskoden inte bygger gradienter i onödan. PyTorchs prestandariktlinjer rekommenderar att inaktivera gradientberäkning för validering eller inferens när gradienter inte behövs, eftersom autograd annars sparar mellanliggande buffertar. Ett typiskt mönster är:

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)

Under träning, använd optimizer.zero_grad(set_to_none=True) om inte din algoritm förlitar sig på den beteendemässiga skillnaden mellan en nollgradient och en None-gradient. PyTorchs optimerardokumentation anger att inställning av gradienter till None generellt har ett lägre minnesavtryck och kan beskedligt förbättra prestanda. PyTorch optimerare zero_grad-dokumentation.

Steg 3: Behåll en större effektiv batch med AMP och gradientackumulering

Använd Automatic Mixed Precision när modellen stöder det

Automatic Mixed Precision (AMP) kör kvalificerade operationer i lägre precision medan operationer som behöver mer räckvidd eller precision behålls i lämpliga typer. PyTorch dokumenterar att AMP kan förbättra prestanda och minska minnesavtrycket för många CUDA-arbetsbelastningar, men det är inte numeriskt lämpligt för varje modell. I synnerhet varnar PyTorch för att vissa modeller förtränade i bfloat16 kan överlöpa i float16.

Ett aktuellt CUDA AMP-träningsmönster är:

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 rekommenderar att köra framåtsteget och förlustfunktionen under autocast, och sedan lämna autocast-kontexten innan bakåtsteget. Om float16 producerar instabilitet, undersök om bfloat16 stöds och är lämpligt för din hårdvara och modell istället för att anta att alla blandade precisionslägen beter sig identiskt.

Använd gradientackumulering när du behöver en större effektiv batch

Gradientackumulering behandlar flera mindre mikrobatcher innan optimeraren uppdateras. Om mikrobatchen är 4 och du ackumulerar 8 steg, är den effektiva batchen för en optimeraruppdatering 32 prover per arbetare, förutsatt att varje mikrobatch har fyra prover och att den dataparallella uppsättningen inte ändar denna aritmetik.

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)

För produktionskod, hantera också ett sista partiellt ackumuleringsfönster när antalet batcher inte är delbart med accum_steps. Om du använder distribuerad träning kan gradientssynkroniseringsbeteendet ändra minnes-/prestandaavvägningen, så följ den distribuerade API:ns ackumuleringsriktlinjer istället för att kopiera en enkel-GPU-loop oförändrad.

Steg 4: Byt beräkning mot minne, undersök sedan retention och fragmentering

Aktiveringscheckpointing

Aktiveringscheckpointing minskar minnet genom att inte hålla utvalda framåtaktiveringar levande tills bakåtsteget. Istället omräknar PyTorch dem under bakåtsteget. Detta byter ytterligare beräkning mot ett lägre aktiveringsminnesavtryck. Den aktuella PyTorch-checkpoint-dokumentationen rekommenderar att uttryckligen skicka use_reentrant=False. PyTorch aktiveringscheckpointing-dokumentation.

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)

Checkpointa lager med stora sparade aktiveringar och acceptabel omräkningskostnad. Anta inte att checkpointing av varje operation är optimalt; det kan avsevärt sänka träningstakten.

Leta efter tensorer som håller beräkningsgrafer levande

Om minnet växer varje iteration istället för att toppa på ungefär samma nivå, inspektera Python-referenser. Ett vanligt mönster är att lagra grafanslutna tensorer i en lista:

# Riskabelt om behållna under många steg:
loss_history.append(loss)

# Lagra ett Python-tal istället:
loss_history.append(loss.item())

Samma problem kan uppstå när du cachar modellutdata, uppmärksamhetskartor, dolda tillstånd eller valideringstensorer utan att koppla bort dem eller flytta dem från GPU:n. Radera referenser du inte längre behöver, och använd detach() endast när du avsiktligt vill ha en tensor kopplad från autograd.

Justera allokatören endast efter att statistiken pekar på fragmentering

Nuvarande PyTorch-dokumentation föredrar miljövariabeln PYTORCH_ALLOC_CONF. Den äldre PYTORCH_CUDA_ALLOC_CONF förblir ett alias för bakåtkompatibilitet. Denna namndetalj ändrades i den aktuella dokumentationen, så nya konfigurationer bör använda det föredragna namnet. PyTorch CUDA-miljövariabler.

Två allokatoralternativ är särskilt relevanta:

• expandable_segments:True är experimentellt och är utformat för att minska oanvändbara minnesskivor när allokstorlekar ändras, såsom arbetsbelastningar vars batch- eller tensorstorlekar varierar.
• max_split_size_mb kan minska fragmentering med den inbyggda allokatören, men PyTorch beskriver uttryckligen det som en sista utväg för arbetsbelastningar som misslyckas med OOM medan de visar en stor mängd inaktiva delade block. Det kan också skada prestanda och ignoreras av cudaMallocAsync-backenden.

# Exempel för en arbetsbelastning med varierande allokstorlekar:
export PYTORCH_ALLOC_CONF=expandable_segments:True

Kopiera inte allokatorflaggor från en annan maskin utan att kontrollera memory_summary() eller en ögonblicksbild. Ett verkligt kapacitetsproblem–där levande tensorer redan fyller GPU:n–kommer inte att lösas av fragmenteringsjustering.

När en GPU fortfarande inte kan rymda modellen

Om ett enda prov vid den minsta praktiska indatabreddan fortfarande ger OOM, kan problemet vara modellen och optimerartillståndet snarare än batchen. Vid den punkten, överväg en mindre arkitektur, parametrar med lägre precision där numeriskt lämpligt, CPU/offload-strategier eller sharded distribuerad träning.

PyTorchs Fully Sharded Data Parallel (FSDP) kan sharda modellparametrar över dataparallella arbetare, och dess FULL_SHARD-strategi shardar också gradienter och optimerartillstånd. Detta kan minska per-GPU-minne jämfört med fullt replikerad dataparallellism, till priset av kommunikation och mer komplext träningsbeteende. PyTorch FSDP-dokumentation.

Praktisk ordning av åtgärder

Checklista: hur du vet att OOM faktiskt är åtgärdat

• Kör flera representativa träningsiterationer, inte bara ett framgångsrikt framåtsteg.
• Återställ och registrera max_memory_allocated() så att du känner till den nya toppen.
• Bekräfta att GPU-minnet når ett stabilt intervall istället för att öka varje iteration.
• Validera förlust och gradienter efter att blandad precision aktiverats.
• Bekräfta att gradientackumulering bevarar den optimeraruppdateringsschema du avsåg.
• Kör ett valideringssteg under torch.no_grad() när gradienter inte krävs.
• Om du ändrade allokatorinställningar, jämför minnesstatistik och genomströmning före och efter.
• Kalla inte problemet löst enbart för att nvidia-smi visar mindre reserverat minne efter empty_cache(); träningsarbetsbelastningen själv måste slutföra vid sin normala topp.

En CUDA OOM bör behandlas som ett minnesbudgetproblem, inte som en enda PyTorch-bugg. Mät toppen, minska den levande arbetsuppsättningen först, använd sedan blandad precision, ackumulering och checkpointing som medvetna avvägningar. Gå bara över till allokatorjustering när allokatorstatistiken indikerar fragmentering, och gå över till sharding eller en annan modell när modellen själv inte längre får plats bekvämt på en GPU.