Accelerare l’addestramento di modelli di grandi dimensioni utilizzando PyTorch Fully Sharded Data Parallel

'Accelerating training of large-scale models using PyTorch Fully Sharded Data Parallel'

In questo post vedremo come possiamo sfruttare la libreria Accelerate per allenare modelli di grandi dimensioni che consentono agli utenti di sfruttare le ultime funzionalità di PyTorch FullyShardedDataParallel (FSDP).

Con l’aumento costante della scala, delle dimensioni e dei parametri dei modelli di Machine Learning (ML), i praticanti di ML trovano difficile allenare o persino caricare tali modelli di grandi dimensioni sul proprio hardware. Da un lato, si è scoperto che i modelli di grandi dimensioni imparano rapidamente (efficienti in termini di dati e calcolo) e sono significativamente più performanti rispetto ai modelli più piccoli [1]; d’altra parte, diventa proibitivo allenare tali modelli sulla maggior parte dell’hardware disponibile.

L’allenamento distribuito è la chiave per consentire l’allenamento di tali modelli di ML di grandi dimensioni. Ci sono state importanti avanzate recenti nel campo dell’Allenamento Distribuito su Larga Scala. Di seguito sono riportati alcuni dei progressi più notevoli:

1. Parallelismo dei dati utilizzando ZeRO – Zero Redundancy Optimizer [2]
   1. Fase 1: Shards degli stati dell’ottimizzatore tra i lavoratori/GPU paralleli ai dati
   2. Fase 2: Shards degli stati dell’ottimizzatore + gradienti tra i lavoratori/GPU paralleli ai dati
   3. Fase 3: Shards degli stati dell’ottimizzatore + gradienti + parametri del modello tra i lavoratori/GPU paralleli ai dati
   4. Scarico su CPU: Scarica i gradienti + gli stati dell’ottimizzatore sulla CPU basandosi su ZERO Fase 2 [3]
2. Parallelismo dei tensori [4]: Forma di parallelismo del modello in cui viene effettuata la suddivisione dei parametri di singoli strati con un numero elevato di parametri tra gli acceleratori/GPU in modo intelligente per ottenere il calcolo parallelo evitando costosi sincronizzazioni di comunicazione.
3. Parallelismo del pipeline [5]: Forma di parallelismo del modello in cui i diversi strati del modello vengono distribuiti su diversi acceleratori/GPU e viene utilizzata la tecnica del pipelining per mantenere tutti gli acceleratori in esecuzione contemporaneamente. Qui, ad esempio, il secondo acceleratore/GPU elabora il primo micro-batch mentre il primo acceleratore/GPU elabora il secondo micro-batch.
4. Parallelismo 3D [3]: Utilizza il parallelismo dei dati utilizzando ZERO + il parallelismo dei tensori + il parallelismo del pipeline per allenare modelli enormi nell’ordine dei miliardi di parametri. Ad esempio, il modello di linguaggio BigScience con 176 miliardi di parametri utilizza questo approccio [6].

In questo post daremo uno sguardo al parallelismo dei dati utilizzando ZeRO e più specificamente all’ultima funzionalità di PyTorch, FullyShardedDataParallel (FSDP). DeepSpeed e FairScale hanno implementato le idee di base del paper ZeRO. Questi sono già stati integrati nel Trainer di transformers e sono accompagnati da un ottimo blog “Fit More and Train Faster With ZeRO via DeepSpeed and FairScale” [10]. PyTorch ha recentemente incorporato upstream il Fairscale FSDP in PyTorch Distributed con ulteriori ottimizzazioni.

• Generazione di testo più veloce con TensorFlow e XLA
• Proximal Policy Optimization (PPO)
• Benvenuti a fastai nell’Hugging Face Hub

Esamineremo il compito di Modellazione del Linguaggio Causale utilizzando le varianti del modello GPT-2 Large (762M) e XL (1.5B).

Di seguito è riportato il codice per il pre-allenamento del modello GPT-2. È simile all’esempio ufficiale di modellazione del linguaggio causale qui con l’aggiunta di 2 argomenti n_train (2000) e n_val (500) per evitare la pre-elaborazione/l’allenamento sull’intero set di dati al fine di eseguire velocemente benchmark di proof of concept.

run_clm_no_trainer.py

Configurazione FSDP di esempio dopo l’esecuzione del comando accelerate config:

compute_environment: LOCAL_MACHINE
deepspeed_config: {}
distributed_type: FSDP
fsdp_config:
  min_num_params: 2000
  offload_params: false
  sharding_strategy: 1
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: 'no'
num_machines: 1
num_processes: 2
use_cpu: false

