Fully Sharded Data Parallel

Fully sharded data parallel (FSDP) est conçu pour l'entraînement distribué de grands modèles préentraînés allant jusqu'à 1T de paramètres. FSDP y parvient en fragmentant les paramètres du modèle, les gradients et les états de l'optimiseur entre les processus de parallélisme de données, et il peut également décharger les paramètres du modèle fragmentés vers un CPU. L'efficacité mémoire offerte par FSDP vous permet d'augmenter la taille des lots ou des modèles pour l'entraînement.

Ces deux fonctionnalités sont prises en charge dans 🤗 Accelerate, et vous pouvez les utiliser avec 🤗 PEFT.

Utiliser PEFT et FSDP

Cette section du guide vous aide à apprendre à utiliser notre script d’entraînement DeepSpeed pour effectuer du SFT. Vous configurerez le script pour réaliser le SFT (supervised fine-tuning) du modèle Llama-70B avec LoRA et FSDP sur 8 GPU H100 80 Go sur une seule machine. Vous pouvez le configurer pour passer à l’échelle sur plusieurs machines en modifiant la configuration accelerate.

Configuration

Commencez par exécuter la commande suivante pour créer un fichier de configuration FSDP avec 🤗 Accelerate. L'option --config_file permet d'enregistrer le fichier de configuration à un emplacement spécifique, sinon il est enregistré sous la forme d'un fichier default_config.yaml dans le cache de 🤗 Accelerate.

Le fichier de configuration est utilisé pour définir les options par défaut lors du lancement du script d'entraînement.

accelerate config --config_file fsdp_config.yaml

On vous posera quelques questions sur votre configuration, et vous devrez configurer les arguments suivants. Dans cet exemple, vous répondrez au questionnaire comme indiqué sur l’image ci-dessous.

Creating Accelerate's config to use FSDP

Une fois cela fait, la configuration correspondante doit ressembler à celle ci-dessous et vous pouvez la trouver dans le dossier config à fsdp_config.yaml :

compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch: BACKWARD_PRE
  fsdp_cpu_ram_efficient_loading: true
  fsdp_forward_prefetch: false
  fsdp_offload_params: false
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_state_dict_type: SHARDED_STATE_DICT
  fsdp_sync_module_states: true
  fsdp_use_orig_params: false
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 8
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

Commande de lancement

La commande de lancement est disponible dans run_peft_fsdp.sh et est également affichée ci-dessous :

accelerate launch --config_file "configs/fsdp_config.yaml"  train.py \
--seed 100 \
--model_name_or_path "meta-llama/Llama-2-70b-hf" \
--dataset_name "smangrul/ultrachat-10k-chatml" \
--chat_template_format "chatml" \
--add_special_tokens False \
--append_concat_token False \
--splits "train,test" \
--max_seq_len 2048 \
--num_train_epochs 1 \
--logging_steps 5 \
--log_level "info" \
--logging_strategy "steps" \
--eval_strategy "epoch" \
--save_strategy "epoch" \
--push_to_hub \
--hub_private_repo True \
--hub_strategy "every_save" \
--bf16 True \
--packing True \
--learning_rate 1e-4 \
--lr_scheduler_type "cosine" \
--weight_decay 1e-4 \
--warmup_steps 0 \
--max_grad_norm 1.0 \
--output_dir "llama-sft-lora-fsdp" \
--per_device_train_batch_size 8 \
--per_device_eval_batch_size 8 \
--gradient_accumulation_steps 4 \
--gradient_checkpointing True \
--use_reentrant False \
--dataset_text_field "content" \
--use_flash_attn True \
--use_peft_lora True \
--lora_r 8 \
--lora_alpha 16 \
--lora_dropout 0.1 \
--lora_target_modules "all-linear" \
--use_4bit_quantization False

Notez que nous utilisons LoRA avec rank=8, alpha=16 et que nous visons toutes les couches linéaires. Nous passons le fichier de configuration FSDP et nous affinons le modèle Llama 70B sur un sous-ensemble du jeu de données ultrachat.

Les parties importantes

Plongeons un peu plus en profondeur dans le script pour voir ce qui se passe et comprendre son fonctionnement.

La première chose à savoir est que le script utilise FSDP pour l'entraînement distribué, car la configuration FSDP a été transmise. La classe SFTTrainer gère toute la création du modèle PEFT à l'aide de la configuration peft transmise. Ensuite, lorsque vous appelez trainer.train(), le Trainer utilise en interne 🤗 Accelerate pour préparer le modèle, l'optimiseur et le trainer en utilisant la configuration FSDP afin de créer un modèle enveloppé par FSDP, qui est ensuite entraîné. Le principal extrait de code est ci-dessous :

