Gradient Accumulation

L'accumulo del gradiente è una tecnica di addestramento che esegue diversi batch piccoli, detti micro-batch, nel passaggio in avanti e nella retropropagazione, somma i loro gradienti e aggiorna solo successivamente i pesi del modello. L'aggiornamento approssima un singolo passaggio su un batch grande quanto tutti i micro-batch combinati, mentre la GPU conserva in memoria solo le attivazioni di un micro-batch alla volta. In questo modo un dispositivo con memoria limitata può addestrare il modello con la dimensione di batch prevista dalla configurazione, senza ridurre la risoluzione delle immagini o le dimensioni del modello. La dimensione effettiva del batch è pari alla dimensione del micro-batch moltiplicata per il numero di passaggi di accumulo e per il numero di dispositivi; quindi, micro-batch di 4 immagini accumulati per 8 passaggi su 4 GPU equivalgono a un batch di 128.

Come funziona l'accumulo dei gradienti#

In PyTorch e nella maggior parte degli altri framework, ogni chiamata a backward() aggiunge nuovi gradienti a quelli già memorizzati sui parametri, invece di sostituirli. L'accumulo sfrutta direttamente questo comportamento: esegui il passaggio in avanti e quello all'indietro per ciascun micro-batch, chiama l'ottimizzatore una volta raggiunto il numero stabilito di micro-batch, poi azzera i gradienti e inizia il ciclo successivo:

La maggior parte delle funzioni di perdita calcola la media sui campioni di un batch, quindi la perdita di ogni micro-batch viene divisa per il numero di passaggi per trasformare la somma dei gradienti in una media. Con questa riscalatura, quattro micro-batch da 8 producono lo stesso gradiente di un batch da 32:

import torch
from torch import nn

torch.manual_seed(0)
model = nn.Linear(16, 1)
x, y = torch.randn(32, 16), torch.randn(32, 1)
loss_fn = nn.MSELoss()

# Reference: one batch of 32 samples
loss_fn(model(x), y).backward()
batch_grad = model.weight.grad.clone()
model.zero_grad()

# Four micro-batches of 8: backward() adds into .grad, and dividing each loss by 4 turns the sum into a mean
steps = 4
for x_micro, y_micro in zip(x.chunk(steps), y.chunk(steps)):
    (loss_fn(model(x_micro), y_micro) / steps).backward()

print("Matches the full-batch gradient:", torch.allclose(model.weight.grad, batch_grad, atol=1e-6))
print(f"Largest difference: {(model.weight.grad - batch_grad).abs().max():.1e}")

Stampa Matches the full-batch gradient: True, con una differenza massima di 1.6e-07, dovuta all'arrotondamento in virgola mobile quando si sommano gli stessi numeri in un ordine diverso.

In cosa l'accumulo differisce da un batch grande#

L'equivalenza vale per le funzioni di perdita che calcolano la media di campioni indipendenti, ma viene meno in tre casi comuni:

• Batch normalization: calcola la media e la varianza di ogni micro-batch, quindi le statistiche ottenute da 4 immagini sono più rumorose di quelle calcolate su 128.
• Sequenze di lunghezza variabile: quando un modello linguistico calcola la media della perdita sui token di ogni micro-batch, i micro-batch con pochi token hanno lo stesso peso di quelli più lunghi.
• Clipping del gradiente: va applicato una sola volta al gradiente accumulato, subito prima del passaggio dell'ottimizzatore; applicarlo a ogni micro-batch modifica il risultato.

Hugging Face e Unsloth hanno risolto la discrepanza nel conteggio dei token nell'ottobre 2024 per la messa a punto dei modelli linguistici di grandi dimensioni (LLM), dividendo la somma delle perdite per il numero totale di token dell'intera finestra di accumulo. L'accumulo non offre neppure un vantaggio in termini di velocità: lo stesso numero di campioni continua comunque a passare attraverso la rete un micro-batch alla volta, e i micro-batch di piccole dimensioni sfruttano la GPU in modo meno efficiente. Nell'addestramento distribuito, DistributedDataParallel di PyTorch offre un contesto no_sync() che salta l'all-reduce dei gradienti per tutti i micro-batch tranne l'ultimo, riducendo così il traffico di comunicazione. Un batch effettivo più grande richiede in genere un tasso di apprendimento più elevato: i ricercatori di Facebook hanno addestrato ResNet-50 su ImageNet in un'ora con batch di 8.192 immagini su 256 GPU, scalando linearmente il tasso di apprendimento in base alla dimensione del batch e aumentando gradualmente il tasso di apprendimento durante la fase iniziale di riscaldamento.

Tecniche correlate#

• Checkpointing del gradiente: ricalcola le attivazioni durante la retropropagazione invece di memorizzarle, risparmiando memoria per ogni micro-batch al costo di un maggiore calcolo.
• Precisione mista: memorizza le attivazioni in formati a 16 bit, lasciando spazio a micro-batch più grandi. Lo scaler della perdita ripristina la scala dei gradienti una volta sola, dopo l'ultimo micro-batch.
• Parallelismo a pipeline: suddivide ogni batch in micro-batch che attraversano le fasi del modello su GPU diverse e ne accumula i gradienti prima di un singolo passaggio dell'ottimizzatore.

Per impostazione predefinita, l'addestramento di Ultralytics YOLO26 accumula i gradienti: in modalità di addestramento, il trainer esegue un passaggio dell'ottimizzatore ogni round(nbs / batch) batch, quindi i valori predefiniti batch=16 e la dimensione nominale del batch nbs=64 producono quattro batch per aggiornamento, mentre batch=8 ne produce otto. Il numero aumenta gradualmente da uno durante il riscaldamento e il decadimento dei pesi viene scalato di batch × accumulation / nbs, quindi gli iperparametri predefiniti, impostati per un batch di 64 immagini, restano validi anche su una GPU che può contenere solo poche immagini alla volta.