Sådan løser du “PyTorch CUDA Out of Memory” under modeltræning

En PyTorch-træningskørsel kan fungere i flere trin og derefter stoppe med torch.OutOfMemoryError eller en besked som CUDA out of memory. Tried to allocate .... Den umiddelbare årsag er enkel: den næste CUDA-allokering kunne ikke få plads. Det nyttige spørgsmål er hvorfor den ikke fik plads.

Under træning kan GPU-hukommelsen indeholde modelparametre, gradienter, optimizer-tilstand, input-tensorer, midlertidige arbejdsområder og aktiveringer gemt til backward-pass. PyTorch bruger også en caching-allocator, så hukommelse vist som “reserved” er ikke identisk med hukommelse, der i øjeblikket er optaget af levende tensorer. Denne distinktion er vigtig, når man skal afgøre, om man skal reducere arbejdsbyrden eller undersøge allocator-fragmentering.

Denne guide følger den aktuelle PyTorch-dokumentation og bruger de nuværende AMP-API-navne. Især dokumenterer PyTorch nu torch.amp.autocast("cuda") og torch.amp.GradScaler("cuda"); de ældre torch.cuda.amp.* indgangspunkter er forældede. Se PyTorch Automatic Mixed Precision-dokumentation.

Hurtig triage: hvilken type OOM har du at gøre med?

Trin 1: Mål hukommelse, før du ændrer træningsopskriften

Start med at registrere batch-størrelse, input-dimensioner, præcision og det punkt, hvor fejlen opstår. Undersøg derefter både levende tensor-hukommelse og allocator-reserveret hukommelse. PyTorch eksponerer memory_allocated(), memory_reserved(), peak-varianter og memory_summary(). Den aktuelle CUDA-hukommelsesstyringsdokumentation forklarer, at caching-allocatoren holder genanvendelige blokke, hvilket er grunden til, at ubrugt reserveret hukommelse stadig kan fremstå som brugt i GPU-overvågningsværktøjer. PyTorch CUDA-hukommelsesstyring.

import torch

torch.cuda.reset_peak_memory_stats()

# Kør ét repræsentativt træningstrin her.

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

Hvis en simpel opsummering ikke er nok, kan PyTorch fange allocator-snapshots til dybere analyse. Dets hukommelsesværktøjer kan registrere allokeringshistorik og producere et snapshot, der kan inspiceres med PyTorch-hukommelsesvisualisereren. PyTorch bemærker, at disse værktøjer ser hukommelse administreret af PyTorch-allocatoren; allokeringer foretaget direkte af andre CUDA-biblioteker vises muligvis ikke der. PyTorch-guide til forståelse af CUDA-hukommelsesforbrug.

Behandl ikke torch.cuda.empty_cache() som en generel OOM-løsning

torch.cuda.empty_cache() frigør ubrugte cachede blokke, så andre GPU-applikationer kan bruge dem. PyTorch angiver eksplicit, at det ikke frigør hukommelse optaget af levende tensorer og derfor ikke øger mængden af GPU-hukommelse tilgængelig for PyTorch til tensorer, der stadig er levende. Det kan være nyttigt mellem separate eksperimenter eller efter sletning af store objekter, men det er ikke en erstatning for at reducere det aktive hukommelsesaftryk.

Trin 2: Reducer det aktive arbejdsset først

Den mest pålidelige første løsning er normalt et mindre micro-batch: antallet af prøver behandlet af ét forward/backward-pass. Aktiveringshukommelse vokser typisk med batch-størrelse, billedopløsning, sekvenslængde og andre input-dimensioner. Hvis modellen træner ved batch-størrelse 32 men fejler ved 64, er reduktion af batchet ikke en workaround i negativ forstand; det er en direkte reduktion af peak-hukommelsesbehov.

