PEFT 文件
完全分片資料並行
並獲得增強的文件體驗
開始使用
完全分片資料並行
完全分片資料並行 (Fully sharded data parallel, FSDP) 是為分散式訓練高達 1T 引數的大型預訓練模型而開發的。FSDP 透過將模型引數、梯度和最佳化器狀態分片到資料並行程序中來實現這一點,它還可以將分片的模型引數解除安裝到 CPU。FSDP 帶來的記憶體效率使您能夠將訓練擴充套件到更大的批次或模型尺寸。
這兩個功能在 🤗 Accelerate 中都得到支援,您可以將它們與 🤗 PEFT 一起使用。
使用 PEFT 和 FSDP
本指南的這一部分將幫助您學習如何使用我們的 DeepSpeed 訓練指令碼來執行 SFT。您將配置該指令碼,以在單臺機器上的 8xH100 80GB GPU 上使用 LoRA 和 FSDP 對 Llama-70B 模型進行 SFT(監督微調)。您可以透過更改 accelerate 配置來將其擴充套件到多臺機器。
配置
首先執行以下命令,使用 🤗 Accelerate 建立一個 FSDP 配置檔案。`--config_file` 標誌允許您將配置檔案儲存到特定位置,否則它將作為 `default_config.yaml` 檔案儲存在 🤗 Accelerate 快取中。
配置檔案用於在啟動訓練指令碼時設定預設選項。
accelerate config --config_file fsdp_config.yaml
您將被問到一些關於您的設定的問題,並配置以下引數。在本示例中,您將按照下圖所示回答問卷。

完成此操作後,相應的配置應如下所示,您可以在 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
啟動命令
啟動命令可在 run_peft_fsdp.sh 找到,如下所示:
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_ratio 0.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
請注意,我們使用的是 LoRA,秩為 8,alpha 為 16,並針對所有線性層。我們傳遞了 FSDP 配置檔案,並在 ultrachat 資料集 的子集上微調 70B Llama 模型。
重要部分
讓我們深入瞭解一下這個指令碼,以便您能看到發生了什麼,並理解其工作原理。
首先要知道的是,該指令碼使用 FSDP 進行分散式訓練,因為已傳遞了 FSDP 配置。SFTTrainer 類處理了使用傳入的 peft 配置建立 PEFT 模型的所有繁重工作。之後,當您呼叫 `trainer.train()` 時,Trainer 內部使用 🤗 Accelerate,根據 FSDP 配置準備模型、最佳化器和訓練器,以建立 FSDP 包裝的模型,然後對其進行訓練。主要程式碼片段如下:
# 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_trainable_parameters()
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()
這裡,目前使用 FSDP 與 PEFT 時需要注意的一點是,`use_orig_params` 需要設定為 `False` 才能實現 GPU 記憶體節省。由於 `use_orig_params=False`,FSDP 的自動包裝策略需要改變,以便可訓練和不可訓練的引數被分開包裝。這由下面的程式碼片段完成,它使用了 PEFT 的工具函式 `fsdp_auto_wrap_policy`。
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)
記憶體使用
在上面的示例中,每個 GPU 消耗的記憶體為 72-80 GB (90-98%),如下面的螢幕截圖所示。在結尾處 GPU 記憶體的輕微增加,是因為使用 `FULL_STATE_DICT` 狀態字典型別而不是 `SHARDED_STATE_DICT` 來儲存模型,這樣模型就擁有可以在推理時使用 `from_pretrained` 方法正常載入的介面卡權重。

使用 PEFT QLoRA 和 FSDP 在多 GPU 上微調大型模型
在本節中,我們將探討如何使用 QLoRA 和 FSDP 在 2 塊 24GB GPU 上微調 70B Llama 模型。Answer.AI 與 bitsandbytes 和 Hugging Face 🤗 合作,開源了能夠使用 FSDP+QLoRA 的程式碼,並在其富有洞察力的部落格文章 你現在可以在家訓練 70b 語言模型 中解釋了整個過程。這現在已整合到 Hugging Face 生態系統中。
為此,我們首先需要 `bitsandbytes>=0.43.3`、`accelerate>=1.0.1`、`transformers>4.44.2`、`trl>0.11.4` 和 `peft>0.13.0`。使用 Accelerate 配置時,我們需要設定 `fsdp_cpu_ram_efficient_loading=true`、`fsdp_use_orig_params=false` 和 `fsdp_offload_params=true`(CPU 解除安裝)。如果不使用 accelerate 啟動器,您可以交替設定環境變數 `export FSDP_CPU_RAM_EFFICIENT_LOADING=true`。這裡,我們將使用 accelerate 配置,下面的配置可以在 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
啟動命令如下所示,可在 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_ratio 0.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"
注意傳遞的新引數 `bnb_4bit_quant_storage_dtype`,它表示用於打包 4 位引數的資料型別。例如,當它被設定為 `bfloat16` 時,16/4 = 4 個 4 位引數在量化後被打包在一起。當使用 `bfloat16` 進行混合精度訓練時,`bnb_4bit_quant_storage_dtype` 可以是 `bfloat16` 用於純 `bfloat16` 微調,也可以是 `float32` 用於自動混合精度(這會消耗更多 GPU 記憶體)。當使用 `float16` 進行混合精度訓練時,`bnb_4bit_quant_storage_dtype` 應設定為 `float32` 以進行穩定的自動混合精度訓練。
在訓練程式碼方面,重要的程式碼更改如下:
...
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",
+ torch_dtype=quant_storage_dtype or torch.float32,
)
請注意,`AutoModelForCausalLM` 的 `torch_dtype` 與 `bnb_4bit_quant_storage` 的資料型別相同。就是這樣。其他所有事情都由 Trainer 和 TRL 處理。
記憶體使用
在上面的示例中,每個 GPU 消耗的記憶體為 19.6 GB,而 CPU RAM 使用量約為 107 GB。停用 CPU 解除安裝時,GPU 記憶體使用量為 35.6 GB/GPU。因此,原本需要 16 塊 80GB GPU 進行完整微調,8 塊 80GB GPU 使用 FSDP+LoRA,以及幾塊 80GB GPU 使用 DDP+QLoRA 的任務,現在只需要 2 塊 24GB GPU。這使得大型模型的微調更加普及。
更多資源
您還可以參考 llama-recipes 倉庫和 Llama 入門指南,瞭解如何使用 FSDP 和 PEFT 進行微調。
注意事項
- 當前不支援在使用 PEFT 和 FSDP 時進行合併,這會引發錯誤。
- 目前尚未測試傳遞 `modules_to_save` 配置引數。
- 目前尚未測試使用 CPU 解除安裝時的 GPU 記憶體節省情況。
- 使用 FSDP+QLoRA 時,`paged_adamw_8bit` 當前在儲存檢查點時會導致錯誤。
- 使用 FSDP 進行 DoRA 訓練應該可以工作(儘管速度比 LoRA 慢)。如果與 bitsandbytes (QDoRA) 結合使用,4 位量化也應該可以工作,但 8 位量化存在已知問題,不推薦使用。