Multi-GPU FSDP

Qui, sperimentiamo l’impostazione Single-Node Multi-GPU. Confrontiamo le prestazioni di Distributed Data Parallel (DDP) e FSDP in diverse configurazioni. Prima, viene utilizzato il modello GPT-2 Large (762M) in cui DDP funziona con determinate dimensioni di batch senza generare errori Out Of Memory (OOM). Successivamente, viene utilizzato il modello GPT-2 XL (1.5B) in cui DDP fallisce con l’errore OOM anche con una dimensione di batch pari a 1. Osserviamo che FSDP consente dimensioni di batch più grandi per il modello GPT-2 Large e consente l’allenamento del modello GPT-2 XL con una dimensione di batch decente, a differenza di DDP.

Configurazione hardware: 2X24GB NVIDIA Titan RTX GPU.

Comando per l’addestramento del modello GPT-2 Large (762M parametri):

export BS=#`provare con diverse dimensioni di batch finché non si ottiene un errore OOM,
#cioè iniziare con una dimensione di batch più grande e diminuire finché non entra nella GPU`

time accelerate launch run_clm_no_trainer.py \
--model_name_or_path gpt2-large \
--dataset_name wikitext \
--dataset_config_name wikitext-2-raw-v1 \
--per_device_train_batch_size $BS 
--per_device_eval_batch_size $BS 
--num_train_epochs 1 
--block_size 12

Esempio di esecuzione FSDP:

Tabella 1: Benchmarking FSDP sul modello GPT-2 Large (762M)

Rispetto a DDP, dalla Tabella 1 possiamo osservare che FSDP consente dimensioni di batch più grandi, fino a 2X-3X senza e con l’impostazione di CPU offload, rispettivamente. In termini di tempo di addestramento, DDP con precisione mista è il più veloce seguito da FSDP utilizzando ZERO Stage 2 e Stage 3, rispettivamente. Poiché il compito di modellazione del linguaggio causale ha sempre una lunghezza di sequenza di contesto fissa (–block_size), il miglioramento del tempo di addestramento con FSDP non è stato così grande. Per le applicazioni con batching dinamico, FSDP, che consente dimensioni di batch più grandi, avrà probabilmente un notevole miglioramento in termini di tempo di addestramento. Attualmente il supporto di FSDP per la precisione mista presenta alcune problematiche con il transformer. Una volta risolto questo problema, il miglioramento del tempo di addestramento sarà ulteriormente notevole.

Offloading della CPU per consentire l’addestramento di modelli enormi che non entrano nella memoria della GPU

Comando per l’addestramento del modello GPT-2 XL (1.5B parametri):

export BS=#`provare con diverse dimensioni di batch finché non si ottiene un errore OOM,
#cioè iniziare con una dimensione di batch più grande e diminuire finché non entra nella GPU`

time accelerate launch run_clm_no_trainer.py \
--model_name_or_path gpt2-xl \
--dataset_name wikitext \
--dataset_config_name wikitext-2-raw-v1 \
--per_device_train_batch_size $BS 
--per_device_eval_batch_size $BS 
--num_train_epochs 1 
--block_size 12

Tabella 2: Benchmarking FSDP sul modello GPT-2 XL (1.5B)

Dalla Tabella 2, possiamo osservare che DDP (con e senza fp16) non riesce nemmeno ad eseguire con una dimensione di batch di 1 e provoca un errore OOM di CUDA. FSDP con Zero-Stage 3 può essere eseguito su 2 GPU con una dimensione di batch di 5 (dimensione di batch effettiva =10 (5 X 2)). FSDP con CPU offload può aumentare ulteriormente la dimensione massima del batch a 14 per GPU quando si utilizzano 2 GPU. FSDP con CPU offload consente di addestrare il modello GPT-2 1.5B su una singola GPU con una dimensione di batch di 10. Ciò consente ai professionisti di ML con risorse di calcolo minime di addestrare modelli così grandi, democratizzando così l’addestramento di modelli di grandi dimensioni.

Capacità e limitazioni dell’integrazione FSDP

Esaminiamo il supporto attuale che Accelerate fornisce per l’integrazione FSDP e le limitazioni note.

Versione PyTorch richiesta per il supporto FSDP: PyTorch Nightly (o 1.12.0 se leggete questo dopo il rilascio) in quanto il salvataggio del modello con FSDP attivato è disponibile solo con correzioni recenti.

Configurazione tramite CLI:

1. Strategia di sharding: [1] FULL_SHARD, [2] SHARD_GRAD_OP
2. Numero minimo di parametri: numero minimo di parametri di FSDP per l’incapsulamento automatico predefinito.
3. Offload dei parametri: decide se trasferire i parametri e i gradienti alla CPU.

