Affiner Llama 2 70B en utilisant PyTorch FSDP

Raffiner Llama 2 70B avec PyTorch FSDP

Introduction

Dans cet article de blog, nous examinerons comment affiner Llama 2 70B en utilisant PyTorch FSDP et les meilleures pratiques associées. Nous utiliserons les bibliothèques Hugging Face Transformers, Accelerate et TRL. Nous apprendrons également comment utiliser Accelerate avec SLURM.

La parallélisation des données entièrement réparties (FSDP) est un paradigme dans lequel les états de l’optimiseur, les gradients et les paramètres sont répartis sur les dispositifs. Pendant la passe avant, chaque unité FSDP effectue une opération de rassemblement complet pour obtenir les poids complets, puis les calculs sont effectués et les fragments provenant des autres dispositifs sont abandonnés. Après la passe avant, la perte est calculée, puis la passe arrière est effectuée. Pendant la passe arrière, chaque unité FSDP effectue une opération de rassemblement complet pour obtenir les poids complets, puis les calculs sont effectués pour obtenir les gradients locaux. Ces gradients locaux sont moyennés et répartis sur les dispositifs via une opération de réduction-éparpillement de sorte que chaque dispositif puisse mettre à jour les paramètres de son fragment. Pour plus d’informations sur ce qu’est PyTorch FSDP, veuillez consulter cet article de blog : Accélérer la formation de grands modèles en utilisant PyTorch Fully Sharded Data Parallel.

(Source: lien)

Matériel utilisé

Nombre de nœuds : 2. Le minimum requis est 1. Nombre de GPU par nœud : 8Type de GPU : A100Mémoire du GPU : 80 GoConnexion intra-nœud : NVLinkRAM par nœud : 1 TBNombre de cœurs CPU par nœud : 96Connexion inter-nœud : Elastic Fabric Adapter

Défis de l’affinage de LLaMa 70B

Nous avons rencontré trois principaux défis lors de la tentative d’affinage de LLaMa 70B avec FSDP :

FSDP enveloppe le modèle après le chargement du modèle pré-entraîné. Si chaque processus/rang dans un nœud charge le modèle Llama-70B, cela nécessiterait 70*4*8 Go ~ 2 To de RAM CPU, où 4 est le nombre d’octets par paramètre et 8 est le nombre de GPU sur chaque nœud. Cela conduirait à une saturation de la RAM CPU et à la terminaison des processus.
La sauvegarde de l’ensemble des points de contrôle intermédiaires en utilisant FULL_STATE_DICT avec le déchargement du CPU sur le rang 0 prend beaucoup de temps et entraîne souvent des erreurs de délai d’attente NCCL en raison d’une attente indéfinie pendant la diffusion. Cependant, à la fin de l’entraînement, nous voulons obtenir l’ensemble complet de l’état du modèle au lieu de l’ensemble d’état fragmenté qui est uniquement compatible avec FSDP.
Nous devons améliorer la vitesse et réduire l’utilisation de la VRAM pour former plus rapidement et économiser les coûts de calcul.

Regardons comment résoudre les défis ci-dessus et affiner un modèle de 70B !

Avant de commencer, voici toutes les ressources requises pour reproduire nos résultats :

Base de code : https://github.com/pacman100/DHS-LLM-Workshop/tree/main/chat_assistant/training avec le correctif de singe flash-attn V2
Configuration FSDP : https://github.com/pacman100/DHS-LLM-Workshop/blob/main/chat_assistant/training/configs/fsdp_config.yaml
Script SLURM launch.slurm : https://gist.github.com/pacman100/1cb1f17b2f1b3139a63b764263e70b25
Modèle : meta-llama/Llama-2-70b-chat-hf
Ensemble de données : smangrul/code-chat-assistant-v1 (mélange de LIMA+GUANACO avec un formatage approprié dans un format prêt à être entraîné)

Prérequis

Suivez d’abord ces étapes pour installer Flash Attention V2 : Dao-AILab/flash-attention: Fast and memory-efficient exact attention (github.com). Installez les dernières versions nocturnes de PyTorch avec CUDA ≥11.8. Installez les autres dépendances selon le fichier DHS-LLM-Workshop/code_assistant/training/requirements.txt. Ici, nous installerons 🤗 Accelerate et 🤗 Transformers depuis la branche principale.

Affinage

Réponse au défi 1

Les PRs huggingface/transformers#25107 et huggingface/accelerate#1777 résolvent le premier défi et ne nécessitent aucun changement de code de la part de l’utilisateur. Voici ce qu’ils font :

