# 完全分片数据并行 [完全分片数据并行](https://pytorch.org/docs/stable/fsdp.html)（FSDP）专为训练多达 1T 参数的大型预训练模型而设计。FSDP 通过将模型参数、梯度和优化器状态分片到数据并行进程中来实现这一点，并且还可以将分片后的模型参数卸载到 CPU。FSDP 提供的内存效率使您能够将训练扩展到更大的批次或模型规模。 这两项功能均受 🤗 Accelerate 支持，并且可以与 🤗 PEFT 配合使用。 # 使用 PEFT 和 FSDP 本指南部分将帮助你学习如何使用我们的 DeepSpeed [训练脚本](https://github.com/huggingface/peft/blob/main/examples/sft/train.py) 来执行 SFT。你将配置该脚本，在单台机器上的 8 块 H100 80GB GPU 上，使用 LoRA 和 FSDP 对 Llama-70B 模型进行 SFT（监督微调）。你可以通过修改 accelerate 配置，将其扩展至多台机器。 ## 配置 首先运行以下命令，使用 🤗 Accelerate [创建 FSDP 配置文件](https://huggingface.co/docs/accelerate/quicktour#launching-your-distributed-script)。`--config_file` 标志允许你将配置文件保存到特定位置，否则它会作为 `default_config.yaml` 文件保存在 🤗 Accelerate 缓存中。 配置文件用于在启动训练脚本时设置默认选项。 ```bash accelerate config --config_file fsdp_config.yaml ``` 系统会询问几个关于您环境配置的问题，请配置以下参数。在本示例中，您将按照下图所示回答问卷。 创建 Accelerate 的配置以使用 FSDP 完成此操作后，相应的配置应如下所示，您可以在 [fsdp_config.yaml](https://github.com/huggingface/peft/blob/main/examples/sft/configs/fsdp_config.yaml) 所在的 config 文件夹中找到它： ```yml 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](https://github.com/huggingface/peft/blob/main/examples/sft/run_peft_fsdp.sh)，如下所示： ```bash 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 数据集](https://huggingface.co/datasets/HuggingFaceH4/ultrachat_200k) 的子集上微调 70B Llama 模型。 ## 重要部分 让我们更深入地了解一下这个脚本，看看其中发生了什么，并理解其工作原理。 首先要了解的是，由于传入了 FSDP 配置，该脚本使用 FSDP 进行分布式训练。[SFTTrainer](https://huggingface.co/docs/trl/main/en/sft_trainer#trl.SFTTrainer) 类负责处理使用传入的 peft 配置创建 PEFT 模型的所有繁重工作。之后，当你调用 `trainer.train()` 时，Trainer 内部会使用 🤗 Accelerate 根据 FSDP 配置来准备模型、优化器和训练器，从而创建 FSDP 包装的模型，然后对其进行训练。主要代码片段如下： ```python # 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 内存使用情况 # 使用 PEFT QLoRA 和 FSDP 在多 GPU 上微调大型模型 在本节中，我们将探讨如何使用 QLoRA 和 FSDP 在 2X24GB GPU 上微调 70B 的 Llama 模型。[Answer.AI](https://www.answer.ai/) 与 bitsandbytes 和 Hugging Face 🤗 合作，开源了支持 FSDP+QLoRA 的代码，并在其富有洞察力的博客文章 [You can now train a 70b language model at home](https://www.answer.ai/posts/2024-03-06-fsdp-qlora.html) 中详细解释了整个过程。目前，该功能已集成到 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](https://github.com/huggingface/peft/blob/main/examples/sft/configs/fsdp_config_qlora.yaml) 中找到： ```yml 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](https://github.com/huggingface/peft/blob/main/examples/sft/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`，以实现稳定的自动混合精度训练。 就训练代码而言，重要的代码更改如下： ```diff ... 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。这使得大型模型的微调更加易于获取。 ## 更多资源 您还可以参考 [llama-recipes](https://github.com/facebookresearch/llama-recipes/?tab=readme-ov-file#fine-tuning) 仓库以及 [Getting started with Llama](https://llama.meta.com/get-started/#fine-tuning) 指南，了解如何使用 FSDP 和 PEFT 进行微调。 ## 注意事项 1. 在使用 PEFT 和 FSDP 时，合并目前不受支持，并会引发错误。 2. 目前，将 `modules_to_save` 配置参数传入尚未经过测试。 3. 目前，使用 CPU Offloading 时的 GPU 内存节省效果尚未经过测试。 4. 在使用 FSDP+QLoRA 时，`paged_adamw_8bit` 目前在保存检查点时会导致错误。 5. 使用 FSDP 进行 DoRA 训练应该可以正常工作（尽管速度低于 LoRA）。如果与 bitsandbytes（QDoRA）结合使用，4 位量化也应该可以正常工作，但 8 位量化存在已知问题，不建议使用。 6. 在使用 FSDP2 时，`set_adapter()` 和 `set_requires_grad()` 会对 [PeftModel](/docs/peft/main/zh/package_reference/peft_model#peft.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 目前尚未涵盖。