Πώς να διορθώσετε το σφάλμα “PyTorch CUDA Out of Memory” κατά την εκπαίδευση μοντέλων

Μια εκτέλεση εκπαίδευσης PyTorch μπορεί να λειτουργήσει για αρκετά βήματα και στη συνέχεια να σταματήσει με torch.OutOfMemoryError ή ένα μήνυμα όπως CUDA out of memory. Tried to allocate .... Η άμεση αιτία είναι απλή: η επόμενη κατανομή CUDA δεν μπορούσε να χωρέσει. Η χρήσιμη ερώτηση είναι γιατί δεν μπορούσε να χωρέσει.

Κατά την εκπαίδευση, η μνήμη GPU μπορεί να περιέχει παραμέτρους μοντέλου, κλίσεις, κατάσταση optimizer, tensors εισόδου, προσωρινούς χώρους εργασίας και ενεργοποιήσεις που αποθηκεύονται για το backward. Το PyTorch χρησιμοποιεί επίσης έναν caching allocator, οπότε η μνήμη που εμφανίζεται ως “reserved” δεν είναι ίδια με τη μνήμη που καταλαμβάνεται επί του παρόντος από ενεργά tensors. Αυτή η διάκριση είναι σημαντική όταν αποφασίζετε αν θα μειώσετε το φορτίο εργασίας ή θα διερευνήσετε τον κατακερματισμό του allocator.

Αυτός ο οδηγός ακολουθεί την τρέχουσα τεκμηρίωση του PyTorch και χρησιμοποιεί τα τρέχοντα ονόματα API του AMP. Συγκεκριμένα, το PyTorch τώρα τεκμηριώνει τα torch.amp.autocast("cuda") και torch.amp.GradScaler("cuda"). Οι παλαιότερες καταχωρίσεις torch.cuda.amp.* έχουν καταργηθεί. Δείτε την τεκμηρίωση PyTorch Automatic Mixed Precision.

Γρήγορη διάγνωση: τι είδους OOM αντιμετωπίζετε;

Βήμα 1: Μετρήστε τη μνήμη πριν αλλάξετε τη συνταγή εκπαίδευσης

Ξεκινήστε καταγράφοντας το μέγεθος του batch, τις διαστάσεις εισόδου, την ακρίβεια και το σημείο όπου συμβαίνει η αποτυχία. Στη συνέχεια, ελέγξτε τόσο τη μνήμη των ενεργών tensors όσο και τη δεσμευμένη μνήμη του allocator. Το PyTorch εκθέτει τις συναρτήσεις memory_allocated(), memory_reserved(), τις παραλλαγές κορυφής και το memory_summary(). Η τρέχουσα τεκμηρίωση διαχείρισης μνήμης CUDA εξηγεί ότι ο caching allocator διατηρεί επαναχρησιμοποιήσιμα μπλοκ, γι' αυτό η μη χρησιμοποιημένη δεσμευμένη μνήμη μπορεί να εμφανίζεται ως χρησιμοποιημένη στα εργαλεία παρακολούθησης GPU. Διαχείριση μνήμης PyTorch CUDA.

import torch

torch.cuda.reset_peak_memory_stats()

# Εκτελέστε εδώ ένα αντιπροσωπευτικό βήμα εκπαίδευσης.

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

Αν μια απλή περίληψη δεν είναι αρκετή, το PyTorch μπορεί να καταγράψει στιγμιότυπα του allocator για βαθύτερη ανάλυση. Τα εργαλεία μνήμης του μπορούν να καταγράψουν το ιστορικό κατανομών και να παράγουν ένα στιγμιότυπο που μπορεί να επιθεωρηθεί με τον οπτικοποιητή μνήμης του PyTorch. Το PyTorch σημειώνει ότι αυτά τα εργαλεία βλέπουν τη μνήμη που διαχειρίζεται ο allocator του PyTorch. Οι κατανομές που γίνονται απευθείας από άλλες βιβλιοθήκες CUDA μπορεί να μην εμφανίζονται εκεί. Οδηγός PyTorch για την κατανόηση της χρήσης μνήμης CUDA.

Μην θεωρείτε το torch.cuda.empty_cache() ως γενική λύση για OOM

