高效微调 Llama 3：使用 PyTorch FSDP 和 Q-Lora

开源 LLM，如 Meta Llama 3、Mistral AI Mistral & Mixtral 模型或 AI21 Jamba，现在成为了 OpenAI 的竞争对手。然而，大多数时候，你需要在自己的数据上微调模型，以充分释放模型的潜力。使用 Q-Lora 微调较小的 LLM（如 Mistral）变得非常方便，只需一块 GPU 就能完成。但高效微调更大的模型，如 Llama 3 70b 或 Mixtral，直到现在仍然是一个挑战。

本文将向你展示如何使用 PyTorch FSDP 和 Q-Lora 微调 Llama 3，并借助 Hugging Face 的 TRL、Transformers、peft 和 datasets。在 FSDP 之外，我们还将通过 Pytorch SDPA 实现使用 Flash Attention v2。

步骤概述

设置开发环境
创建并准备数据集
使用 PyTorch FSDP、Q-Lora 和 SDPA 微调 LLM
测试模型并运行推理

注意：本文在 NVIDIA H100 和 NVIDIA A10G GPUs 上创建并验证。配置和代码优化针对每块具有 24GB 内存的 4xA10G GPU。我希望这使得示例对大多数人尽可能可行。如果你有更多的计算资源，可以在步骤 3 中更改配置（yaml）。

FSDP + Q-Lora 背景

通过 Answer.AI、Q-Lora 创作者 Tim Dettmers 和 Hugging Face 的合作，我们很高兴分享对 Q-Lora 和 PyTorch FSDP（全分片数据并行）的支持。FSDP 和 Q-Lora 现在允许你在 2x 消费级 GPU（24GB）上微调 Llama 2 70b 或 Mixtral 8x7B。如果想了解更多关于此合作的背景，可以查看 You can now train a 70b language model at home。

PyTorch FSDP 是一种数据/模型并行技术，它将模型分片到多个 GPU 上，减少内存需求，并更高效地训练更大的模型。

Q-LoRA 是一种微调方法，它利用量化和低秩适配器来有效地减少计算需求和内存占用。

1. 设置开发环境

我们的第一步是安装 Hugging Face 库和 PyTorch，包括 trl、transformers 和 datasets。如果你还没听说过 trl，不要担心。它是一个基于 transformers 和 datasets 的新库，使微调、RLHF、对齐开放 LLM 更加容易。

# 安装用于 FSDP 和 FA/SDPA 的 Pytorch
%pip install "torch==2.2.2" tensorboard

# 安装 Hugging Face 库
%pip install  --upgrade "transformers==4.40.0" "datasets==2.18.0" "accelerate==0.29.3" "evaluate==0.4.1" "bitsandbytes==0.43.1" "huggingface_hub==0.22.2" "trl==0.8.6" "peft==0.10.0"

接下来，我们需要登录 Hugging Face 以访问 Llama 3 70b 模型。如果你还没有账户并接受了条款，可以创建一个账户。

!huggingface-cli login --token ""

2. 创建并准备数据集

环境设置好后，我们可以开始创建和准备数据集。一个微调数据集应包含你想要解决的任务的多样化示例。如果你想了解更多关于如何创建数据集的信息，请查看 How to Fine-Tune LLMs in 2024 with Hugging Face。

我们将使用 HuggingFaceH4/no_robots 数据集，这是由熟练的人类注释者创建的高质量数据集，包含 10,000 个指令和示例。这些数据可用于监督微调（SFT），以使语言模型更好地遵循指令。No Robots 是根据 OpenAI 的 InstructGPT 论文中的指令数据集建模的，主要由单回合指令组成。