Créer le modèle sans poids sur tous les rangs (en utilisant le périphérique meta).
Charger le dictionnaire d’état uniquement sur rank==0 et définir les poids du modèle avec ce dictionnaire d’état sur le rang 0.
Pour tous les autres rangs, faire torch.empty(*param.size(), dtype=dtype) pour chaque paramètre sur le périphérique meta.
Ainsi, rank==0 aura chargé le modèle avec le bon dictionnaire d’état tandis que tous les autres rangs auront des poids aléatoires.
Définir sync_module_states=True afin que l’objet FSDP s’occupe de les diffuser à tous les rangs avant le début de l’entraînement.

Ci-dessous se trouve un extrait de sortie pour un modèle de 7B sur 2 GPUs mesurant la mémoire utilisée et les paramètres du modèle à différentes étapes. Nous pouvons observer que lors du chargement du modèle pré-entraîné, le rang 0 et le rang 1 ont une mémoire totale maximale du CPU de 32744 Mo et 1506 Mo, respectivement. Par conséquent, seul le rang 0 charge le modèle pré-entraîné, ce qui permet une utilisation efficace de la RAM du CPU. Les journaux complets peuvent être trouvés ici

accelerator.process_index=0 Mémoire GPU avant le chargement : 0
accelerator.process_index=0 Mémoire GPU consommée à la fin du chargement (fin-début) : 0
accelerator.process_index=0 Mémoire GPU maximale consommée pendant le chargement (max-début) : 0
accelerator.process_index=0 Mémoire GPU totale maximale consommée pendant le chargement (max) : 0
accelerator.process_index=0 Mémoire CPU avant le chargement : 926
accelerator.process_index=0 Mémoire CPU consommée à la fin du chargement (fin-début) : 26415
accelerator.process_index=0 Mémoire CPU maximale consommée pendant le chargement (max-début) : 31818
accelerator.process_index=0 Mémoire CPU totale maximale consommée pendant le chargement (max) : 32744

accelerator.process_index=1 Mémoire GPU avant le chargement : 0
accelerator.process_index=1 Mémoire GPU consommée à la fin du chargement (fin-début) : 0
accelerator.process_index=1 Mémoire GPU maximale consommée pendant le chargement (max-début) : 0
accelerator.process_index=1 Mémoire GPU totale maximale consommée pendant le chargement (max) : 0
accelerator.process_index=1 Mémoire CPU avant le chargement : 933
accelerator.process_index=1 Mémoire CPU consommée à la fin du chargement (fin-début) : 10
accelerator.process_index=1 Mémoire CPU maximale consommée pendant le chargement (max-début) : 573
accelerator.process_index=1 Mémoire CPU totale maximale consommée pendant le chargement (max) : 1506

Réponse au défi 2

Cela est résolu en choisissant le type de dictionnaire d’état SHARDED_STATE_DICT lors de la création de la configuration FSDP. SHARDED_STATE_DICT enregistre chaque fragment par GPU séparément, ce qui permet de sauvegarder rapidement ou de reprendre l’entraînement à partir d’un point de contrôle intermédiaire. Lorsque FULL_STATE_DICT est utilisé, le premier processus (rang 0) rassemble l’ensemble du modèle sur le CPU, puis le sauvegarde dans un format standard.

Créons la configuration accelerate via la commande suivante :

accelerate config --config_file "fsdp_config.yaml"

La configuration résultante est disponible ici : fsdp_config.yaml. Ici, la stratégie de fragmentation est FULL_SHARD. Nous utilisons TRANSFORMER_BASED_WRAP pour la politique d’enveloppement automatique et il utilise _no_split_module pour trouver le nom du bloc Transformer pour l’enveloppement automatique FSDP imbriqué. Nous utilisons SHARDED_STATE_DICT pour sauvegarder les points de contrôle intermédiaires et les états de l’optimiseur dans ce format recommandé par l’équipe PyTorch. Assurez-vous d’activer la diffusion des paramètres du module à partir du rang 0 au début, comme mentionné dans le paragraphe ci-dessus sur la réponse au défi 1. Nous activons l’entraînement en précision mixte bf16.

Pour obtenir le point de contrôle final correspondant à l’ensemble du dictionnaire d’état du modèle, le code suivant est utilisé :

si trainer.is_fsdp_enabled:
    trainer.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")

trainer.save_model(script_args.output_dir) # ou bien, trainer.push_to_hub() si l'ensemble des ckpt est inférieur à 50 Go, car la limite LFS par fichier est de 50 Go

Réponse au défi 3