# trainer
trainer = SFTTrainer(
    model=model,
    processing_class=tokenizer,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    peft_config=peft_config,
)
trainer.accelerator.print(f"{trainer.model}")
if model_args.use_peft_lora:
    # handle PEFT+FSDP case
    trainer.model.print_healthcheck()
    if getattr(trainer.accelerator.state, "fsdp_plugin", None):
        from peft.utils.other import fsdp_auto_wrap_policy

        fsdp_plugin = trainer.accelerator.state.fsdp_plugin
        fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(trainer.model)

# train
checkpoint = None
if training_args.resume_from_checkpoint is not None:
    checkpoint = training_args.resume_from_checkpoint
trainer.train(resume_from_checkpoint=checkpoint)

# saving final model
if trainer.is_fsdp_enabled:
    trainer.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")
trainer.save_model()

Ici, un point principal à retenir actuellement lors de l'utilisation de FSDP avec PEFT est que use_orig_params doit être False pour réaliser des économies de mémoire GPU. En raison de use_orig_params=False, la politique d'auto wrap pour FSDP doit être modifiée afin que les paramètres entraînables et non entraînables soient enveloppés séparément. Cela est réalisé par le fragment de code ci-dessous qui utilise la fonction utilitaire fsdp_auto_wrap_policy de PEFT :

if getattr(trainer.accelerator.state, "fsdp_plugin", None):
    from peft.utils.other import fsdp_auto_wrap_policy

    fsdp_plugin = trainer.accelerator.state.fsdp_plugin
    fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(trainer.model)

Utilisation de la mémoire

Dans l’exemple ci-dessus, la mémoire consommée par GPU est de 72 à 80 Go (90 à 98 %), comme le montre la capture d’écran ci-dessous. La légère augmentation de la mémoire GPU en fin de processus correspond à l’enregistrement du modèle en utilisant le type de state dict FULL_STATE_DICT au lieu de SHARDED_STATE_DICT, afin que le modèle possède des poids d’adaptateurs qui peuvent être chargés normalement avec la méthode from_pretrained lors de l’inférence :

GPU memory usage for the training run

Utiliser PEFT QLoRA et FSDP pour l’ajustement fin de grands modèles sur plusieurs GPU

Dans cette section, nous allons voir comment utiliser QLoRA et FSDP pour l’ajustement fin du modèle llama 70B sur 2 GPU de 24 Go. Answer.AI, en collaboration avec bitsandbytes et Hugging Face 🤗, a publié sous licence open source du code permettant l’utilisation de FSDP+QLoRA et a expliqué l’ensemble du processus dans leur article de blog éclairant You can now train a 70b language model at home. Cette fonctionnalité est désormais intégrée à l’écosystème Hugging Face.

Pour cela, nous avons d'abord besoin de bitsandbytes>=0.43.3, accelerate>=1.0.1, transformers>4.44.2, trl>0.11.4 et peft>0.13.0. Il est nécessaire de définir fsdp_cpu_ram_efficient_loading=true, fsdp_use_orig_params=false et fsdp_offload_params=true (offloading CPU) lors de l'utilisation de la configuration Accelerate. Si vous n'utilisez pas le lanceur accelerate, vous pouvez alternativement définir la variable d'environnement export FSDP_CPU_RAM_EFFICIENT_LOADING=true. Ici, nous utiliserons la configuration accelerate et la configuration suivante peut être trouvée dans fsdp_config_qlora.yaml :

compute_environment: LOCAL_MACHINE                                                                                                                                           
debug: false                                                                                                                                                                 
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch: BACKWARD_PRE
  fsdp_cpu_ram_efficient_loading: true
  fsdp_forward_prefetch: false
  fsdp_offload_params: true
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_state_dict_type: SHARDED_STATE_DICT
  fsdp_sync_module_states: true
  fsdp_use_orig_params: false
machine_rank: 0
main_training_function: main
mixed_precision: 'no'
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

La commande de lancement est donnée ci-dessous et est disponible dans run_peft_qlora_fsdp.sh :

accelerate launch --config_file "configs/fsdp_config_qlora.yaml"  train.py \
--seed 100 \
--model_name_or_path "meta-llama/Llama-2-70b-hf" \
--dataset_name "smangrul/ultrachat-10k-chatml" \
--chat_template_format "chatml" \
--add_special_tokens False \
--append_concat_token False \
--splits "train,test" \
--max_seq_len 2048 \
--num_train_epochs 1 \
--logging_steps 5 \
--log_level "info" \
--logging_strategy "steps" \
--eval_strategy "epoch" \
--save_strategy "epoch" \
--push_to_hub \
--hub_private_repo True \
--hub_strategy "every_save" \
--bf16 True \
--packing True \
--learning_rate 1e-4 \
--lr_scheduler_type "cosine" \
--weight_decay 1e-4 \
--warmup_steps 0 \
--max_grad_norm 1.0 \
--output_dir "llama-sft-qlora-fsdp" \
--per_device_train_batch_size 2 \
--per_device_eval_batch_size 2 \
--gradient_accumulation_steps 2 \
--gradient_checkpointing True \
--use_reentrant True \
--dataset_text_field "content" \
--use_flash_attn True \
--use_peft_lora True \
--lora_r 8 \
--lora_alpha 16 \
--lora_dropout 0.1 \
--lora_target_modules "all-linear" \
--use_4bit_quantization True \
--use_nested_quant True \
--bnb_4bit_compute_dtype "bfloat16" \
--bnb_4bit_quant_storage_dtype "bfloat16"