Η συνάρτηση torch.cuda.empty_cache() απελευθερώνει μη χρησιμοποιημένα cached μπλοκ ώστε άλλες εφαρμογές GPU να μπορούν να τα χρησιμοποιήσουν. Το PyTorch δηλώνει ρητά ότι δεν ελευθερώνει μνήμη που καταλαμβάνεται από ενεργά tensors και επομένως δεν αυξάνει την ποσότητα μνήμης GPU διαθέσιμη στο PyTorch για tensors που είναι ακόμα ζωντανά. Μπορεί να είναι χρήσιμο μεταξύ ξεχωριστών πειραμάτων ή μετά τη διαγραφή μεγάλων αντικειμένων, αλλά δεν αντικαθιστά τη μείωση του ενεργού αποτυπώματος μνήμης.

Βήμα 2: Μειώστε πρώτα το ενεργό σύνολο εργασίας

Η πιο αξιόπιστη πρώτη διόρθωση είναι συνήθως ένα μικρότερο micro-batch: ο αριθμός των δειγμάτων που επεξεργάζονται ένα forward/backward pass. Η μνήμη ενεργοποίησης συνήθως αυξάνεται με το μέγεθος του batch, την ανάλυση εικόνας, το μήκος ακολουθίας και άλλες διαστάσεις εισόδου. Αν το μοντέλο εκπαιδεύεται με μέγεθος batch 32 αλλά αποτυγχάνει στο 64, η μείωση του batch δεν είναι workaround με αρνητική έννοια· είναι μια άμεση μείωση της ζήτησης κορυφαίας μνήμης.

Για εικόνες, η μείωση της χωρικής ανάλυσης ή του μεγέθους crop μπορεί να κάνει μεγάλη διαφορά. Για transformers και άλλα μοντέλα ακολουθιών, η μείωση του μήκους ακολουθίας μπορεί να είναι ακόμα πιο σημαντική επειδή ορισμένα ενδιάμεσα tensors αυξάνονται έντονα με το μήκος ακολουθίας. Η ακριβής κλιμάκωση εξαρτάται από την αρχιτεκτονική, οπότε μετρήστε αντί να υποθέσετε.

Επίσης, βεβαιωθείτε ότι ο κώδικας αξιολόγησης δεν δημιουργεί κλίσεις χωρίς λόγο. Οι οδηγίες απόδοσης του PyTorch συνιστούν την απενεργοποίηση του υπολογισμού κλίσεων για επικύρωση ή εξαγωγή (inference) όταν οι κλίσεις δεν χρειάζονται, επειδή το autograd αλλιώς αποθηκεύει ενδιάμεσους buffers. Ένα τυπικό μοτίβο είναι:

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)

Κατά την εκπαίδευση, χρησιμοποιήστε optimizer.zero_grad(set_to_none=True) εκτός αν ο αλγόριθμός σας βασίζεται στη διαφορά συμπεριφοράς μεταξύ μηδενικής κλίσης και κλίσης None. Η τεκμηρίωση optimizer του PyTorch δηλώνει ότι η ρύθμιση των κλίσεων σε None έχει γενικά χαμηλότερο αποτύπωμα μνήμης και μπορεί να βελτιώσει modestly την απόδοση. Τεκμηρίωση PyTorch optimizer zero_grad.

Βήμα 3: Διατηρήστε ένα μεγαλύτερο αποτελεσματικό batch με AMP και συσσώρευση κλίσεων

Χρησιμοποιήστε Automatic Mixed Precision όταν το μοντέλο το υποστηρίζει

Η Automatic Mixed Precision (AMP) εκτελεί επιλέξιμες λειτουργίες σε χαμηλότερη ακρίβεια διατηρώντας τις λειτουργίες που χρειάζονται μεγαλύτερη εμβέλεια ή ακρίβεια σε κατάλληλους τύπους. Το PyTorch τεκμηριώνει ότι το AMP μπορεί να βελτιώσει την απόδοση και να μειώσει το αποτύπωμα μνήμης για πολλά φορτία εργασίας CUDA, αλλά δεν είναι αριθμητικά κατάλληλο για κάθε μοντέλο. Συγκεκριμένα, το PyTorch προειδοποιεί ότι ορισμένα μοντέλα προ-εκπαιδευμένα σε bfloat16 μπορεί να έχουν overflow σε float16.

Ένα τρέχον μοτίβο εκπαίδευσης CUDA AMP είναι:

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 συνιστά την εκτέλεση του forward pass και της απώλειας υπό autocast, αφήνοντας στη συνέχεια το πλαίσιο autocast πριν από το backward. Αν το float16 παράγει αστάθεια, διερευνήστε αν το bfloat16 υποστηρίζεται και είναι κατάλληλο για το υλικό και το μοντέλο σας, αντί να υποθέσετε ότι όλες οι λειτουργίες μικτής ακρίβειας συμπεριφέρονται το ίδιο.

