Come funziona?
Come ho detto, durante l'addestramento su più GPU, ogni processo ha copie esatte degli stessi dati durante l'addestramento con DDP. Possiamo ottimizzarlo, implementando diversi miglioramenti:
Stato dell'ottimizzatore shard (ZeRO 1)
Durante l'addestramento con DDP, ogni processo conserva una copia completa degli stati dell'ottimizzatore. Con ZeRO1, suddividiamo questi stati di ottimizzazione su tutti i ranghi in modo tale che ciascun rango contenga solo una parte degli stati di ottimizzazione. Durante il passaggio all'indietro, ogni rango deve solo raccogliere gli stati dell'ottimizzatore rilevanti per i suoi parametri per effettuare una fase di ottimizzazione. Questa riduzione della ridondanza aiuta a conservare la memoria.
💡 Nel caso di Adam, che contiene parametri a circa il doppio della dimensione del modello, suddividere lo stato dell'ottimizzatore tra 8 ranghi significa che ciascun rango memorizza solo un quarto (2/8) della dimensione totale dello stato.
Gradienti di frammenti (ZeRO 2)
Shardiamo gli stati dell'ottimizzatore. Ora modificheremo il passaggio dell'ottimizzatore anche per frammentare i gradienti. Se un rango ha stati di ottimizzazione per una parte di parametri, allora:
- aggregare tutti i gradienti rilevanti per gli stati detenuti dal rango
- calcolare la fase di ottimizzazione
- invia la fase di ottimizzazione per una parte dei parametri a tutti gli altri ranghi
Come hai notato, ora non è necessario che ogni rango contenga una replica completa dei gradienti. Possiamo inviare i gradienti a un rango pertinente non appena sono disponibili. Quindi, possiamo ridurre ulteriormente il consumo massimo di memoria.
Parametri del modello Shard (ZeRO 3)
Sta per diventare epico.
Perché dobbiamo archiviare una copia completa del modello su ogni rango? Suddividiamo i parametri del modello tra tutti i ranghi. Quindi, recupereremo i parametri richiesti giusto in tempo durante le operazioni avanti e indietro.
💡 In caso di modelli di grandi dimensioni, queste ottimizzazioni possono ridurre drasticamente il consumo di memoria
Come utilizzare il PFSDP?
Abbastanza semplice in realtà. Tutto ciò di cui abbiamo bisogno è avvolgere il modello con FSDP:
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributed.fsdp import FullyShardedDataParallel as FSDPmodel = FSDP(model)
# it's critical to get parameters from the wrapped model
# as only a portion of them returned (sharded part)
optimizer = optim.Adam(model.parameters())
# consuct training as usual
train(model, optimizer)
Puoi anche specificare la strategia di partizionamento orizzontale di FSDP. Ad esempio, possiamo selezionare il SHARD_GRAD_OP strategia per ottenere un comportamento simile a quello di ZeRO2. Puoi conoscere altre strategie qui:
Inoltre, puoi eseguire il wrapper con i sottomoduli FSDP. Nell'esempio precedente viene utilizzato un solo modulo FSDP, il che ridurrà l'efficienza di calcolo e l'efficienza della memoria. Il modo in cui funziona è che, supponiamo che il tuo modello contenga 100 livelli lineari. Se esegui FSDP (modello), ci sarà solo un'unità FSDP che avvolge l'intero modello. In tal caso, l'allgather raccoglierebbe i parametri completi per tutti i 100 livelli lineari e quindi non salverà la memoria CUDA per lo sharding dei parametri.
È possibile racchiudere i sottomoduli in modo esplicito o definire una politica di wrap automatico. Per saperne di più su FSDP, leggi la guida PyTorch:
Fonte: towardsdatascience.com