Flash Attention et l’activation de la vérification des gradients sont nécessaires pour accélérer l’entraînement et réduire l’utilisation de la VRAM afin de permettre un ajustement fin et d’économiser des ressources de calcul. Le code actuel utilise actuellement un “monkey patching” et l’implémentation se trouve dans chat_assistant/training/llama_flash_attn_monkey_patch.py.

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness présente une méthode pour calculer une attention exacte tout en étant plus rapide et plus efficace en termes de mémoire, en exploitant les connaissances de la hiérarchie de mémoire du matériel/GPU sous-jacent – Plus la bande passante/vitesse de la mémoire est élevée, plus sa capacité est réduite car elle devient plus coûteuse.

Si nous suivons le blog Making Deep Learning Go Brrrr From First Principles, nous pouvons constater que le module Attention sur le matériel actuel est limité par la mémoire/la bande passante. La raison en est que l’Attention consiste principalement en des opérations élémentaires comme indiqué ci-dessous à gauche. Nous pouvons observer que les opérations de masquage, de softmax et de dropout prennent la majeure partie du temps, au lieu des multiplications matricielles qui constituent la majeure partie des FLOPs.

(Source : lien)

C’est précisément le problème auquel Flash Attention s’attaque. L’idée est de supprimer les lectures/écritures HBM redondantes. Pour ce faire, tout est conservé dans la SRAM, toutes les étapes intermédiaires sont effectuées, puis le résultat final est écrit dans la HBM, également appelée Fusion de noyaux. L’illustration ci-dessous montre comment cela permet de surmonter le goulot d’étranglement lié à la mémoire.

(Source : lien)

Le carrelage (tiling) est utilisé lors des passes avant et arrière pour découper le calcul des scores/softmax NxN en blocs afin de surmonter la limitation de la taille de la mémoire SRAM. Pour activer le carrelage, un algorithme softmax en ligne est utilisé. La recomposition est utilisée lors de la passe arrière afin d’éviter de stocker la matrice de softmax/score NxN complète lors de la passe avant. Cela réduit considérablement la consommation de mémoire.

Pour une compréhension simplifiée et approfondie de Flash Attention, veuillez consulter les billets de blog ELI5 : FlashAttention et Making Deep Learning Go Brrrr From First Principles, ainsi que l’article original FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.

Mise en œuvre globale

Pour exécuter l’entraînement en utilisant le lanceur Accelerate avec SLURM, consultez ce gist launch.slurm. Ci-dessous, une commande équivalente montrant comment utiliser le lanceur Accelerate pour exécuter l’entraînement. Notez que nous remplaçons les valeurs main_process_ip, main_process_port, machine_rank, num_processes et num_machines du fichier fsdp_config.yaml. Ici, un autre point important à noter est que le stockage est partagé entre tous les nœuds.

accelerate launch \
    --config_file configs/fsdp_config.yaml \
    --main_process_ip $MASTER_ADDR \
    --main_process_port $MASTER_PORT \
    --machine_rank \$MACHINE_RANK \
    --num_processes 16 \
    --num_machines 2 \
    train.py \
    --model_name "meta-llama/Llama-2-70b-chat-hf" \
    --dataset_name "smangrul/code-chat-assistant-v1" \
    --max_seq_len 2048 \
    --max_steps 500 \
    --logging_steps 25 \
    --eval_steps 100 \
    --save_steps 250 \
    --bf16 True \
    --packing True \
    --output_dir "/shared_storage/sourab/experiments/full-finetune-llama-chat-asst" \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 1 \
    --dataset_text_field "content" \
    --use_gradient_checkpointing True \
    --learning_rate 5e-5  \
    --lr_scheduler_type "cosine" \
    --weight_decay 0.01 \
    --warmup_ratio 0.03 \
    --use_flash_attn True

Le fine-tuning a été effectué en ~13.5 heures et ci-dessous se trouve le graphique de la perte d’entraînement. Calculons l’utilisation des FLOPS du modèle (MFU) pour l’exécution de l’entraînement.