Χρησιμοποιήστε συσσώρευση κλίσεων όταν χρειάζεστε μεγαλύτερο αποτελεσματικό batch

Η συσσώρευση κλίσεων επεξεργάζεται αρκετά μικρότερα micro-batches πριν ενημερώσει τον optimizer. Αν το micro-batch είναι 4 και συσσωρεύετε 8 βήματα, το αποτελεσματικό batch για μία ενημέρωση optimizer είναι 32 δείγματα ανά worker, υποθέτοντας ότι κάθε micro-batch έχει τέσσερα δείγματα και η ρύθμιση data-parallel δεν αλλάζει αυτή την αριθμητική.

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)

Για παραγωγικό κώδικα, χειριστείτε επίσης ένα τελικό μερικό παράθυρο συσσώρευσης όταν ο αριθμός των batches δεν διαιρείται ακριβώς με το accum_steps. Αν χρησιμοποιείτε κατανεμημένη εκπαίδευση, η συμπεριφορά συγχρονισμού κλίσεων μπορεί να αλλάξει το trade-off μνήμης/απόδοσης, οπότε ακολουθήστε τις οδηγίες συσσώρευσης του κατανεμημένου API αντί να αντιγράψετε έναν βρόχο μονής GPU χωρίς αλλαγές.

Βήμα 4: Ανταλλάξτε υπολογισμό για μνήμη, στη συνέχεια διερευνήστε τη διατήρηση και τον κατακερματισμό

Activation checkpointing

Το activation checkpointing μειώνει τη μνήμη μη διατηρώντας επιλεγμένες forward ενεργοποιήσεις ζωντανές μέχρι το backward. Αντίθετα, το PyTorch τις υπολογίζει ξανά κατά το backward. Αυτό ανταλλάσσει επιπλέον υπολογισμό για χαμηλότερο αποτύπωμα μνήμης ενεργοποιήσεων. Η τρέχουσα τεκμηρίωση checkpoint του PyTorch συνιστά τη ρητή παράδοση του use_reentrant=False. Τεκμηρίωση PyTorch activation checkpointing.

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 σε στρώματα με μεγάλες αποθηκευμένες ενεργοποιήσεις και αποδεκτό κόστος επανυπολογισμού. Μην υποθέσετε ότι το checkpointing κάθε λειτουργίας είναι βέλτιστο· μπορεί να επιβραδύνει σημαντικά την εκπαίδευση.

Αναζητήστε tensors που κρατούν ζωντανά γραφήματα υπολογισμού

Αν η μνήμη αυξάνεται σε κάθε iteration αντί να κορυφώνεται περίπου στο ίδιο επίπεδο, ελέγξτε τις αναφορές Python. Ένα κοινό μοτίβο είναι η αποθήκευση tensors συνδεδεμένων με γράφημα σε μια λίστα:

# Επικίνδυνο αν διατηρείται για πολλά βήματα:
loss_history.append(loss)

# Αποθηκεύστε έναν αριθμό Python αντίθετα:
loss_history.append(loss.item())

Το ίδιο πρόβλημα μπορεί να συμβεί όταν cacheάρετε εξόδους μοντέλου, χάρτες προσοχής (attention maps), κρυφές καταστάσεις ή tensors επικύρωσης χωρίς να τα αποσυνδέσετε (detach) ή να τα μετακινήσετε εκτός GPU. Διαγράψτε αναφορές που δεν χρειάζεστε πλέον και χρησιμοποιήστε detach() μόνο όταν θέλετε σκόπιμα ένα tensor αποσυνδεδεμένο από το autograd.

Ρυθμίστε τον allocator μόνο αφού τα στατιστικά δείξουν κατακερματισμό

Η τρέχουσα τεκμηρίωση του PyTorch προτιμά τη μεταβλητή περιβάλλοντος PYTORCH_ALLOC_CONF. Η παλαιότερη PYTORCH_CUDA_ALLOC_CONF παραμένει alias για backward compatibility. Αυτή η λεπτομέρεια ονοματολογίας άλλαξε στην τρέχουσα τεκμηρίωση, οπότε οι νέες ρυθμίσεις πρέπει να χρησιμοποιούν το προτιμώμενο όνομα. Μεταβλητές περιβάλλοντος PyTorch CUDA.

Δύο επιλογές allocator είναι ιδιαίτερα σχετικές:

