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

最新推荐文章于 2024-09-26 15:06:22 发布

asd8705

最新推荐文章于 2024-09-26 15:06:22 发布

阅读量237

点赞数

文章标签： llama 人工智能

原文链接：https://www.philschmid.de/fsdp-qlora-llama3

版权

像Meta的Llama 3、Mistral AI的Mistral与Mixtral模型或AI21的Jamba这样的开源大模型现在成为了OpenAI的竞争者。然而，大多数情况下，你需要在自己的数据上对模型进行微调，以释放模型的全部潜力。对于较小的模型，如Mistral，通过使用Q-Lora，现在在单个GPU上进行微调变得非常方便。但直到现在，高效地对较大的模型如Llama 3 70b或Mixtral进行微调一直是一个挑战。

这篇博客文章将引导您了解如何使用PyTorch FSDP、Q-Lora以及Hugging Face的TRL、Transformers、peft和datasets对Llama 3进行微调。除了使用FSDP之外，我们还将通过Pytorch SDPA实现使用Flash Attention v2。

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

注意：此博客是在NVIDIA H100和NVIDIA A10G GPU上创建和验证的。配置和代码针对每块GPU（共4块）优化，每块GPU配备24GB内存。希望此示例尽可能对大多数人来说都是可访问的。如果您有更多计算资源，可以在步骤3的配置（yaml）中进行更改。

FSDP + Q-Lora 背景

在Answer.AI、Tim Dettmers（Q-Lora的创建者）以及Hugging Face的共同合作下，我们很自豪地宣布支持Q-Lora和PyTorch FSDP（全分布式数据并行）。FSDP和Q-Lora现在允许您在2块消费级GPU（每块24GB内存）上微调Llama 2 70b或Mixtral 8x7B。如果您想了解这次合作的背景，可以阅读文章《现在您可以在家中训练70b语言模型》。Hugging Face的PEFT是实现这种效果的地方，有关更多信息，请参阅PEFT文档。

PyTorch FSDP是一种数据/模型并行技术，它将模型分片到GPU上，减少了内存需求，并使更大模型的训练更加高效。
Q-LoRA是一种微调方法，它利用量化和低秩适配器来有效地减少计算需求和内存占用。

1、设置开发环境

我们的第一步是安装Hugging Face库和Pyroch，包括trl、transformers和datasets。如果你之前没听说过trl，不用担心。它是在transformers和datasets之上的一个新库，使得更容易进行微调、rlhf（奖励强化学习和反馈）以及对开放的LLM（大型语言模型）进行对齐。

# Install Pytorch for FSDP and FA/SDPA
%pip install "torch==2.2.2" tensorboard
 
# Install Hugging Face libraries
%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、创建并准备数据集

在我们的环境配置完成后，我们可以开始创建并准备我们的数据集。用于微调的数据集应该包含解决你想要解决的任务的各种示例。如果你想了解更多关于如何创建数据集的信息，可以查看2024年使用Hugging Face微调LLM的指南。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
 
# Convert dataset to OAI messages
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
 
# Load dataset from the hub
dataset = load_dataset("HuggingFaceH4/no_robots")
 
# Add system message to each conversation
columns_to_remove = list(dataset["train"].features)
columns_to_remove.remove("messages")
dataset = dataset.map(create_conversation, remove_columns=columns_to_remove,batched=False)
 
# Filter out conversations which are corrupted with wrong turns, keep which have even number of turns after adding system message
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)
 
# save datasets to disk
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微调大模型

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

我们已经准备了一个名为run_fsdp_qlora.py的脚本，它将从磁盘加载数据集，准备模型和分词器，并开始训练。此脚本使用trl中的SFTTrainer来微调我们的模型。SFTTrainer使得监督微调开源大型语言模型变得简单，支持：

数据集格式化，包括对话和指令格式（✅ 使用）
仅在完成时进行训练，忽略提示（❌ 不使用）
为更高效的训练打包数据集（✅ 使用）
PEFT（参数高效微调）支持，包括Q-LoRA（✅ 使用）
为对话微调准备模型和分词器（❌ 不使用，参见下方）
注意：我们使用了与Anthropic/Vicuna类似的聊天模板，其中包含User:和Assistant:的角色。这是因为在基础Llama 3中，特殊标记（如<|begin_of_text|>或<|reserved_special_token_XX|>）没有经过训练。这意味着如果想要使用这些标记作为模板的一部分，我们需要训练它们，这需要更多的内存，因为我们需要更新嵌入层和lm_head。如果你拥有更多的计算资源，可以修改run_fsdp_qlora.py脚本中的LLAMA_3_CHAT_TEMPLATE。

