Accélérer la formation de modèles volumineux en utilisant DeepSpeed

'Formation plus rapide de modèles volumineux avec DeepSpeed'

Dans cet article, nous verrons comment nous pouvons tirer parti de la bibliothèque Accelerate pour former de grands modèles, ce qui permet aux utilisateurs de tirer parti des fonctionnalités ZeRO de DeeSpeed.

Marre des erreurs de mémoire insuffisante (OOM) lors de l’entraînement de grands modèles ? Nous avons ce qu’il vous faut. Les grands modèles sont très performants [1], mais difficiles à entraîner avec le matériel disponible. Pour tirer le meilleur parti du matériel disponible pour l’entraînement de grands modèles, on peut tirer parti du parallélisme des données en utilisant ZeRO – Zero Redundancy Optimizer [2].

Voici une brève description du parallélisme des données en utilisant ZeRO avec un diagramme provenant de cet article de blog

(Source : lien)

• Décollage ! Comment se lancer dans votre premier projet de ML 🚀
• Politique de gradient avec PyTorch
• Commencer avec l’analyse de sentiment sur Twitter

a. Étape 1 : Partage des états de l’optimiseur entre les travailleurs/GPU parallèles de données

b. Étape 2 : Partage des états de l’optimiseur + des gradients entre les travailleurs/GPU parallèles de données

c. Étape 3 : Partage des états de l’optimiseur + des gradients + des paramètres du modèle entre les travailleurs/GPU parallèles de données

d. Décharge de l’optimiseur : Décharge des gradients + des états de l’optimiseur vers le CPU/le disque, en s’appuyant sur ZERO Étape 2

e. Décharge des paramètres : Décharge des paramètres du modèle vers le CPU/le disque, en s’appuyant sur ZERO Étape 3

Dans cet article de blog, nous verrons comment tirer parti du parallélisme des données en utilisant ZeRO avec Accelerate. DeepSpeed, FairScale et PyTorch FullyShardedDataParallel (FSDP) ont mis en œuvre les idées centrales de l’article ZERO. Ils ont déjà été intégrés dans 🤗 transformers Trainer et 🤗 accelerate accompagnés de formidables blogs Fit More and Train Faster With ZeRO via DeepSpeed and FairScale [4] et Accelerate Large Model Training using PyTorch Fully Sharded Data Parallel [5]. Nous laissons l’explication de ce qui se passe en coulisses à ces blogs et nous nous concentrons principalement sur l’utilisation de DeepSpeed ZeRO avec Accelerate.

Configuration matérielle : 2X24GB NVIDIA Titan RTX GPUs. 60GB RAM.

Nous examinerons la tâche de l’adaptation d’un modèle d’encodeur uniquement pour la classification de texte. Nous utiliserons le modèle pré-entraîné microsoft/deberta-v2-xlarge-mnli (900M params) pour l’adaptation sur l’ensemble de données MRPC GLUE.

Le code est disponible ici run_cls_no_trainer.py. Il est similaire à l’exemple officiel de classification de texte ici, avec l’ajout d’une logique pour mesurer le temps d’entraînement et d’évaluation. Comparons les performances entre Distributed Data Parallel (DDP) et DeepSpeed ZeRO Stage-2 dans une configuration multi-GPU.

Pour activer DeepSpeed ZeRO Stage-2 sans aucun changement de code, veuillez exécuter accelerate config et tirer parti du plugin Accelerate DeepSpeed.

Exemple de 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

Maintenant, exécutez la commande ci-dessous pour l’entraînement :

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" \

Dans notre configuration mono-nœud multi-GPU, la taille de lot maximale que DDP prend en charge sans erreur OOM est de 8. En revanche, DeepSpeed Zero-Stage 2 permet une taille de lot de 40 sans rencontrer d’erreurs OOM. Par conséquent, DeepSpeed permet de traiter 5X plus de données par GPU par rapport à DDP. Voici un aperçu des graphiques de l’exécution wandb ainsi qu’un tableau de comparaison des performances entre DDP et DeepSpeed.

---

Tableau 1 : Évaluation des performances de DeepSpeed ZeRO Stage-2 sur le modèle DeBERTa-XL (900M)

---

Avec une taille de batch plus grande, nous observons une accélération d’environ 3,5 fois du temps total d’entraînement sans aucune baisse des métriques de performance, le tout sans changer de code. Hourra ! 🤗.

Pour pouvoir ajuster plus d’options, vous devrez utiliser un fichier de configuration DeepSpeed et apporter des modifications de code minimales. Voyons comment faire cela.

Tout d’abord, nous nous intéresserons à la tâche de fine-tuning d’un modèle de séquence-à-séquence pour entraîner notre propre Chatbot. Plus précisément, nous affinerons facebook/blenderbot-400M-distill sur l’ensemble de données smangrul/MuDoConv (Multi-Domain Conversation). L’ensemble de données contient des conversations provenant de 10 sources de données différentes couvrant des personas, des ancrages dans des contextes émotionnels spécifiques, des objectifs (par exemple, réservation de restaurant) et des sujets généraux de Wikipédia (par exemple, Cricket).

