PEFT 文档

完全分片数据并行

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

完全分片数据并行

完全分片数据并行 (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

您将被问到一些关于您的设置的问题,并配置以下参数。在本示例中,您将按照下图所示回答问卷。

创建 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

请注意,我们使用的是 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` 方法正常加载的适配器权重。

训练运行的 GPU 内存使用情况

使用 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 进行微调。

注意事项

  1. 当前不支持在使用 PEFT 和 FSDP 时进行合并,这会引发错误。
  2. 目前尚未测试传递 `modules_to_save` 配置参数。
  3. 目前尚未测试使用 CPU 卸载时的 GPU 内存节省情况。
  4. 使用 FSDP+QLoRA 时,`paged_adamw_8bit` 当前在保存检查点时会导致错误。
  5. 使用 FSDP 进行 DoRA 训练应该可以工作(尽管速度比 LoRA 慢)。如果与 bitsandbytes (QDoRA) 结合使用,4 位量化也应该可以工作,但 8 位量化存在已知问题,不推荐使用。
< > 在 GitHub 上更新

© . This site is unofficial and not affiliated with Hugging Face, Inc.