Per un maggiore controllo, gli utenti possono sfruttare il FullyShardedDataParallelPlugin in cui possono specificare auto_wrap_policy, backward_prefetch e ignored_modules.

Dopo aver creato un’istanza di questa classe, gli utenti possono passarla durante la creazione dell’oggetto Accelerator.

Per ulteriori informazioni su queste opzioni, si prega di fare riferimento al codice PyTorch FullyShardedDataParallel.

Successivamente, vedremo l’importanza della configurazione min_num_params. Di seguito è riportato un estratto da [8] che illustra l’importanza della politica di avvolgimento automatico di FSDP.

(Fonte: link)

Quando si utilizza la default_auto_wrap_policy, uno strato viene avvolto nel modulo FSDP se il numero di parametri in tale strato è superiore a min_num_params. Il codice per il fine-tuning del modello BERT-Large (330M) sul compito GLUE MRPC è l’esempio completo ufficiale di NLP che illustra come utilizzare correttamente la funzionalità FSDP con l’aggiunta di utility per il monitoraggio dell’utilizzo di memoria picco.

fsdp_with_peak_mem_tracking.py

Sfruttiamo la funzionalità di monitoraggio supportata in Accelerate per registrare l’utilizzo di memoria picco durante l’addestramento e la valutazione insieme alle metriche di valutazione. Di seguito è riportato uno screenshot dei grafici da wandb run.

Possiamo osservare che DDP occupa il doppio di memoria rispetto a FSDP con avvolgimento automatico. FSDP senza avvolgimento automatico occupa più memoria rispetto a FSDP con avvolgimento automatico ma considerevolmente meno rispetto a DDP. FSDP con avvolgimento automatico con min_num_params=2k occupa leggermente meno memoria rispetto all’impostazione con min_num_params=1M. Questo sottolinea l’importanza della politica di avvolgimento automatico di FSDP e gli utenti dovrebbero sperimentare con il parametro min_num_params per trovare l’impostazione che consente di risparmiare considerevolmente memoria senza generare un’eccessiva comunicazione. Il team di PyTorch sta lavorando a uno strumento di auto-tuning per questa configurazione, come menzionato in [8].

Alcune avvertenze da tenere presente

• PyTorch FSDP avvolge automaticamente i sottomoduli, appiattisce i parametri e frammenta i parametri sul posto. A causa di ciò, qualsiasi ottimizzatore creato prima dell’avvolgimento del modello viene interrotto e occupa più memoria. Pertanto, è altamente raccomandato ed efficiente preparare il modello prima di creare l’ottimizzatore. Accelerate avvolgerà automaticamente il modello e creerà un ottimizzatore per voi nel caso di un singolo modello, con un messaggio di avviso.

  Avviso FSDP: Quando si utilizza FSDP, è efficiente e consigliato chiamare il metodo prepare per il modello prima di creare l’ottimizzatore

Tuttavia, di seguito è riportato il modo consigliato per preparare il modello e l’ottimizzatore durante l’utilizzo di FSDP:

model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", return_dict=True)
+ model = accelerator.prepare(model)

optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)

- model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(model,
-        optimizer, train_dataloader, eval_dataloader, lr_scheduler
-    )

+ optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+         optimizer, train_dataloader, eval_dataloader, lr_scheduler
+        )

• Nel caso di un singolo modello, se si è creato un ottimizzatore con più gruppi di parametri e si è chiamato il metodo prepare con essi insieme, i gruppi di parametri saranno persi e verrà visualizzato il seguente avviso:

  Avviso FSDP: Quando si utilizza FSDP, diversi gruppi di parametri verranno fusi in un unico gruppo a causa dell’avvolgimento del modulo nidificato e dell’appiattimento del parametro.

Ciò accade perché i gruppi di parametri creati prima dell’avvolgimento non avranno significato dopo l’avvolgimento a causa dell’appiattimento dei parametri dei moduli FSDP nidificati in array 1D (che possono contenere molti strati). Ad esempio, di seguito sono riportati i parametri nominati del modello FSDP su GPU 0 (quando si utilizzano 2 GPU. Circa 55M (110M/2) parametri in array 1D in quanto questo avrà la 1a frammentazione dei parametri). Qui, se non si è applicato alcun decadimento del peso per i parametri nominati del modello BERT-Base non avvolto, come [bias, LayerNorm.weight], non sarà possibile applicarlo al modello FSDP avvolto di seguito in quanto non vi sono parametri nominati con una di quelle stringhe e i parametri di quei livelli sono concatenati con i parametri di vari altri livelli. Ulteriori dettagli sono menzionati in questo problema (I gradienti dei parametri originali del modello non vengono impostati, il che significa che non possono essere ottimizzati separatamente (motivo per cui non possiamo supportare più gruppi di parametri)).

