完全分片数据并行
完全分片数据并行 (FSDP) 旨在用于训练高达 1T 参数的大型预训练模型。FSDP 通过将模型参数、梯度和优化器状态跨数据并行进程分片来实现这一点,它还可以将分片模型参数卸载到 CPU。FSDP 提供的内存效率使您能够将训练扩展到更大的批次或模型大小。
这两个功能都受 🤗 Accelerate 支持,您可以将它们与 🤗 PEFT 结合使用。
使用 PEFT 和 FSDP
本指南的这一部分将帮助您学习如何使用我们的 DeepSpeed 训练脚本 执行 SFT。您将配置脚本以使用 LoRA 和 FSDP 对 Llama-70B 模型进行 SFT(监督微调),并在单台机器上的 8 个 H100 80GB GPU 上进行训练。您可以通过更改加速配置将其配置为扩展到多台机器。
配置
首先运行以下命令以使用 🤗 Accelerate 创建 FSDP 配置文件。--config_file
标志允许您将配置文件保存到特定位置,否则它将作为 🤗 Accelerate 缓存中的 default_config.yaml
文件保存。
配置文件用于在您启动训练脚本时设置默认选项。
accelerate config --config_file fsdp_config.yaml
系统会询问您一些有关设置的问题,并配置以下参数。在本例中,您将按照下图所示回答问卷。
创建 Accelerate 的配置以使用 FSDP完成后,相应的配置应如下所示,您可以在 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
请注意,我们正在使用 rank=8、alpha=16 的 LoRA,并将其应用于所有线性层。我们正在传递 FSDP 配置文件,并在 ultrachat 数据集 的子集上微调 70B Llama 模型。
重要部分
让我们深入了解一下脚本,以便您了解正在发生的事情以及它的工作原理。
首先要知道的是,该脚本使用 FSDP 进行分布式训练,因为已传递了 FSDP 配置。SFTTrainer
类负责使用传递的 peft 配置创建 PEFT 模型的所有繁重工作。之后,当您调用 trainer.train()
时,Trainer 在内部使用 🤗 Accelerate 来准备模型、优化器和训练器,并使用 FSDP 配置创建 FSDP 包装模型,然后进行训练。主要代码片段如下
# trainer
trainer = SFTTrainer(
model=model,
tokenizer=tokenizer,
args=training_args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
peft_config=peft_config,
packing=data_args.packing,
dataset_kwargs={
"append_concat_token": data_args.append_concat_token,
"add_special_tokens": data_args.add_special_tokens,
},
dataset_text_field=data_args.dataset_text_field,
max_seq_length=data_args.max_seq_length,
)
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
方法正常加载的适配器权重。
在多块 GPU 上使用 PEFT QLoRA 和 FSDP 微调大型模型
在本节中,我们将探讨如何使用 QLoRA 和 FSDP 在 2 个 24GB GPU 上微调 70B Llama 模型。Answer.AI 与 bitsandbytes 和 Hugging Face 🤗 合作开源了代码,支持使用 FSDP+QLoRA,并在其有见地的博文中解释了整个过程 您现在可以在家训练 70b 语言模型。现在已集成到 Hugging Face 生态系统中。
为此,我们首先需要 bitsandbytes>=0.43.0
、accelerate>=0.28.0
、transformers>4.38.2
、trl>0.7.11
和 peft>0.9.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 位量化存在已知问题,不建议使用。