fsdp_qlora

使用量化LoRA + FSDP训练大语言模型。

阅读我们的公告博客文章。

你应该将此脚本视为alpha/预览版本。如果你不熟悉测试和调试模型，我们建议再等几个月，让社区更充分地测试这种方法。

集成

FSDP+QLoRA已集成到:

• Axolotl: 实验性支持

安装

以下步骤应该可以正常工作(在Cuda 11.7、11.8和12.1上测试通过):

• 克隆 https://github.com/AnswerDotAI/fsdp_qlora
• pip install llama-recipes fastcore "transformers!=4.38.*,!=4.39.*" --extra-index-url https://download.pytorch.org/whl/test/cu118 作为获取大多数依赖项的简单方法(将118替换为你所需的Cuda版本)
• 安装bitsandbytes pip install bitsandbytes>=0.43.0
• 运行 huggingface-cli login (访问Llama 2)
• 可选库:
  • HQQ量化: 按照HQQ安装说明进行操作。我们的训练脚本使用 HQQBackend.ATEN_BACKPROP，所以还要确保构建自定义内核 cd hqq/kernels && python setup_cuda.py install。
  • Weights and Biases日志记录: pip install wandb
• 推荐使用Pytorch >= 2.2以利用原生flash-attention 2内核。

在双24GB GPU上微调Llama-2 70B

安装完成后，运行 cd fsdp_qlora，然后运行以下命令，开始在最大序列长度为512个token的Alpaca数据集上微调Llama-2 70B。

python train.py \
--model_name meta-llama/Llama-2-70b-hf \
--batch_size 2 \
--context_length 512 \
--precision bf16 \
--train_type qlora \
--use_gradient_checkpointing true \
--use_cpu_offload true \
--dataset alpaca \
--reentrant_checkpointing true

这个示例命令目前使用略超过128GB的CPU RAM。如果你只有128GB可用，我们建议创建一个10-20GB的交换文件以适应初始的使用峰值。

训练选项

对于量化，我们支持HQQ和bitsandbytes。我们目前正在进行基准测试，以帮助你决定使用哪一个。如果你使用bitsandbytes，请确保传递 --reentrant_checkpointing True 以避免触发bitsandbytes中导致高内存使用的bug(修复正在进行中)。

--train_type full

全参数微调。

export CUDA_VISIBLE_DEVICES=4,5 # 可选设置设备
python train.py \
--world_size 2 \ # 可选，在单机上会自动设置
--master_port 12356 \ # 可选，默认为12355
--model_name meta-llama/Llama-2-7b-hf \
--gradient_accumulation_steps 4 \
--batch_size 8 \
--context_length 512 \
--precision bf16 \
--train_type full \
--use_gradient_checkpointing true \
--use_cpu_offload false \
--use_activation_cpu_offload false \
--log_to wandb \
--dataset alpaca

--train_type lora

使用HF PEFT库进行LoRA微调。

- --train_type full \
+ --train_type lora \

--train_type custom_lora

使用自定义LoRA模块进行LoRA微调。

- --train_type full \
+ --train_type custom_lora \

--train_type qlora

使用bitsanbytes Linear4bit层(NF4量化)和HF PEFT库进行4位量化LoRA微调。

- --train_type full \
+ --train_type qlora \
+ --reentrant_checkpointing true \

--train_type custom_qlora

使用bitsanbytes Linear4bit层(NF4量化)和自定义LoRA模块进行4位量化LoRA微调。

- --train_type full \
+ --train_type custom_qlora \
+ --reentrant_checkpointing true \

--train_type hqq_lora

使用HQQ库和自定义LoRA模块进行4位量化LoRA微调。

- --train_type full \
+ --train_type hqq_lora \

--train_type bnb_dora

使用bitsanbytes Linear4bit层(NF4量化)和自定义DoRA模块进行4位量化DoRA微调。

- --train_type full \
+ --train_type bnb_dora \

--train_type hqq_dora

使用HQQ库和自定义DoRA模块进行4位量化DoRA微调。

- --train_type full \
+ --train_type hqq_dora \

--train_type bnb_llama_pro

使用bitsanbytes Linear4bit层(NF4量化)进行4位量化Llama-Pro微调。