```
{
'_fsdp_wrapped_module.flat_param': torch.Size([494209]),

'_fsdp_wrapped_module._fpw_module.bert.embeddings.word_embeddings._fsdp_wrapped_module.flat_param': torch.Size([11720448]),

'_fsdp_wrapped_module._fpw_module.bert.encoder._fsdp_wrapped_module.flat_param': torch.Size([42527232])
}
```

• In caso di modelli multipli, è necessario preparare i modelli prima di creare gli ottimizzatori altrimenti si verificherà un errore.

• La precisione mista non è attualmente supportata con FSDP in quanto si attende che PyTorch risolva il supporto per essa.

(Fonte: link )

Il flusso di lavoro sopra riportato fornisce una panoramica di ciò che accade dietro le quinte quando FSDP è attivato. Cerchiamo prima di capire come funziona DDP e come FSDP lo migliora. In DDP, ogni worker/acceleratore/GPU ha una replica di tutti i parametri del modello, i gradienti e gli stati dell’ottimizzatore. Ogni worker riceve un batch di dati diverso, passa attraverso il passaggio in avanti, viene calcolata una perdita seguita dal passaggio all’indietro per generare i gradienti. Ora, viene eseguita un’operazione di all-reduce in cui ogni worker riceve i gradienti dagli altri worker e viene effettuata una media. In questo modo, ogni worker ha ora gli stessi gradienti globali che vengono utilizzati dall’ottimizzatore per aggiornare i parametri del modello. Possiamo vedere che avere repliche complete consuma molta memoria ridondante su ogni GPU, il che limita la dimensione del batch e la dimensione dei modelli.

FSDP affronta precisamente questo frammentando gli stati dell’ottimizzatore, i gradienti e i parametri del modello tra i worker paralleli al dato. Inoltre, facilita il trasferimento della CPU di tutti quei tensori, consentendo così il caricamento di modelli di grandi dimensioni che non rientrano nella memoria GPU disponibile. Similmente a DDP, ogni worker riceve un batch di dati diverso. Durante il passaggio in avanti, se è abilitato l’offload della CPU, i parametri dello shard locale vengono prima copiati sulla GPU/acceleratore. Quindi, ogni worker esegue un’operazione di all-gather per un dato modulo/strato FSDP avvolto per ottenere tutti i parametri necessari, si esegue il calcolo seguito dal rilascio/svuotamento degli shard dei parametri degli altri worker. Questo continua per tutti i moduli FSDP. La perdita viene calcolata dopo il passaggio in avanti e durante il passaggio all’indietro, viene eseguita nuovamente un’operazione di all-gather per ottenere tutti i parametri necessari per un dato modulo FSDP, si esegue il calcolo per ottenere i gradienti locali seguiti dal rilascio degli shard degli altri worker. Ora, i gradienti locali vengono mediati e frammentati a ciascun worker rilevante utilizzando l’operazione di reduce-scatter. Questo consente a ciascun worker di aggiornare i parametri del proprio shard locale. Se l’offload della CPU è attivato, i gradienti vengono passati alla CPU per aggiornare i parametri direttamente sulla CPU.

Si prega di consultare [7, 8, 9] per tutti i dettagli approfonditi sul funzionamento di PyTorch FSDP e gli approfondimenti sperimentali effettuati utilizzando questa funzionalità.

Se riscontri problemi con la parte di integrazione di PyTorch FSDP, apri un problema in accelerate .

Ma se hai problemi con la configurazione e la distribuzione di PyTorch FSDP, devi chiedere agli esperti nei rispettivi domini, perciò, per favore, apri un problema di PyTorch.

[1] Train Large, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers

[2] ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

[3] DeepSpeed: Extreme-scale model training for everyone – Microsoft Research

[4] Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

[5] Introducing GPipe, an Open Source Library for Efficiently Training Large-scale Neural Network Models

[6] Which hardware do you need to train a 176B parameters model?

[7] Introducing PyTorch Fully Sharded Data Parallel (FSDP) API | PyTorch

[8] Getting Started with Fully Sharded Data Parallel(FSDP) — PyTorch Tutorials 1.11.0+cu102 documentation

[9] Training a 1 Trillion Parameter Model With PyTorch Fully Sharded Data Parallel on AWS | by PyTorch | PyTorch | Mar, 2022 | VoAGI

[10] Fit More and Train Faster With ZeRO via DeepSpeed and FairScale