完全分片数据并行

完全分片数据并行（FSDP）专为训练多达 1T 参数的大型预训练模型而设计。FSDP 通过将模型参数、梯度和优化器状态分片到数据并行进程中来实现这一点，并且还可以将分片后的模型参数卸载到 CPU。FSDP 提供的内存效率使您能够将训练扩展到更大的批次或模型规模。

这两项功能均受 🤗 Accelerate 支持，并且可以与 🤗 PEFT 配合使用。

使用 PEFT 和 FSDP

本指南部分将帮助你学习如何使用我们的 DeepSpeed 训练脚本来执行 SFT。你将配置该脚本，在单台机器上的 8 块 H100 80GB GPU 上，使用 LoRA 和 FSDP 对 Llama-70B 模型进行 SFT（监督微调）。你可以通过修改 accelerate 配置，将其扩展至多台机器。

配置

首先运行以下命令，使用 🤗 Accelerate 创建 FSDP 配置文件。--config_file 标志允许你将配置文件保存到特定位置，否则它会作为 default_config.yaml 文件保存在 🤗 Accelerate 缓存中。

配置文件用于在启动训练脚本时设置默认选项。

accelerate config --config_file fsdp_config.yaml

系统会询问几个关于您环境配置的问题，请配置以下参数。在本示例中，您将按照下图所示回答问卷。

Creating Accelerate's config to use FSDP

完成此操作后，相应的配置应如下所示，您可以在 fsdp_config.yaml 所在的 config 文件夹中找到它：

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_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

请注意，我们使用 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,
    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()

在此，目前使用 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 memory usage for the training run

使用 PEFT QLoRA 和 FSDP 在多 GPU 上微调大型模型

在本节中，我们将探讨如何使用 QLoRA 和 FSDP 在 2X24GB GPU 上微调 70B 的 Llama 模型。Answer.AI 与 bitsandbytes 和 Hugging Face 🤗 合作，开源了支持 FSDP+QLoRA 的代码，并在其富有洞察力的博客文章 You can now train a 70b language model at home 中详细解释了整个过程。目前，该功能已集成到 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 launcher，也可以替代性地设置环境变量 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_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"

注意新传入的参数 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",
+   dtype=quant_storage_dtype or torch.float32,
)

请注意，AutoModelForCausalLM 的 dtype 与 bnb_4bit_quant_storage 的数据类型相同。就是这样。其余部分均由 Trainer 和 TRL 处理。

内存使用

在上面的示例中，每个 GPU 消耗的内存为 19.6 GB，而 CPU 内存使用量约为 107 GB。当禁用 CPU 卸载时，GPU 内存使用量为 35.6 GB/ GPU。因此，全量微调需要 16 个 80GB GPU，FSDP+LoRA 需要 8 个 80GB GPU，DDP+QLoRA 需要几个 80GB GPU，而现在仅需 2 个 24GB GPU。这使得大型模型的微调更加易于获取。

注意事项

在使用 PEFT 和 FSDP 时，合并目前不受支持，并会引发错误。
目前，将 modules_to_save 配置参数传入尚未经过测试。
目前，使用 CPU Offloading 时的 GPU 内存节省效果尚未经过测试。
在使用 FSDP+QLoRA 时，paged_adamw_8bit 目前在保存检查点时会导致错误。
使用 FSDP 进行 DoRA 训练应该可以正常工作（尽管速度低于 LoRA）。如果与 bitsandbytes（QDoRA）结合使用，4 位量化也应该可以正常工作，但 8 位量化存在已知问题，不建议使用。
在使用 FSDP2 时，set_adapter() 和 set_requires_grad() 会对 PeftModel 内部的 FSDP2 模块进行重新分片（reshard），包括 PeftModel 本身，以确保 requires_grad 的更改能够传递到分片参数。退出 disable_adapter() 时，仅当 FSDP2 替换了其中注册的参数对象（类似于前向传播对分片参数的处理），或者重新启用导致 requires_grad 的值与进入时不同时，才会执行重新分片。对已合并的适配器层参数进行分片的 FSDP2 模块不会被重新分片，因此在调用这些方法之前请先执行 unmerge。如果设置了 reshard_after_forward=False，或者对于使用默认值 None 的根模块，下一次前向传播会再次执行 all-gather。包裹 PeftModel 的 FSDP2 模块不会被重新分片；请在调用这些方法之前对它们调用 reshard()。由于这些调用可能导致基础权重仍处于分片状态，紧随其后执行合并操作无法正常工作（参见 1）。提示学习方法（prompt learning methods）和 adaption prompt 目前尚未涵盖。

Update on GitHub

PEFT

完全分片数据并行

使用 PEFT 和 FSDP

配置

启动命令

重要部分

内存使用

使用 PEFT QLoRA 和 FSDP 在多 GPU 上微调大型模型

内存使用

更多资源

注意事项