Accelerare l’addestramento di modelli di grandi dimensioni utilizzando DeepSpeed

Accelerate large-scale model training with DeepSpeed.

In questo post esamineremo come possiamo sfruttare la libreria Accelerate per l’addestramento di modelli di grandi dimensioni che consente agli utenti di sfruttare le funzionalità ZeRO di DeeSpeed.

Stanco degli errori di memoria esaurita (OOM) durante il tentativo di addestrare modelli di grandi dimensioni? Abbiamo la soluzione. I modelli di grandi dimensioni sono molto performanti [1], ma difficili da addestrare con l’hardware disponibile. Per sfruttare al meglio l’hardware disponibile per l’addestramento di modelli di grandi dimensioni, è possibile sfruttare il parallelismo dei dati utilizzando ZeRO – Zero Redundancy Optimizer [2].

Di seguito è riportata una breve descrizione del parallelismo dei dati utilizzando ZeRO con il diagramma di questo post del blog

(Fonte: link)

a. Fase 1: Shards degli stati dell’ottimizzatore tra i lavoratori/GPU paralleli dei dati

b. Fase 2: Shards degli stati dell’ottimizzatore + gradienti tra i lavoratori/GPU paralleli dei dati

c. Fase 3: Shards degli stati dell’ottimizzatore + gradienti + parametri del modello tra i lavoratori/GPU paralleli dei dati

d. Scarico dell’ottimizzatore: Scarica i gradienti + stati dell’ottimizzatore su CPU/Disk basandosi su ZERO Fase 2

e. Scarico dei parametri: Scarica i parametri del modello su CPU/Disk basandosi su ZERO Fase 3

In questo post del blog esamineremo come sfruttare il parallelismo dei dati utilizzando ZeRO utilizzando Accelerate. DeepSpeed, FairScale e PyTorch FullyShardedDataParallel (FSDP) hanno implementato le idee principali del paper ZERO. Questi sono già stati integrati nel Trainer di 🤗 transformers e nell’acceleratore di 🤗 accelerate accompagnati da ottimi blog Fit More and Train Faster With ZeRO via DeepSpeed and FairScale [4] e Accelerate Large Model Training using PyTorch Fully Sharded Data Parallel [5]. Rimandiamo l’illustrazione di ciò che avviene dietro le quinte a quei blog e ci concentriamo principalmente sull’utilizzo di DeepSpeed ZeRO utilizzando Accelerate.

Configurazione dell’hardware: 2X24GB NVIDIA Titan RTX GPU. 60GB di RAM.

Esamineremo il compito di affinamento del modello solo encoder per la classificazione del testo. Utilizzeremo il modello preaddestrato microsoft/deberta-v2-xlarge-mnli (900M parametri) per l’affinamento sul dataset MRPC GLUE.

Il codice è disponibile qui run_cls_no_trainer.py. È simile all’esempio ufficiale di classificazione del testo qui con l’aggiunta di una logica per misurare il tempo di addestramento e valutazione. Confrontiamo le prestazioni tra Distributed Data Parallel (DDP) e DeepSpeed ZeRO Stage-2 in un ambiente Multi-GPU.

Per abilitare DeepSpeed ZeRO Stage-2 senza apportare alcuna modifica al codice, eseguire accelerate config e sfruttare il plugin Accelerate DeepSpeed.

Esempio di plugin DeepSpeed ZeRO Stage-2

compute_environment: LOCAL_MACHINE
deepspeed_config:
 gradient_accumulation_steps: 1
 gradient_clipping: 1.0
 offload_optimizer_device: none
 offload_param_device: none
 zero3_init_flag: false
 zero_stage: 2
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false

Ora, eseguire il comando seguente per l’addestramento:

accelerate launch run_cls_no_trainer.py \
  --model_name_or_path "microsoft/deberta-v2-xlarge-mnli" \
  --task_name "mrpc" \
  --ignore_mismatched_sizes \
  --max_length 128 \
  --per_device_train_batch_size 40 \
  --learning_rate 2e-5 \
  --num_train_epochs 3 \
  --output_dir "/tmp/mrpc/deepspeed_stage2/" \
  --with_tracking \
  --report_to "wandb" \

Nella nostra configurazione Single-Node Multi-GPU, la dimensione massima del batch che DDP supporta senza errori di memoria esaurita (OOM) è 8. Al contrario, DeepSpeed Zero-Stage 2 consente una dimensione del batch di 40 senza incorrere in errori di memoria esaurita (OOM). Pertanto, DeepSpeed consente di adattare 5X più dati per GPU rispetto a DDP. Di seguito è riportato lo snapshot dei grafici da wandb run insieme a una tabella di confronto delle prestazioni tra DDP e DeepSpeed.