Le code est disponible ici run_seq2seq_no_trainer.py. La pratique actuelle pour mesurer efficacement l’« Engagingness » et l’« Humanness » des Chatbots est via des évaluations humaines qui sont coûteuses [6]. Pour cet exemple, la métrique suivie est le score BLEU (qui n’est pas idéal mais est la métrique conventionnelle pour de telles tâches). On peut adapter le code pour entraîner des modèles T5 plus grands si vous avez accès à des GPU prenant en charge la précision bfloat16, sinon vous rencontrerez des valeurs de perte NaN. Nous effectuerons un test rapide sur 10000 exemples d’entraînement et 1000 exemples d’évaluation car nous nous intéressons à DeepSpeed vs DDP.

Nous exploiterons la configuration DeepSpeed Zero Stage-2 zero2_config_accelerate.json (donnée ci-dessous) pour l’entraînement. Pour des informations détaillées sur les différentes fonctionnalités de configuration, veuillez consulter la documentation 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
}

Pour activer DeepSpeed ZeRO Stage-2 avec la configuration ci-dessus, veuillez exécuter accelerate config et fournir le chemin du fichier de configuration lorsque cela vous est demandé. Pour plus de détails, consultez la documentation officielle 🤗 accelerate pour le fichier de configuration DeepSpeed.

Exemple de fichier de configuration 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

Maintenant, exécutez la commande ci-dessous pour l’entraînement :

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"

Lors de l’utilisation de la configuration DeepSpeed, si l’utilisateur a spécifié optimizer et scheduler dans la configuration, l’utilisateur devra utiliser accelerate.utils.DummyOptim et accelerate.utils.DummyScheduler. Ce sont les seuls changements mineurs que l’utilisateur doit effectuer. Ci-dessous, nous montrons un exemple des changements minimaux requis lors de l’utilisation de la configuration 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
+ )

---

Tableau 2 : Évaluation du modèle DeepSpeed ZeRO Stage-2 sur BlenderBot (400M)

Dans notre configuration à un seul nœud et plusieurs GPU, la taille maximale de lot que DDP prend en charge sans erreur de mémoire insuffisante est de 100. En revanche, DeepSpeed Zero-Stage 2 permet une taille de lot de 200 sans rencontrer d’erreurs de mémoire insuffisante. Par conséquent, DeepSpeed permet de traiter 2X plus de données par GPU par rapport à DDP. Nous observons une accélération d’environ 1,44X dans l’entraînement et une accélération d’environ 1,23X dans l’évaluation car nous sommes en mesure de traiter plus de données sur le même matériel disponible. Comme ce modèle est de taille VoAGI, l’accélération n’est pas aussi excitante, mais cela s’améliorera avec des modèles plus grands. Vous pouvez discuter avec le Chatbot entraîné en utilisant l’ensemble de données complet sur 🤗 Space smangrul/Chat-E. Vous pouvez donner une personnalité au bot, donner une conversation de référence à une émotion particulière, l’utiliser dans des tâches orientées vers un objectif ou de manière libre. Ci-dessous se trouve une conversation amusante avec le chatbot 💬. Vous pouvez trouver des captures d’écran de conversations supplémentaires utilisant différents contextes ici.

---

Déchargement CPU/disque pour permettre l’entraînement de modèles gigantesques qui ne tiennent pas dans la mémoire GPU

Avec une seule GPU NVIDIA Titan RTX de 24 Go, il est impossible d’entraîner le modèle GPT-XL (1,5 milliard de paramètres) même avec une taille de lot de 1. Nous allons voir comment nous pouvons utiliser DeepSpeed ZeRO Stage-3 avec le déchargement CPU des états de l’optimiseur, des gradients et des paramètres pour entraîner le modèle GPT-XL.

Nous utiliserons la configuration de déchargement CPU de la scène zéro3_offload_config_accelerate.json de DeepSpeed Zero Stage-3 (indiquée ci-dessous) pour l’entraînement. Le reste du processus d’utilisation de la configuration avec 🤗 accelerate est similaire à l’expérience ci-dessus.

{
    "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
}

Exemple de fichier de configuration DeepSpeed ZeRO Stage-3 CPU Offload

compute_environment: LOCAL_MACHINE
deepspeed_config:
 deepspeed_config_file: /chemin/vers/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

Maintenant, exécutez la commande suivante pour l’entraînement :

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"\

---

Tableau 3 : Évaluation des performances de DeepSpeed ZeRO Stage-3 CPU Offload sur le modèle GPT-XL (1,5 milliard)

---

DDP entraînera une erreur OOM même avec une taille de lot de 1. En revanche, avec DeepSpeed ZeRO Stage-3 CPU offload, nous pouvons entraîner avec une taille de lot de 16.

Enfin, veuillez noter que, 🤗 Accelerate intègre uniquement DeepSpeed. Par conséquent, si vous rencontrez des problèmes ou avez des questions concernant l’utilisation de DeepSpeed, veuillez soumettre un problème sur le référentiel GitHub de DeepSpeed.

[1] Former en utilisant de grands modèles puis les compresser : repenser la taille des modèles pour un entraînement et une inférence efficaces des Transformers

[2] ZeRO : Optimisations de la mémoire pour l’entraînement de modèles à un trillion de paramètres

[3] DeepSpeed : Entraînement de modèles à grande échelle pour tous – Microsoft Research

[4] Plus d’ajustement et d’entraînement plus rapide avec ZeRO via DeepSpeed et FairScale

[5] Accélérer l’entraînement de modèles volumineux en utilisant PyTorch Fully Sharded Data Parallel

[6] Recettes pour construire un chatbot à domaine ouvert

We will continue to update IPGirl; if you have any questions or suggestions, please contact us!