要创建llama-pro权重，运行以下命令：

python scripts/block_expansion.py \
--model_name meta-llama/Llama-2-7b-hf \
--output_dir /path/to/llama_pro_weights_directory \
--expansion_rate 0.1

- --train_type full \
+ --train_type bnb_llama_pro \
+ --llama_pro_path /path/to/llama_pro_weights_directory \

--train_type hqq_llama_pro

使用HQQ库进行4位量化Llama-Pro微调。

要创建llama-pro权重，运行以下命令：

python scripts/block_expansion.py \
--model_name meta-llama/Llama-2-7b-hf \
--output_dir /path/to/llama_pro_weights_directory \
--expansion_rate 0.1

- --train_type full \
+ --train_type hqq_llama_pro \
+ --llama_pro_path /path/to/llama_pro_weights_directory \

低内存加载

在量化LoRA训练期间，我们使用自定义量化和加载代码，以避免在将整个模型分片到多个GPU之前将其完全加载到GPU内存中。当使用以下任何训练选项 "qlora", "custom_qlora", "hqq_lora" 时，这是我们训练脚本的默认行为。其他训练选项已经在最大程度上针对低内存加载进行了优化。

我们迭代加载权重，在GPU上对其进行量化，并根据其等级将其放回CPU或meta设备，同时每次处理几个层。我们在所有GPU上执行此操作以初始化量化参数(如零点和缩放因子)，同时在FSDP初始化期间使用 sync_module_states=True 来同步所有GPU上的模型参数和缓冲区。

混合精度训练

--precision bf16 (纯bfloat16)

这将在训练前将所有模型参数转换为 torch.bfloat16，并且不会使用FSDP混合精度。因此，分片和未分片的参数将以bf16存储，前向和后向传递将以bf16进行，梯度归约和更新也将以bf16进行。

--precision fp32 (纯float32)

这将在训练前将所有模型参数转换为 torch.float32，并且不会使用FSDP混合精度。因此，分片和未分片的参数将以fp32存储，前向和后向传递将以fp32进行，梯度归约和更新也将以fp32进行。

--precision mp_fp16_autocast (带自动转换的混合float16)

这将在训练前将所有模型参数转换为 torch.float32，并使用FSDP混合精度，配置如下：

mp_policy = MixedPrecision(param_dtype=torch.float32, reduce_dtype=torch.float32, buffer_dtype=torch.float32)

因此，分片和未分片的参数将以fp32存储。它将使用 autocast(torch.float16) 进行前向和后向传递，以及梯度归约和更新。

--precision mp_bf16_autocast (带自动转换的混合bfloat16)

这将在训练前将所有模型参数转换为 torch.float32，并使用FSDP混合精度，配置如下：

mp_policy = MixedPrecision(param_dtype=torch.float32, reduce_dtype=torch.float32, buffer_dtype=torch.float32)

因此，分片和未分片的参数将以fp32存储。它将使用 autocast(torch.bfloat16) 进行前向和后向传递，以及梯度归约和更新。

--precision mp_bf16_buffers_autocast (bfloat16参数和float32缓冲区，带自动转换)

这将在训练前将所有模型参数转换为 torch.bfloat16，但将缓冲区保持在 torch.float32，并使用FSDP混合精度，配置如下：

mp_policy = MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.bfloat16, buffer_dtype=torch.float32)

因此，分片和未分片的参数将以bf16存储。它将使用 autocast(torch.bfloat16) 进行前向和后向传递，以及梯度归约和更新。缓冲区和自动转换中的符合条件的操作将以bf16执行。

这个选项对RoPE层很重要，因为当转换为较低精度时，特别是在较长的上下文长度下，RoPE层会给出不正确的结果。

与现有训练器的比较

hf_train.py使用TRL的SFTTrainer进行对比运行。为了与我们的脚本匹配，修改数据加载代码以训练所有内容（不仅仅是完成部分），然后运行train.py --train_type qlora --dataset guanaco --batch_size 8 --lr_scheduler cosine --log_to wandb --save_model True --output_dir guanaco_7B --gradient_accumulation_steps 2 --lr 2e-4。SFTTrainer版本必须以较低的批量大小运行（4而不是8），所以我们只进行2个梯度累积步骤，而QLoRA+FSDP版本是4个。