Notez le nouvel argument transmis, bnb_4bit_quant_storage_dtype, qui désigne le type de données pour l’empaquetage des paramètres 4 bits. Par exemple, lorsqu’il est défini sur bfloat16, 16/4 = 4 paramètres 4 bits sont regroupés ensemble après quantification. Lors de l’entraînement en précision mixte avec bfloat16, bnb_4bit_quant_storage_dtype peut être soit bfloat16 pour un affinement pur en bfloat16, soit float32 pour la précision mixte automatique (ce qui consomme davantage de mémoire GPU). Lors de l’entraînement en précision mixte avec float16, bnb_4bit_quant_storage_dtype doit être défini sur float32 pour un entraînement en précision mixte automatique stable.

En ce qui concerne le code d'entraînement, les modifications de code importantes sont les suivantes :

...

bnb_config = BitsAndBytesConfig(
    load_in_4bit=args.use_4bit_quantization,
    bnb_4bit_quant_type=args.bnb_4bit_quant_type,
    bnb_4bit_compute_dtype=compute_dtype,
    bnb_4bit_use_double_quant=args.use_nested_quant,
+   bnb_4bit_quant_storage=quant_storage_dtype,
)

...

model = AutoModelForCausalLM.from_pretrained(
    args.model_name_or_path,
    quantization_config=bnb_config,
    trust_remote_code=True,
    attn_implementation="flash_attention_2" if args.use_flash_attn else "eager",
+   dtype=quant_storage_dtype or torch.float32,
)

Remarquez que le dtype de AutoModelForCausalLM est identique au type de données bnb_4bit_quant_storage. C’est tout. Le reste est géré par Trainer et TRL.

Utilisation de la mémoire

Dans l’exemple ci-dessus, la mémoire consommée par GPU est de 19,6 Go tandis que l’utilisation de la RAM CPU est d’environ 107 Go. En désactivant le déchargement sur CPU, l’utilisation de la mémoire GPU passe à 35,6 Go / GPU. Ainsi, ce qui nécessitait 16 GPU de 80 Go pour un fine-tuning complet, 8 GPU de 80 Go avec FSDP+LoRA, et quelques GPU de 80 Go avec DDP+QLoRA, n’exige désormais que 2 GPU de 24 Go. Cela rend le fine-tuning de grands modèles plus accessible.

Plus de ressources

Vous pouvez également consulter le dépôt llama-recipes et le guide Getting started with Llama pour savoir comment effectuer un fine-tuning en utilisant FSDP et PEFT.

Limites

La fusion lors de l'utilisation de PEFT et FSDP n'est actuellement pas prise en charge et déclenchera une erreur.
La transmission du paramètre de configuration modules_to_save n'est pas testée à ce jour.
L'économie de mémoire GPU lors de l'utilisation du déchargement sur CPU n'a pas encore été testée.
Lors de l'utilisation de FSDP+QLoRA, paged_adamw_8bit provoque actuellement une erreur lors de l'enregistrement d'un point de contrôle.
L’entraînement DoRA avec FSDP devrait fonctionner (bien qu’à une vitesse inférieure à celle de LoRA). S’il est combiné à bitsandbytes (QDoRA), la quantification 4 bits devrait également fonctionner, mais la quantification 8 bits présente des problèmes connus et n’est pas recommandée.
Avec FSDP2, set_adapter() et set_requires_grad() ré-ensharent les modules FSDP2 à l’intérieur du PeftModel, y compris le PeftModel lui-même, afin que les modifications de requires_grad atteignent les paramètres sharded. La sortie de disable_adapter() ne le fait que si FSDP2 a remplacé les objets de paramètres enregistrés à l’intérieur, comme le fait un forward avec les paramètres sharded, ou si la réactivation modifie requires_grad par rapport à sa valeur à l’entrée. Les modules FSDP2 qui sharent les paramètres d’une couche d’adaptateur fusionné ne sont pas ré-ensharents ; il faut donc dé-fusionner avant d’appeler ces méthodes. Avec reshard_after_forward=False, ou pour le module racine avec la valeur par défaut None, le forward suivant effectue à nouveau un all-gather. Les modules FSDP2 qui enveloppent le PeftModel ne sont pas ré-ensharents ; appelez reshard() sur eux avant ces méthodes. Étant donné que ces appels peuvent laisser les poids de base sharded, la fusion immédiatement après ne fonctionne pas correctement (voir 1). Les méthodes de prompt learning et l’adaption prompt ne sont pas encore prises en charge.

Update on GitHub