For billeder kan sænkning af rumlig opløsning eller crop-størrelse gøre en stor forskel. For transformers og andre sekvensmodeller kan reduktion af sekvenslængde være endnu vigtigere, fordi nogle mellemliggende tensorer vokser stærkt med sekvenslængden. Den præcise skalering afhænger af arkitekturen, så mål i stedet for at antage.

Sørg også for, at evalueringskode ikke opbygger gradienter unødigt. PyTorchs ydeevneguidelines anbefaler at deaktivere gradientberegning for validering eller inferens, når gradienter ikke er nødvendige, fordi autograd ellers gemmer mellemliggende buffere. Et typisk mønster er:

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, brug optimizer.zero_grad(set_to_none=True), medmindre din algoritme afhænger af den adfærdsmæssige forskel mellem en nul-gradient og en None-gradient. PyTorchs optimizer-dokumentation angiver, at indstilling af gradienter til None generelt har et lavere hukommelsesaftryk og kan forbedre ydeevnen beskedent. PyTorch optimizer zero_grad-dokumentation.

Trin 3: Bevar en større effektiv batch med AMP og gradientakkumulering

Brug Automatic Mixed Precision, når modellen understøtter det

Automatic Mixed Precision (AMP) kører kvalificerede operationer i lavere præcision, mens operationer, der kræver større rækkevidde eller præcision, holdes i passende typer. PyTorch dokumenterer, at AMP kan forbedre ydeevnen og reducere hukommelsesaftrykket for mange CUDA-arbejdsbyrder, men det er ikke numerisk egnet til enhver model. Især advarer PyTorch om, at nogle modeller fortrænet i bfloat16 kan overløbe i float16.

Et nuværende CUDA AMP-træningsmønster er:

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 anbefaler at køre forward-pass og loss under autocast, og derefter forlade autocast-konteksten før backward. Hvis float16 producerer ustabilitet, så undersøg om bfloat16 er understøttet og passende for dit hardware og din model, i stedet for at antage, at alle mixed-precision-tilstande opfører sig identisk.

Brug gradientakkumulering, når du har brug for en større effektiv batch

Gradientakkumulering behandler flere mindre micro-batches, før optimizeren opdateres. Hvis micro-batchet er 4, og du akkumulerer 8 trin, er den effektive batch for én optimizer-opdatering 32 prøver pr. worker, forudsat at hvert micro-batch har fire prøver, og at data-parallel opsætningen ikke ændrer denne 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)

For produktionskode skal du også håndtere et sidste delvist akkumuleringsvindue, når antallet af batches ikke er deleligt med accum_steps. Hvis du bruger distribueret træning, kan gradient-synkroniseringsadfærd ændre hukommelse/ydelse-trade-offen, så følg den distribuerede APIs akkumuleringsvejledning i stedet for at kopiere en single-GPU-løkke uændret.

Trin 4: Byt beregning for hukommelse, og undersøg derefter bevarelse og fragmentering

Aktiverings-checkpointing

Aktiverings-checkpointing reducerer hukommelse ved ikke at holde udvalgte forward-aktiveringer levende indtil backward. I stedet genberegner PyTorch dem under backward. Dette bytter ekstra beregning for et lavere aktiverings-hukommelsesaftryk. Den aktuelle PyTorch-checkpoint-dokumentation anbefaler eksplicit at sende use_reentrant=False. PyTorch aktiverings-checkpointing-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)

Checkpoint lag med store gemte aktiveringer og acceptabel genberegningstid. Antag ikke, at checkpointing af hver operation er optimal; det kan bremse træningen betydeligt.

Se efter tensorer, der holder beregningsgrafer levende

Hvis hukommelsen vokser ved hver iteration i stedet for at toppe på omtrent samme niveau, så undersøg Python-referencer. Et almindeligt mønster er at gemme graf-forbundne tensorer i en liste:

# Risikabelt hvis bevaret i mange trin:
loss_history.append(loss)

# Gem et Python-tal i stedet:
loss_history.append(loss.item())

