Gradient Checkpointing

Il gradient checkpointing è una tecnica di addestramento che consente di risparmiare memoria, memorizzando solo attivazioni intermedie selezionate dal passaggio in avanti e ricalcolando le altre durante la backpropagation. Chiamato anche activation checkpointing, scambia calcoli aggiuntivi con un minore utilizzo di memoria di picco. Nonostante il nome, la tecnica memorizza i checkpoint delle attivazioni anziché i gradienti dei parametri o i file del modello, risultando particolarmente utile quando i tensori di attivazione impediscono a una rete neurale di entrare nella GPU memory disponibile.

Come funziona il Gradient Checkpointing#

Durante un normale passaggio in avanti, una rete neurale calcola tensori intermedi chiamati attivazioni. Il sistema di differenziazione automatica conserva le attivazioni necessarie per calcolare i gradienti in un secondo momento, come descritto nelle PyTorch autograd mechanics. Reti profonde, batch di grandi dimensioni, immagini ad alta risoluzione e lunghe sequenze di input possono fare in modo che questi tensori salvati consumino una memoria considerevole.

Il gradient checkpointing divide la rete in segmenti:

1. Il passaggio in avanti memorizza gli input o le attivazioni dei confini per i segmenti selezionati.
2. Altre attivazioni intermedie all'interno di tali segmenti vengono scartate.
3. Durante il passaggio all'indietro, ogni segmento con checkpoint viene eseguito nuovamente in avanti per ricostruire i valori mancanti.
4. Le attivazioni ricostruite vengono utilizzate immediatamente per calcolare i gradienti.

La PyTorch activation checkpointing API espone questo comportamento tramite torch.utils.checkpoint. Concetti equivalenti compaiono come JAX gradient checkpointing and rematerialization e TensorFlow tape checkpointing.

La posizione dei checkpoint determina il compromesso. L'inserimento di checkpoint in un numero maggiore di regioni generalmente fa risparmiare più memoria ma ripete più operazioni. Una posizione selettiva attorno ai blocchi ad alto contenuto di attivazione può fornire un bilanciamento migliore rispetto al ricalcolo dell'intera rete.

Compromessi e tecniche correlate#

Il gradient checkpointing di solito lascia immutati i parametri del modello, i gradienti e gli stati dell'optimizer. Il suo obiettivo principale è la memoria delle attivazioni e l'addestramento diventa più lento perché alcuni calcoli in avanti vengono eseguiti due volte. Il risultato esatto dipende dall'architettura, dai confini dei checkpoint, dalla forma del batch e dall'hardware.

Differisce da diverse tecniche correlate:

• Gradient accumulation elabora più microbatch prima di aggiornare i pesi del modello, creando un batch effettivo più grande senza caricare ogni campione contemporaneamente. Il gradient checkpointing invece riduce le attivazioni conservate per ogni microbatch.
• Mixed precision utilizza tipi di dati a precisione inferiore per operazioni selezionate. Il PyTorch automatic mixed precision workflow può ridurre la memoria e accelerare le operazioni compatibili, mentre il checkpointing aggiunge deliberatamente calcoli.
• La riduzione della Batch size riduce la memoria elaborando meno campioni insieme. Il checkpointing può consentire a un batch o a una risoluzione di input più grandi di rimanere fattibili.
• Training checkpoints salva i pesi e lo stato dell'optimizer per il recupero o per una successiva inferenza. La PyTorch model checkpoint guide descrive questo meccanismo di persistenza, che non è correlato al ricalcolo delle attivazioni.

Il codice con checkpoint dovrebbe essere funzionalmente coerente tra le sue esecuzioni in avanti originali e ricalcolate. Uno stato mutabile, trasferimenti di dispositivi o casualità non controllata all'interno di una regione con checkpoint possono causare errori o gradienti non corretti. Nei transformer in stile decoder, un key-value cache potrebbe anche dover essere disabilitato durante l'addestramento con checkpoint perché lo stato di inferenza memorizzato nella cache può entrare in conflitto con la ricostruzione del grafo in avanti.

Esempio in PyTorch#

Il seguente esempio inserisce un checkpoint in un blocco a elevato utilizzo di memoria durante un passaggio di addestramento:

import torch
from torch import nn
from torch.utils.checkpoint import checkpoint

torch.manual_seed(0)

block = nn.Sequential(
    nn.Linear(1024, 4096),
    nn.ReLU(),
    nn.Linear(4096, 1024),
)
optimizer = torch.optim.AdamW(block.parameters())
inputs = torch.randn(8, 1024, requires_grad=True)
targets = torch.zeros_like(inputs)

optimizer.zero_grad(set_to_none=True)
outputs = checkpoint(block, inputs, use_reentrant=False)
loss = nn.functional.mse_loss(outputs, targets)
loss.backward()
optimizer.step()

Vengono conservati solo gli input del blocco e le informazioni sui confini richieste; le sue attivazioni interne vengono ricostruite durante loss.backward(). I progetti reali dovrebbero confrontare le esecuzioni con e senza checkpoint utilizzando una misurazione come PyTorch peak allocated GPU memory.

Applicazioni nel mondo reale#

• Computer vision ad alta risoluzione: la segmentazione medica, il rilevamento aereo e l'ispezione industriale possono essere addestrati su immagini di grandi dimensioni le cui feature map consumano più memoria rispetto ai pesi del modello. L'inserimento di checkpoint in fasi selezionate del backbone può preservare la risoluzione dell'immagine o consentire campioni aggiuntivi per batch anziché ridurre drasticamente la scala degli input.

• Addestramento di transformer a sequenza lunga: l'archiviazione delle attivazioni cresce rapidamente man mano che la lunghezza della sequenza e il numero di layer aumentano. Il ricalcolo dei blocchi dei transformer può far sì che contesti più lunghi o microbatch più grandi rientrino sullo stesso acceleratore. La NVIDIA activation recomputation guide illustra il ricalcolo completo e selettivo per i layer dei transformer.

Guida pratica#

Usa il gradient checkpointing quando la profilazione mostra che le attivazioni, anziché i parametri o lo stato dell'optimizer, dominano la memoria. Inizia con grandi blocchi ripetuti e fai il benchmark della memoria di picco, del tempo di iterazione e del comportamento di validazione prima di espandere la copertura.

Per le Ultralytics YOLO training configurations documentate, i controlli di memoria di prima linea includono precisione mista automatica, dimensione dell'immagine e batch size fisica. Il Ultralytics AutoBatch reference spiega la selezione automatica del batch in base alla memoria GPU disponibile. Quando l'hardware locale rimane insufficiente, Ultralytics Platform cloud training fornisce GPU cloud configurabili per esecuzioni di addestramento gestite. Il gradient checkpointing può integrare questi controlli quando un'architettura PyTorch personalizzata richiede una gestione più fine della memoria delle attivazioni.