转换保存的模型

如果指定--save_model True，适配器层将被保存为状态字典。要转换为常规的Hugging Face格式并上传到hub，请参见：Converting the State Dict.ipynb

如果使用"custom_qlora", "hqq_lora"训练选项，则只会保存可训练的LoRA参数。在推理之前，你需要再次加载和量化基础模型，并单独加载保存的LoRA参数。

你也可以尝试将基础模型权重与训练后的LoRA权重合并后再量化，看看是否与训练期间分开保存参数的效果相似。要在HQQ中使用torch.compile，请参见https://github.com/mobiusml/hqq/issues/18。

限制

虽然QLoRA微调与FSDP兼容，但这个alpha版本和我们的示例脚本还存在一些不完善之处。

首先，当前版本的Transformer AutoModel.from_pretrained不能用于将模型加载到量化权重中，因为它不支持新的quant_storage或quantization标志。加载预训练模型需要编写或使用自定义的模型加载代码。我们在演示脚本中提供了如何加载和量化QLoRA模型以进行微调的示例。

我们正在与Hugging Face积极合作，以在未来的Transformers和PEFT版本中解决这个不兼容问题。

其次，虽然FSDP的混合精度与QLoRA兼容，但实践者需要小心设置MixedPrecision.param_type以匹配Linear4Bit.quant_storage的dtype。否则，FSDP的混合精度可能会将量化权重转换为不同的精度，实际上将它们变成随机权重。我们的示例脚本展示了如何避免这个潜在的陷阱，我们也很乐意协助模型训练库在使用QLoRA训练时正确地向用户公开FSDP的混合精度选项。

示例：Llama 70B 4-A100 40GB训练

# BnB QLoRA
export CUDA_VISIBLE_DEVICES=4,5,6,7
python train.py \
--world_size 4 \
--master_port 12356 \
--model_name meta-llama/Llama-2-70b-hf \
--gradient_accumulation_steps 4 \
--batch_size 2 \
--context_length 512 \
--precision bf16_buffers_autocast \
--train_type custom_qlora \
--use_gradient_checkpointing true \
--reentrant_checkpointing true
--use_cpu_offload false \
--log_to stdout \
--dataset alpaca

# HQQ QLoRA
export CUDA_VISIBLE_DEVICES=4,5,6,7
python train.py \
--world_size 4 \
--master_port 12356 \
--model_name meta-llama/Llama-2-70b-hf \
--gradient_accumulation_steps 4 \
--batch_size 2 \
--context_length 512 \
--precision bf16_buffers_autocast \
--train_type hqq_lora \
--use_gradient_checkpointing true \
--use_cpu_offload false \
--log_to stdout \
--dataset alpaca

**注意：**对于大批量或长上下文训练，HQQ LoRA比使用重入检查点的BnB LoRA略微更节省内存。因此，如果遇到OOM问题，可以尝试使用HQQ LoRA。

SLURM训练

查看fsdp_multi_node.sh以获取使用SLURM进行多节点训练的示例训练脚本。

添加对新模型的支持

首先，从Transformers导入新模型的transformer、attention和MLP层：

from transformers.models.mistral.modeling_mistral import MistralDecoderLayer, MISTRAL_ATTENTION_CLASSES, MistralMLP

然后在get_wrapping_policy函数中，将attention、MLP和transformer层添加到self_attn_policy_fn、mlp_policy_fn和transformer_wrap_policy包装策略方法中：

def get_wrapping_policy(custom_policy:bool=False):

    def self_attn_policy_fn(module):
        return isinstance(module, tuple(*LLAMA_ATTENTION_CLASSES.values(), *MISTRAL_ATTENTION_CLASSES.values()))

    def mlp_policy_fn(module):
        return isinstance(module, (LlamaMLP, MistralMLP))

    transformer_wrap_policy = functools.partial(
        transformer_auto_wrap_policy,
        transformer_layer_cls=(LlamaDecoderLayer, MistralDecoderLayer),
    )

最后，通过将transformer层添加到check_fn中来添加梯度检查点支持：

if args["use_gradient_checkpointing"]:
    check_fn = lambda submodule: isinstance(submodule, (LlamaDecoderLayer, MistralDecoderLayer))