Tabella 1: Benchmarking di DeepSpeed ZeRO Stage-2 sul modello DeBERTa-XL (900M)


Con questa dimensione di batch più grande, osserviamo un incremento di velocità di ~ 3,5 volte nel tempo totale di addestramento senza alcuna diminuzione delle metriche di prestazione, tutto ciò senza modificare alcun codice. Evviva! 🤗.

Per essere in grado di modificare più opzioni, sarà necessario utilizzare un file di configurazione DeepSpeed e apportare modifiche minime al codice. Vediamo come fare questo.

Prima di tutto, esamineremo il compito di addestrare un modello di sequenza su sequenza per addestrare il nostro Chatbot. In particolare, addestreremo facebook/blenderbot-400M-distill sul dataset smangrul/MuDoConv (Multi-Domain Conversation). Il dataset contiene conversazioni provenienti da 10 diverse fonti di dati che coprono persone, contesti emotivi specifici, obiettivi (ad esempio, prenotazione di un ristorante) e argomenti generali di Wikipedia (ad esempio, cricket).

Il codice è disponibile qui run_seq2seq_no_trainer.py . La pratica attuale per misurare efficacemente l’coinvolgimento e l’umanità dei Chatbot è attraverso valutazioni umane che sono costose [6]. Pertanto, per questo esempio, la metrica seguita è il punteggio BLEU (che non è ideale ma è la metrica convenzionale per compiti simili). È possibile adattare il codice per addestrare modelli T5 più grandi se si dispone di GPU che supportano la precisione bfloat16, altrimenti si otterranno valori di perdita NaN. Eseguiamo un benchmark rapido su 10000 campioni di addestramento e 1000 campioni di valutazione poiché siamo interessati a DeepSpeed vs DDP.

Sfrutteremo la configurazione DeepSpeed Zero Stage-2 zero2_config_accelerate.json (riportata di seguito) per l’addestramento. Per informazioni dettagliate sulle varie funzionalità di configurazione, fare riferimento alla documentazione di DeeSpeed.

{
    "fp16": {
        "enabled": "true",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 15,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "weight_decay": "auto",
            "torch_adam": true,
            "adam_w_mode": true
        }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto",
            "total_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 2,
        "allgather_partitions": true,
        "allgather_bucket_size": 2e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    },
    "gradient_accumulation_steps": 1,
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

Per abilitare DeepSpeed ZeRO Stage-2 con la configurazione precedente, eseguire accelerate config e fornire il percorso del file di configurazione quando richiesto. Per ulteriori dettagli, fare riferimento alla documentazione ufficiale di 🤗 accelerate per il file di configurazione DeepSpeed.

Esempio di File di Configurazione DeepSpeed ZeRO Stage-2

compute_environment: LOCAL_MACHINE
deepspeed_config:
 deepspeed_config_file: /path/to/zero2_config_accelerate.json
 zero3_init_flag: false
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false

Ora, eseguire il comando seguente per l’addestramento:

accelerate launch run_seq2seq_no_trainer.py \
    --dataset_name "smangrul/MuDoConv" \
    --max_source_length 128 \
    --source_prefix "chatbot: " \
    --max_target_length 64 \
    --val_max_target_length 64 \
    --val_min_target_length 20 \
    --n_val_batch_generations 5 \
    --n_train 10000 \
    --n_val 1000 \
    --pad_to_max_length \
    --num_beams 10 \
    --model_name_or_path "facebook/blenderbot-400M-distill" \
    --per_device_train_batch_size 200 \
    --per_device_eval_batch_size 100 \
    --learning_rate 1e-6 \
    --weight_decay 0.0 \
    --num_train_epochs 1 \
    --gradient_accumulation_steps 1 \
    --num_warmup_steps 100 \
    --output_dir "/tmp/deepspeed_zero_stage2_accelerate_test" \
    --seed 25 \
    --logging_steps 100 \
    --with_tracking \
    --report_to "wandb" \
    --report_name "blenderbot_400M_finetuning"

Quando si utilizza la configurazione DeepSpeed, se l’utente ha specificato optimizer e scheduler nella configurazione, l’utente dovrà utilizzare accelerate.utils.DummyOptim e accelerate.utils.DummyScheduler. Questi sono gli unici piccoli cambiamenti che l’utente deve fare. Di seguito mostriamo un esempio delle modifiche minime richieste quando si utilizza la configurazione DeepSpeed:

- optimizer = torch.optim.Adam(optimizer_grouped_parameters, lr=args.learning_rate)
+ optimizer = accelerate.utils.DummyOptim(optimizer_grouped_parameters, lr=args.learning_rate)