Les GPU A100 effectuent ~3,12e14 FLOPS par seconde (en float32 ou bfloat16).
Nombre de jetons entraînés lors des expériences ci-dessus = longueur de la séquence * taille du lot * nombre d’étapes d’entraînement = (2048 * 16 *
1. = 16,384,000 = 1.64e7
Calcul approximatif pour les expériences ci-dessus = 6 * P (num_params) * D (num_tokens) = 6 * 7e10 * 1.64e7 = 6.89e18 FLOPS
FLOPS d’entraînement par seconde = Calcul approximatif / temps d’entraînement = 6.89e17 / (13.5 * 3600) = 6.89e18 / 4.86e4 = 1.42e14 FLOPS par seconde
MFU = FLOPS d’entraînement par seconde / Performance maximale A100 = 1.42e14 / 3.12e14 = 0.4551 = 45.51% de la performance maximale

Voici un exemple de conversation utilisant le modèle ci-dessus :

Invite du système : Vous êtes un assistant serviable, respectueux et honnête. Répondez toujours de manière aussi utile que possible, tout en restant sûr. Vos réponses ne doivent contenir aucun contenu nocif, non éthique, raciste, sexiste, toxique, dangereux ou illégal. Veillez à ce que vos réponses soient socialement impartiales et positives.

Si une question n'a pas de sens ou n'est pas cohérente sur le plan factuel, expliquez pourquoi au lieu de répondre quelque chose d'incorrect. Si vous ne connaissez pas la réponse à une question, veuillez ne pas partager de fausses informations.

- Humain : Qu'est-ce que l'apprentissage profond ? Expliquez comme un pirate.

+ Assistant : Arrr Grumete ! Êtes-vous à la recherche d'une explication appropriée de l'apprentissage profond ?
+ Très bien alors, écoutez-moi bien, matelot ! L'apprentissage profond est un type de machine learning qui utilise des réseaux de neurones artificiels pour analyser des données.
+ Oui, vous avez bien entendu, matelots ! Les réseaux de neurones artificiels sont des modèles du cerveau humain, mais au lieu d'utiliser des neurones, ils utilisent des nœuds interconnectés ou des "neurones artificiels" pour traiter les informations.
+ Ces nœuds travaillent ensemble pour former des couches d'unités de traitement qui peuvent apprendre à partir de grandes quantités de données.
+ Plus ces systèmes reçoivent de données, plus ils deviennent bons pour reconnaître des motifs et prendre des décisions par eux-mêmes.
+ Alors hissez les voiles et mettez le cap sur l'aventure avec l'apprentissage profond, mes matelots ! C'est un trésor de connaissances qui n'attend que d'être découvert.
+ Que les vents vous soient favorables et que les mers vous suivent !

- Humain : Maintenant, expliquez comme un chef.

+ Assistant : Certainement ! Voici une explication de l'apprentissage profond du point de vue culinaire :

+ Pensez à une recette comme une séquence d'étapes utilisées pour transformer des ingrédients bruts en un délicieux plat.
+ De même, dans l'apprentissage profond, il y a plusieurs couches d'"ingrédients" (ou caractéristiques) qui sont combinées et transformées par le biais de différentes opérations pour produire une sortie finale ou une prédiction.
+ Chaque couche s'appuie sur la précédente, affinant et améliorant la représentation des données d'entrée jusqu'à ce qu'elle atteigne le résultat souhaité.
+ Tout comme les différentes épices rehaussent les saveurs des aliments, chaque couche ajoute de la profondeur et de la complexité à la compréhension du modèle de l'entrée, lui permettant de faire des prédictions plus précises.
+ Et tout comme les chefs doivent équilibrer soigneusement la quantité de sel, de sucre et autres assaisonnements dans une recette, le processus d'entraînement dans l'apprentissage profond consiste à ajuster les poids et les biais du réseau pour obtenir le mélange parfait d'exactitude et d'efficacité.
+ Avec assez de pratique et d'expérimentation, n'importe qui peut devenir un chef émérite... Pardon, un maître de l'apprentissage profond ! Bon appétit !

La conversation complète est formatée comme suit :

<|système|> message du système <|endoftext|> <|prompteur|> Q1 <|endoftext|> <|assistant|> A1 <|endoftext|> ...

Conclusion

Nous avons réussi à affiner de manière efficace le modèle Llama de 70 milliards de paramètres en utilisant PyTorch FSDP dans un environnement multi-nœud multi-GPU tout en relevant différents défis. Nous avons vu comment les 🤗 Transformers et les 🤗 Accelerates prennent désormais en charge une initialisation efficace des grands modèles lors de l’utilisation de FSDP pour résoudre le problème de la mémoire RAM du CPU qui se remplit. Cela a été suivi de pratiques recommandées pour la sauvegarde/chargement des points de contrôle intermédiaires et la sauvegarde du modèle final de manière à pouvoir l’utiliser facilement. Pour permettre une formation plus rapide et réduire l’utilisation de la mémoire GPU, nous avons souligné l’importance de l’attention éclair et de la vérification des gradients. Dans l’ensemble, nous pouvons voir comment une simple configuration utilisant 🤗 Accelerate permet d’affiner de tels grands modèles dans un environnement multi-nœud multi-GPU.

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

Was this article helpful?

93 out of 132 found this helpful

Les chercheurs de Microsoft présentent InstructDiffusion un cadre d’IA unificateur et générique pour aligner les tâches de vision par ordinateur avec les instructions humaines.