Gradient checkpointing

Che cos'è

Un passaggio di addestramento ha due metà. Nel forward pass, ogni layer trasforma il proprio input in un output, producendo una serie di risultati intermedi lungo il percorso: le attivazioni. Il backward pass lavora poi attraverso i layer a ritroso per calcolare i gradienti, e i gradienti di ogni layer necessitano delle attivazioni di quel layer. L'approccio semplice le mantiene tutte in memoria finché il backward pass non le raggiunge. Con esempi lunghi, le attivazioni possono occupare più memoria dei pesi stessi del modello.

Il gradient checkpointing, noto anche come activation checkpointing o ricalcolo, conserva solo alcune attivazioni, i checkpoint, e scarta le restanti. Quando il backward pass necessita di un'attivazione scartata, esegue nuovamente il forward pass dal checkpoint più vicino per ricrearla. Chen et al. (2016) hanno dimostrato che questo riduce la memoria delle attivazioni per una rete a n layer all'ordine di √n, al costo di circa un forward pass aggiuntivo per batch. Nei loro esperimenti, una rete a 1.000 layer è passata da 48 GB a 7 GB con un aumento del tempo di esecuzione di circa il 30%.

Come funziona in un modello linguistico

Nelle librerie dei transformer, i checkpoint si trovano solitamente ai confini dei layer. L'input di ogni layer viene conservato; tutto ciò che viene calcolato all'interno del layer viene scartato e ricostruito quando il backward pass vi giunge. Pertanto, in qualsiasi momento esiste solo il set completo di attivazioni di un singolo layer, anziché di tutti.

activations per GPT-style layer, 16-bit, efficient attention   (Korthikanti et al., 2022)
  ≈ 34 × batch × tokens × hidden size   bytes

illustrative 2B-class model: hidden size 2,048 · 28 layers · batch 2 × 1,024 tokens
  all kept                 34 × 2 × 1,024 × 2,048 × 28   ≈ 4.0 GB
  checkpointed, inputs      2 × 2 × 1,024 × 2,048 × 28   ≈ 0.23 GB
  + one layer rebuilt at a time                          ≈ 0.14 GB

Il numero di checkpoint da conservare è un parametro regolabile. Conservarne di più significa dover ricalcolare meno dati ma occupare più memoria; conservarne di meno comporta l'opposto. Il checkpointing a ogni confine di layer è l'impostazione predefinita abituale poiché è semplice e rimuove già la maggior parte della memoria delle attivazioni: ciò che rimane è un piccolo input per layer più il singolo layer in fase di ricostruzione.

Il costo è in termini di calcolo. Il forward pass viene di fatto eseguito due volte: una volta normalmente e una volta pezzo per pezzo durante il backward pass. Il backward pass costa già circa il doppio del forward pass, quindi il lavoro aggiuntivo è al massimo circa un terzo di uno step; in pratica, l'addestramento risulta circa il 20–30% più lento. Il risultato non cambia: i valori ricalcolati sono gli stessi, quindi la curva di loss corrisponde a quella di un'esecuzione senza checkpointing, al netto di minime differenze numeriche.

Perché i vocabolari ampi lo rendono più importante

Alla fine del modello, ogni posizione in ogni esempio ottiene un punteggio per ogni token nel vocabolario: i logit. La loro dimensione non dipende dalla profondità del modello, ma solo dal numero di posizioni addestrate e dall'ampiezza del vocabolario. I modelli più recenti hanno vocabolari da 150.000 a 250.000 token, pertanto i soli logit possono occupare gigabyte.

logits = batch × tokens × vocabulary size × bytes per number

2 × 1,024 × 248,000 × 4 bytes (32-bit)   ≈ 2.0 GB   before their gradients

Usarlo bene

• Da preferire rispetto alla riduzione dei dati. Attivare il checkpointing costa tempo; ridurre la lunghezza massima taglia la parte finale degli esempi lunghi. Quando la memoria è il limite, il rallentamento è solitamente il compromesso migliore.
• Utilizzare la memoria liberata. Un batch più grande o una lunghezza massima superiore possono recuperare parte della velocità persa per epoca.
• Combinarlo con QLoRA per i modelli più grandi che una macchina può supportare: QLoRA riduce i pesi, il checkpointing riduce le attivazioni.
• Sapere cosa non viene modificato. Pesi, gradienti e stato dell'ottimizzatore mantengono la stessa dimensione, e non ha alcun effetto sull'esecuzione del modello finito.
• Controllare le impostazioni predefinite. Molte procedure di fine-tuning lo attivano già. Se la memoria è sufficiente e la velocità è importante, disattivarlo rende ogni step più rapido.

Errori comuni

• Ridurre prima la lunghezza massima, tagliando le estremità degli esempi lunghi, quando il solo checkpointing li avrebbe fatti rientrare nei limiti.
• Aspettarsi che riduca i pesi o lo stato dell'ottimizzatore. Influisce solo sulle attivazioni.
• Interpretare il rallentamento del 20–30% come un difetto della macchina. È il prezzo della memoria risparmiata.
• Presumere che elimini il costo dei logit. Con un vocabolario enorme, potrebbero comunque essere necessari un batch o una lunghezza massima inferiori.

Riferimenti e letture di approfondimento

Le fonti esterne si aprono in una nuova scheda.

1. 01Training Deep Nets with Sublinear Memory Cost↗Chen et al., 2016 · arXiv · arxiv.orgIl paper che ha reso il ricalcolo uno standard: memoria √n per un forward pass aggiuntivo.
2. 02Reducing Activation Recomputation in Large Transformer Models↗Korthikanti et al., 2022 · arXiv · arxiv.orgDa dove deriva la stima delle attivazioni per layer e come ricalcolare solo le parti più dispendiose.
3. 03torch.utils.checkpoint↗PyTorch docs · docs.pytorch.org
4. 04GPU memory usage↗Hugging Face Transformers docs · huggingface.coPesi, gradienti, stato dell'ottimizzatore e attivazioni: dove viene allocata la memoria di addestramento.

Concetti correlati