Det samme problem kan opstå, når du cachede model-output, opmærksomhedskort, skjulte tilstande eller valideringstensorer uden at detach dem eller flytte dem fra GPU'en. Slet referencer, du ikke længere har brug for, og brug detach() kun, når du bevidst ønsker en tensor frakoblet fra autograd.

Justér allocatoren kun efter statistikken peger på fragmentering

Nuværende PyTorch-dokumentation foretrækker miljøvariablen PYTORCH_ALLOC_CONF. Den ældre PYTORCH_CUDA_ALLOC_CONF forbliver et alias for bagudkompatibilitet. Denne navngivningsdetalje ændrede sig i den aktuelle dokumentation, så nye konfigurationer bør bruge det foretrukne navn. PyTorch CUDA-miljøvariabler.

To allocator-indstillinger er især relevante:

• expandable_segments:True er eksperimentelt og er designet til at reducere ubrugelige hukommelsessplinter, når allokeringsstørrelser ændrer sig, såsom arbejdsbyrder, hvor batch- eller tensorstørrelser varierer.
• max_split_size_mb kan reducere fragmentering med den native allocator, men PyTorch beskriver det eksplicit som en sidste udvej for arbejdsbyrder, der fejler med OOM, mens de viser en stor mængde inaktive split-blokke. Det kan også skade ydeevnen og ignoreres af cudaMallocAsync-backenden.

# Eksempel for en arbejdsbyrde med varierende allokeringsstørrelser:
export PYTORCH_ALLOC_CONF=expandable_segments:True

Kopiér ikke allocator-flags fra en anden maskine uden at tjekke memory_summary() eller et snapshot. Et ægte kapacitetsproblem—hvor levende tensorer allerede fylder GPU'en—vil ikke blive løst af fragmenteringstuning.

Når én GPU stadig ikke kan rumme modellen

Hvis én prøve ved den mindste praktiske inputstørrelse stadig giver OOM, kan problemet være modellen og optimizer-tilstanden snarere end batchet. På det tidspunkt skal du overveje en mindre arkitektur, lavere præcisionsparametre hvor numerisk passende, CPU/offload-strategier eller sharded distribueret træning.

PyTorchs Fully Sharded Data Parallel (FSDP) kan shard modelparametre på tværs af data-parallel workers, og dets FULL_SHARD-strategi sharded også gradienter og optimizer-tilstande. Dette kan reducere per-GPU-hukommelse sammenlignet med fuldt replikeret data parallelism, på bekostning af kommunikation og mere kompleks træningsadfærd. PyTorch FSDP-dokumentation.

Praktisk rækkefølge af handlinger

Tjekliste: hvordan du ved, at OOM faktisk er løst

• Kør flere repræsentative træningsiterationer, ikke kun ét vellykket forward-pass.
• Nulstil og registrér max_memory_allocated(), så du kender den nye peak.
• Bekræft at GPU-hukommelsen når et stabilt interval i stedet for at stige ved hver iteration.
• Valider loss og gradienter efter aktivering af mixed precision.
• Bekræft at gradientakkumulering bevarer den optimizer-opdateringsplan, du tilsigtede.
• Kør et valideringspas under torch.no_grad(), når gradienter ikke er påkrævet.
• Hvis du ændrede allocator-indstillinger, så sammenlign hukommelsesstatistik og gennemstrømning før og efter.
• Kald ikke problemet løst blot fordi nvidia-smi viser mindre reserveret hukommelse efter empty_cache(); selve træningsarbejdsbyrden skal kunne fuldføre ved sin normale peak.

En CUDA OOM bør bedst behandles som et hukommelsesbudgetproblem, ikke som en enkelt PyTorch-fejl. Mål peak, reducer det levende arbejdsset først, brug derefter mixed precision, akkumulering og checkpointing som bevidste trade-offs. Gå kun over til allocator-tuning, når allocator-statistikken indikerer fragmentering, og gå over til sharding eller en anden model, når modellen ikke længere passer komfortabelt på én GPU.