{"messages": [{"role": "system", "content": "You are..."}, {"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}
{"messages": [{"role": "system", "content": "You are..."}, {"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}
{"messages": [{"role": "system", "content": "You are..."}, {"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}

no_robots 数据集有 10,000 个示例，其中 9,500 个用于训练，500 个用于测试。有些示例不包含系统消息。我们将使用 datasets 库加载数据集，添加缺失的系统消息，并将它们保存为单独的 JSON 文件。

from datasets import load_dataset

# 将数据集转换为 OAI 消息格式
system_message = """You are Llama, an AI assistant created by Philipp to be helpful and honest. Your knowledge spans a wide range of topics, allowing you to engage in substantive conversations and provide analysis on complex subjects."""

def create_conversation(sample):
    if sample["messages"][0]["role"] == "system":
        return sample
    else:
      sample["messages"] = [{"role": "system", "content": system_message}] + sample["messages"]
      return sample

# 从 hub 加载数据集
dataset = load_dataset("HuggingFaceH4/no_robots")

# 为每个对话添加系统消息
columns_to_remove = list(dataset["train"].features)
columns_to_remove.remove("messages")
dataset = dataset.map(create_conversation, remove_columns=columns_to_remove, batched=False)

# 过滤掉错误轮次的对话，保留添加系统消息后轮次数为偶数的对话
dataset["train"] = dataset["train"].filter(lambda x: len(x["messages"][1:]) % 2 == 0)
dataset["test"] = dataset["test"].filter(lambda x: len(x["messages"][1:]) % 2 == 0)

# 将数据集保存到磁盘
dataset["train"].to_json("train_dataset.json", orient="records", force_ascii=False)
dataset["test"].to_json("test_dataset.json", orient="records", force_ascii=False)

3. 使用 PyTorch FSDP、Q-Lora 和 SDPA 微调 LLM

我们现在准备使用 PyTorch FSDP、Q-Lora 和 SDPA 微调我们的模型。由于我们运行的是分布式设置，需要使用 torchrun 和 Python 脚本启动训练。

我们准备了一个脚本 run_fsdp_qlora.py，它将从磁盘加载数据集，准备模型、tokenizer 并开始训练。它使用 trl 中的 SFTTrainer 来微调我们的模型。SFTTrainer 简化了监督微调开放 LLM 的过程，支持：

数据集格式化，包括对话和指令格式（✅ 使用）
仅在完成任务时进行训练，忽略提示（❌ 未使用）
打包数据集以提高训练效率（✅ 使用）
PEFT（参数高效微调）支持，包括 Q-LoRA（✅ 使用）
准备模型和 tokenizer 进行对话微调（❌ 未使用，见下文）

注意：我们使用了类似于 Anthropic/Vicuna 的聊天模板，其中包含 User 和 Assistant 角色。这是因为基础 Llama 3 中的特殊标记（或）未经过训练。这意味着如果希望在模板中使用它们，我们需要进行训练，这需要更多的内存，因为需要更新嵌入层和 lm_head。如果你有更多的计算资源，可以修改 run_fsdp_qlora.py 脚本中的 LLAMA_3_CHAT_TEMPLATE。

对于配置，我们使用了新的 TrlParser，允许我们在 yaml 文件中提供超参数，或者通过显式传递它们给 CLI 来覆盖配置文件中的参数，例如 --num_epochs 10。以下是用于在 4x A10G GPU 或 4x24GB GPU 上微调 Llama 3 70B 的配置文件。

%%writefile llama_3_70b_fsdp_qlora.yaml
# 脚本参数
model_id: "meta-llama/Meta-Llama-3-70b" # Hugging Face 模型 id
dataset_path: "."                      # 数据集路径
max_seq_len: 3072                      # 模型和数据集打包的最大序列长度

# 训练参数
output_dir: "./llama-3-70b-hf-no-robot" # 模型检查点的临时输出目录
report_to: "tensorboard"               # 将指标报告给 tensorboard
learning_rate: 0.0002                  # 学习率 2e-4
lr_scheduler_type: "constant"          # 学习率调度器
num_train_epochs: 3                    # 训练 epoch 数
per_device_train_batch_size: 1         # 每个设备上的训练批大小
per_device_eval_batch_size: 1          # 每个设备上的评估批大小
gradient_accumulation_steps: 2         # 执行反向/更新传递前的步数
optim: adamw_torch                     # 使用 torch adamw 优化器
logging_steps: 10                      # 每 10 步记录一次日志
save_strategy: epoch                   # 每个 epoch 保存检查点
evaluation_strategy: epoch             # 每个 epoch 进行评估
max_grad_norm: 0.3                     # 最大梯度范数
warmup_ratio: 0.03                     # 热身比例
bf16: true                             # 使用 bfloat16 精度
tf32: true                             # 使用 tf32 精度
gradient_checkpointing: true           # 使用梯度检查点来节省内存

# FSDP 参数：https://huggingface.co/docs/transformers/main/en/fsdp
fsdp: "full_shard auto_wrap offload"   # 如果有足够的 GPU 内存，移除 offload
fsdp_config:
  backward_prefetch: "backward_pre"
  forward_prefetch: "false"
  use_orig_params: "false"

注意：在训练结束时，GPU 内存使用量会略有增加 (~10%)。这是由于正确保存模型所致。确保 GPU 上有足够的内存来保存模型。参考

为了启动训练，我们将使用 torchrun，以保持示例灵活且易于调整，例如 Amazon SageMaker 或 Google Cloud Vertex AI。对于 torchrun 和 FSDP，我们需要设置环境变量 ACCELERATE_USE_FSDP 和 FSDP_CPU_RAM_EFFICIENT_LOADING，告诉 transformers/accelerate 使用 FSDP 并以内存高效方式加载模型。

注意：要取消 CPU offloading，需要更改 fsdp 的值并移除 offload。这仅适用于 >40GB 的 GPU，因为它需要更多内存。

现在，让我们使用以下命令启动训练：

!ACCELERATE_USE_FSDP=1 FSDP_CPU_RAM_EFFICIENT_LOADING=1 torchrun --nproc_per_node=4 ./scripts/run_fsdp_qlora.py --config llama_3_70b_fsdp_qlora.yaml

预期内存使用情况：

使用 FSDP 完全微调需要 ~16X80GB GPU
FSDP + LoRA 需要 ~8X80GB GPU
FSDP + Q-Lora 需要 ~2x40GB GPU
FSDP + Q-Lora + CPU offloading 需要4x24GB GPU，每个 GPU 占用22 GB/GPU 和127 GB CPU RAM，序列长度为3072，批量大小为1。

使用 Flash Attention 微调 Llama 3 70B 在包含10k样本的数据集上进行3个epoch的训练，在g5.12xlarge上需时45小时。实例成本为5.67$/h，总成本约为255.15$。这看起来很贵，但允许你在小型GPU资源上微调Llama 3 70B。如果我们将训练扩展到4个H100 GPU，训练时间将减少到约1.25小时。如果假设每个H100成本为5-10$/h，总成本将在25$-50$之间。

可以看到，在可访问性和性能之间存在权衡。如果你有更多/更好的计算资源，可以减少训练时间和成本，但即使资源有限，你也可以微调Llama 3 70B。由于需要将模型卸载到CPU，4x A10G GPUs的成本/性能不同，从而降低了整体FLOPS。

注意：在评估和测试过程中，我发现~40个最大步数（80个样本堆叠到3k序列长度）就足够看到初步结果。训练40步约需时1小时或约5美元。

可选：将 LoRA Adapter 合并到原始模型中

使用 QLoRA 时，我们仅训练适配器而不是完整模型。这意味着在训练期间保存模型时，我们只保存适配器权重而不是完整模型权重。如果你想保存完整模型，使其更容易与 Text Generation Inference 一起使用，可以使用 merge_and_unload 方法将适配器权重合并到模型权重中，然后使用 save_pretrained 方法保存模型。这将保存一个默认模型，可用于推理。

注意：可能需要 >192GB CPU 内存。

#### COMMENT IN TO MERGE PEFT AND BASE MODEL ####
# from peft import AutoPeftModelForCausalLM
# # Load PEFT model on CPU
# model = AutoPeftModelForCausalLM.from_pretrained(
#     args.output_dir,
#     torch_dtype=torch.float16,
#     low_cpu_mem_usage=True,
# )
# # Merge LoRA and base model and save
# merged_model = model.merge_and_unload()
# merged_model.save_pretrained(args.output_dir, safe_serialization=True, max_shard_size="2GB")

4. 测试模型并运行推理

在训练完成后，我们希望评估和测试我们的模型。我们将加载来自原始数据集的不同样本，并手动评估模型。评估生成式 AI 模型不是一项简单的任务，因为一个输入可以有多个正确输出。如果你想了解更多关于评估生成式模型的信息，请查看 Evaluate LLMs and RAG a practical example using Langchain and Hugging Face 博客文章。

import torch
from peft import AutoPeftModelForCausalLM
from transformers import AutoTokenizer

peft_model_id = "./llama-3-70b-hf-no-robot"

# 加载带有 PEFT adapter 的模型
model = AutoPeftModelForCausalLM.from_pretrained(
    peft_model_id,
    torch_dtype=torch.float16,
    quantization_config={"load_in_4bit": True},
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(peft_model_id)

让我们加载测试数据集并尝试生成一个指令。

from datasets import load_dataset
from random import randint

# 加载我们的测试数据集
eval_dataset = load_dataset("json", data_files="test_dataset.json", split="train")
rand_idx = randint(0, len(eval_dataset))
messages = eval_dataset[rand_idx]["messages"][:2]

# 测试样本
input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
outputs = model.generate(
    input_ids,
    max_new_tokens=512,
    eos_token_id=tokenizer.eos_token_id,
    do_sample=True,
    temperature=0.6,
    top_p=0.9,
)

response = outputs[0][input_ids.shape[-1]:]
print(f"**Query:**\n{eval_dataset[rand_idx]['messages'][1]['content']}\n")
print(f"**Original Answer:**\n{eval_dataset[rand_idx]['messages'][2]['content']}\n")
print(f"**Generated Answer:**\n{tokenizer.decode(response, skip_special_tokens=True)}")

生成示例：

**Query:**
How long was the Revolutionary War?
**Original Answer:**
The American Revolutionary War lasted just over seven years. The war started on April 19, 1775, and ended on September 3, 1783.
**Generated Answer:**
The Revolutionary War, also known as the American Revolution, was an 18th-century war fought between the Kingdom of Great Britain and the Thirteen Colonies. The war lasted from 1775 to 1783.

看起来不错！🚀现在轮到你了！