- lr_scheduler = get_scheduler(
-     name=args.lr_scheduler_type,
-     optimizer=optimizer,
-     num_warmup_steps=args.num_warmup_steps,
-     num_training_steps=args.max_train_steps,
- )

+ lr_scheduler = accelerate.utils.DummyScheduler(
+     optimizer, total_num_steps=args.max_train_steps, warmup_num_steps=args.num_warmup_steps
+ )

Tabella 2: Benchmarking di DeepSpeed ZeRO Stage-2 sul modello BlenderBot (400M)

Nella nostra configurazione Single-Node Multi-GPU, la dimensione massima del batch che DDP supporta senza errori OOM è 100. Al contrario, DeepSpeed Zero-Stage 2 consente una dimensione del batch di 200 senza incorrere in errori OOM. Pertanto, DeepSpeed consente di adattare 2X più dati per GPU rispetto a DDP. Osserviamo un miglioramento di ~ 1,44X nella fase di addestramento e ~ 1,23X nella fase di valutazione poiché siamo in grado di adattare più dati sull’hardware disponibile. Poiché questo modello è di dimensioni VoAGI, l’aumento di velocità non è così entusiasmante, ma migliorerà con modelli più grandi. Puoi chattare con il Chatbot addestrato utilizzando tutti i dati su 🤗 Space smangrul/Chat-E . Puoi fornire al bot una personalità, una conversazione di base con una particolare emozione, usarlo in compiti orientati a obiettivi o in modo libero. Di seguito è riportata una conversazione divertente con il chatbot 💬. Puoi trovare screenshot di altre conversazioni utilizzando contesti diversi qui .


Offloading CPU/Disk per consentire l’addestramento di modelli enormi che non si adattano alla memoria GPU

Su una singola GPU NVIDIA Titan RTX da 24 GB, non è possibile addestrare il modello GPT-XL (1,5 miliardi di parametri) neanche con una dimensione del batch di 1. Vedremo come possiamo utilizzare DeepSpeed ZeRO Stage-3 con il trasferimento dei dati dell’ottimizzatore, dei gradienti e dei parametri verso la CPU per addestrare il modello GPT-XL.

Sfrutteremo la configurazione di scarico CPU della fase 3 di DeepSpeed Zero zero3_offload_config_accelerate.json (riportata di seguito) per l’addestramento. Il resto del processo di utilizzo della configurazione con 🤗 accelerate è simile all’esperimento precedente.

{
    "fp16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "weight_decay": "auto"
        }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto",
            "total_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "sub_group_size": 1e9,
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": true
    },
    "gradient_accumulation_steps": 1,
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

File di configurazione di esempio ZeRO Stage-3 CPU Offload DeepSpeed

compute_environment: LOCAL_MACHINE
deepspeed_config:
 deepspeed_config_file: /percorso/verso/zero3_offload_config_accelerate.json
 zero3_init_flag: true
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false

Ora, esegui il seguente comando per l’addestramento:

accelerate launch run_clm_no_trainer.py \
--config_name "gpt2-xl" \
--tokenizer_name "gpt2-xl" \
--dataset_name "wikitext" \
--dataset_config_name "wikitext-2-raw-v1" \
--block_size 128 \
--output_dir "/tmp/clm_deepspeed_stage3_offload__accelerate" \
--learning_rate 5e-4 \
--per_device_train_batch_size 16 \
--per_device_eval_batch_size 1 \
--num_train_epochs 1 \
--with_tracking \
--report_to "wandb"\

Tabella 3: Benchmarking DeepSpeed ZeRO Stage-3 CPU Offload sul modello GPT-XL (1.5B)


DDP produrrà un errore OOM anche con una dimensione di batch 1. D’altra parte, con DeepSpeed ZeRO Stage-3 CPU offload, possiamo addestrare con una dimensione di batch di 16.

Infine, ricorda che, 🤗 Accelerate integra solo DeepSpeed, quindi se hai problemi o domande riguardo all’uso di DeepSpeed, ti preghiamo di segnalare un problema su DeepSpeed GitHub.

[1] Addestra grandi modelli, quindi comprimi: ripensa alle dimensioni del modello per un addestramento e un’inferenza efficienti dei Transformers

[2] ZeRO: ottimizzazioni della memoria per modelli di addestramento con trilioni di parametri

[3] DeepSpeed: addestramento di modelli su scala estrema per tutti – Microsoft Research

[4] Adatta di più e addestra più velocemente con ZeRO tramite DeepSpeed e FairScale

[5] Accelerazione dell’addestramento di modelli di grandi dimensioni utilizzando il parallelismo dati completamente suddiviso di PyTorch

[6] Ricette per la creazione di un chatbot a dominio aperto