配置方面，我们使用了新的TrlParser，它允许我们通过yaml文件提供超参数，并可以通过在命令行界面（CLI）显式传递参数来覆盖配置文件中的参数，例如--num_epochs 10。以下是针对在4个A10G GPU或4个24GB GPU上微调Llama 3 70B的配置文件示例。

%%writefile llama_3_70b_fsdp_qlora.yaml
# script parameters
model_id: "meta-llama/Meta-Llama-3-70b" # Hugging Face model id
dataset_path: "."                      # path to dataset
max_seq_len:  3072 # 2048              # max sequence length for model and packing of the dataset
# training parameters
output_dir: "./llama-3-70b-hf-no-robot" # Temporary output directory for model checkpoints
report_to: "tensorboard"               # report metrics to tensorboard
learning_rate: 0.0002                  # learning rate 2e-4
lr_scheduler_type: "constant"          # learning rate scheduler
num_train_epochs: 3                    # number of training epochs
per_device_train_batch_size: 1         # batch size per device during training
per_device_eval_batch_size: 1          # batch size for evaluation
gradient_accumulation_steps: 2         # number of steps before performing a backward/update pass
optim: adamw_torch                     # use torch adamw optimizer
logging_steps: 10                      # log every 10 steps
save_strategy: epoch                   # save checkpoint every epoch
evaluation_strategy: epoch             # evaluate every epoch
max_grad_norm: 0.3                     # max gradient norm
warmup_ratio: 0.03                     # warmup ratio
bf16: true                             # use bfloat16 precision
tf32: true                             # use tf32 precision
gradient_checkpointing: true           # use gradient checkpointing to save memory
# FSDP parameters: https://huggingface.co/docs/transformers/main/en/fsdp
fsdp: "full_shard auto_wrap offload" # remove offload if enough GPU memory
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卸载，需要更改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

预期内存使用情况：

全精细调用（Full-finetuning）使用FSDP需要大约16个80GB的GPU。
FSDP + LoRA需要大约8个80GB的GPU。
FSDP + Q-Lora需要大约2个40GB的GPU。
FSDP + Q-Lora + CPU卸载需要4个24GB的GPU，序列长度为3072，批大小为1时，每GPU需要22GB内存和总计127GB的CPU RAM。

使用Flash Attention对Llama 3 70B进行3个周期的训练，使用10k样本的数据集，在g5.12xlarge实例上需要45小时。该实例的成本为每小时5.67美元，总成本为255.15美元。这听起来可能很昂贵，但允许你在小型GPU资源上微调Llama 3 70B。如果我们将训练规模扩大到4个H100 GPU，训练时间将缩短至约1.25小时。如果我们假设1个H100的成本为每小时5-10美元，则总成本将在25美元至50美元之间。

我们可以看到可访问性和性能之间的权衡。如果你能够访问更多的/更好的计算资源，你可以减少训练时间和成本，但即使使用较小的资源，你也可以微调Llama 3 70B。成本与性能之间的差异在于，对于4个A10G GPU，我们需要将模型卸载到CPU上，这降低了整体浮点运算能力。

备注：在博客文章的评估和测试过程中，我注意到大约40个最大步数（将80个样本堆叠到3k序列长度）足以获得初步结果。对于40步的训练大约需要1小时，或者大约5美元。

可选：将LoRA适配器合并到原始模型中

使用QLoRA时，我们仅训练适配器而不是整个模型。这意味着在训练过程中保存模型时，我们只保存适配器权重，而不保存整个模型。如果你想保存整个模型，这使得使用文本生成推理更加容易，你可以使用`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模型并非易事，因为一个输入可以有多个正确的输出。如果你想了解更多关于评估生成模型的知识，可以查看“使用Langchain和Hugging Face的Evaluate LLMs和RAG”博客文章，其中提供了实践示例。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"
 
# Load Model with 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
 
 
# Load our test dataset
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]
 
# Test on sample
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.

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

如果你想将模型部署到生产环境中，请查看“为生产部署LLM”指南。

How to Fine-Tune LLMs in 2024 with Hugging Face