• Το expandable_segments:True είναι πειραματικό και σχεδιάστηκε για να μειώνει τα άχρηστα κομμάτια μνήμης όταν αλλάζουν τα μεγέθη κατανομής, όπως σε φορτία εργασίας των οποίων τα μεγέθη batch ή tensor ποικίλλουν.
• Το max_split_size_mb μπορεί να μειώσει τον κατακερματισμό με τον εγγενή allocator, αλλά το PyTorch το περιγράφει ρητά ως last resort για φορτία εργασίας που αποτυγχάνουν με OOM ενώ εμφανίζουν μεγάλη ποσότητα ανενεργών split blocks. Μπορεί επίσης να βλάψει την απόδοση και αγνοείται από το backend cudaMallocAsync.

# Παράδειγμα για φορτίο εργασίας με ποικίλα μεγέθη κατανομής:
export PYTORCH_ALLOC_CONF=expandable_segments:True

Μην αντιγράφετε flags allocator από άλλη μηχανή χωρίς να ελέγξετε το memory_summary() ή ένα στιγμιότυπο. Ένα πραγματικό πρόβλημα χωρητικότητας—όπου τα ενεργά tensors γεμίζουν ήδη την GPU—δεν θα λυθεί με ρύθμιση κατακερματισμού.

Όταν μία GPU δεν μπορεί ακόμα να χωρέσει το μοντέλο

Αν ένα μόνο δείγμα στο μικρότερο πρακτικό μέγεθος εισόδου προκαλεί ακόμα OOM, το πρόβλημα μπορεί να είναι το μοντέλο και η κατάσταση του optimizer παρά το batch. Σε αυτό το σημείο, εξετάστε μια μικρότερη αρχιτεκτονική, παραμέτρους χαμηλότερης ακρίβειας όπου είναι αριθμητικά κατάλληλο, στρατηγικές CPU/offload ή κατανεμημένη εκπαίδευση με sharding.

Το Fully Sharded Data Parallel (FSDP) του PyTorch μπορεί να κάνει shard παραμέτρων μοντέλου σε workers data-parallel, και η στρατηγική FULL_SHARD κάνει επίσης shard κλίσεων και καταστάσεων optimizer. Αυτό μπορεί να μειώσει τη μνήμη ανά GPU σε σύγκριση με την πλήρως αναπαραγόμενη data parallelism, με το κόστος της επικοινωνίας και πιο σύνθετης συμπεριφοράς εκπαίδευσης. Τεκμηρίωση PyTorch FSDP.

Πρακτική σειρά ενεργειών

Λίστα ελέγχου: πώς να ξέρετε ότι το OOM έχει πράγματι διορθωθεί

• Εκτελέστε αρκετά αντιπροσωπευτικά βήματα εκπαίδευσης, όχι μόνο ένα επιτυχημένο forward pass.
• Επαναφέρετε και καταγράψτε το max_memory_allocated() ώστε να γνωρίζετε τη νέα κορυφή.
• Επιβεβαιώστε ότι η μνήμη GPU φτάνει σε ένα σταθερό εύρος αντί να αυξάνεται σε κάθε iteration.
• Επικυρώστε την απώλεια και τις κλίσεις μετά την ενεργοποίηση μικτής ακρίβειας.
• Επιβεβαιώστε ότι η συσσώρευση κλίσεων διατηρεί το χρονοδιάγραμμα ενημέρωσης optimizer που προορίζατε.
• Εκτελέστε ένα πέρασμα επικύρωσης υπό torch.no_grad() όταν οι κλίσεις δεν απαιτούνται.
• Αν αλλάξατε ρυθμίσεις allocator, συγκρίνετε τα στατιστικά μνήμης και την απόδοση πριν και μετά.
• Μην θεωρείτε το πρόβλημα λυμένο μόνο επειδή το nvidia-smi δείχνει λιγότερη δεσμευμένη μνήμη μετά το empty_cache(). Το ίδιο το φορτίο εργασίας εκπαίδευσης πρέπει να ολοκληρωθεί στην κανονική του κορυφή.

Ένα CUDA OOM είναι καλύτερο να αντιμετωπίζεται ως πρόβλημα προϋπολογισμού μνήμης, όχι ως ένα μεμονωμένο bug του PyTorch. Μετρήστε την κορυφή, μειώστε πρώτα το ζωντανό σύνολο εργασίας, στη συνέχεια χρησιμοποιήστε μικτή ακρίβεια, συσσώρευση και checkpointing ως σκόπιμα trade-offs. Μεταβείτε στη ρύθμιση του allocator μόνο όταν τα στατιστικά του allocator υποδεικνύουν κατακερματισμό, και μεταβείτε σε sharding ή διαφορετικό μοντέλο όταν το μοντέλο δεν χωράει πλέον άνετα σε μία GPU.