如何用TRL微调大模型(LLMs)

​ 截止到2024年,大模型已经有了飞速发展。ChatGPT的面世,催生了一系列的大模型,包括Meta的Llama 2、Mistrals Mistral & Mixtral 模型、TII Falcon和Google的Flan-T5等等。这些大模型可以运用于很多的任务。例如,聊天机器人、Q&A问答、摘要总结等等。然而,如果你想要定制一个使用于你的大模型。你可能需要根据数据对模型进行微调,以获得比提示更高质量的结果,或通过训练更高效的小型模型来节约成本。

​ 这篇博客将引导你了解如何在2024年使用huggingface的TRL、Transformers和数据集对开放式LLM进行微调。在本博客中,我们将

  • 定义用例
  • 设置开发环境
  • 创建并准备数据集
  • 使用 trl 和 SFTTrainer 微调 LLM
  • 测试和评估 LLM
  • 为生产部署 LLM

注:本博客是为在消费级 GPU(24GB)(如英伟达 A10G 或 RTX 4090/3090)上运行而创建的,但也可轻松调整为在更大的 GPU 上运行。

1.定义用例

​ 在微调 LLM 时,了解您的使用案例和要解决的任务非常重要。这将有助于您选择合适的模型,或帮助您创建一个数据集来微调模型。如果您还没有定义用例。您可能需要回到绘图板上。我想说的是,并非所有用例都需要微调,建议在微调您自己的模型之前,先评估并试用已经微调过的模型或基于 API 的模型。

例如,我们将使用以下的用例:

我们希望对模型进行微调,使其能够根据自然语言指令生成 SQL 查询,然后将其集成到我们的 BI 工具中。我们的目标是缩短创建 SQL 查询所需的时间,让非技术用户更容易创建 SQL 查询。
文本到 SQL 可以作为微调 LLM 的良好用例,因为这是一项复杂的任务,需要大量有关数据和 SQL 语言的(内部)知识。

2.设置开发环境

​ 我们的第一步是安装 Hugging Face Libraries 和 Pyroch,包括 trl、transformers和数据集。如果你还没听说过 trl,别担心。它是transformers和数据集之上的一个新库,能让微调、rlhf、对齐开放式 LLM 变得更容易。

# Install Pytorch & other libraries
!pip install "torch==2.1.2" tensorboard

# Install Hugging Face libraries
!pip install  --upgrade \
  "transformers==4.36.2" \
  "datasets==2.16.1" \
  "accelerate==0.26.1" \
  "evaluate==0.4.1" \
  "bitsandbytes==0.42.0" \
  # "trl==0.7.10" # \
  # "peft==0.7.1" \

# install peft & trl from github
!pip install git+https://github.com/huggingface/trl@a3c5b7178ac4f65569975efadc97db2f3749c65e --upgrade
!pip install git+https://github.com/huggingface/peft@4a1559582281fc3c9283892caea8ccef1d6f5a4f --upgrade

​ 如果您使用的是安培架构(如英伟达 A10G 或 RTX 4090/3090)或更新的 GPU,则可以使用闪存注意力。Flash Attention 是一种对注意力计算进行重新排序的方法,它利用经典技术(平铺、重新计算)大大加快了计算速度,并将内存使用量从序列长度的二次方降低到线性。简而言之,就是将训练速度提高 3 倍。了解更多信息,请访问 FlashAttention

注意:如果您的机器内存不足 96GB,CPU 内核较多,请减少 MAX_JOBS 的数量。在 g5.2xlarge 上,我们使用了 4 个。

import torch; assert torch.cuda.get_device_capability()[0] >= 8, 'Hardware not supported for Flash Attention'
# install flash-attn
!pip install ninja packaging
!MAX_JOBS=4 pip install flash-attn --no-build-isolation

安装 flash attention需要很多时间(10-45 分钟)。

​ 我们将使用Hugging Face Hub作为模型远程版本服务。这意味着,我们将在训练的时候自动向hub推送模型、日志和信息。为此,你需要在 Hugging Face 上注册账号,然后我们可以使用huggingface_hub包的login模块来登录账号,并且在本地磁盘上保存我们的token(access key)。

from huggingface_hub import login

login(
  token="", # ADD YOUR TOKEN HERE
  add_to_git_credential=True
)

3.创建并准备数据集

​ 当你决定要微调大模型时,我们需要准备一个数据集用来微调大模型。数据集应该是你需要解决的任务的各种演示集。有多种方法创建这样一个数据集,例如:

  • 利用现有的开源的数据集。e.g., Spider

  • 利用LLMs创建综合数据集。e.g., Alpaca

  • 收集人类活动产生的数据集。e.g., Dolly

  • 结合上述几种方法生成数据集。e.g., Orca

​ 每种方法都有自己的优缺点,并取决于预算、时间和质量要求。例如,使用现有的数据集是最简单的,但可能不适合你的特定用例,而使用人类可能是最准确的,但可能费时费力。也可以将几种方法结合起来创建指令数据集,例如: Orca: Progressive Learning from Complex Explanation Traces of GPT-4.

​ 在我们的示例中,我们将使用一个名为 sql-create-context 的现有数据集,其中包含自然语言指令、模式定义和相应 SQL 查询的样本。

​ 在最新发布的 trl 中,我们现在支持流行的指令和对话数据集格式。这意味着我们只需将数据集转换为其中一种受支持的格式,其余的工作都将由 trl 来完成。这些格式包括:

  • 对话形式(conversational format)

    {"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": "..."}]}
  • 指令格式(instruction format)

    {"prompt": "", "completion": ""}
    {"prompt": "", "completion": ""}
    {"prompt": "", "completion": ""}

​ 在我们的示例中,我们将使用🤗 数据集库加载我们的开源数据集,然后将其转换为对话格式,在系统消息中为我们的助手加入模式定义。然后,我们会将数据集保存为 jsonl 文件,然后可以用它来微调我们的模型。我们将随机对数据集进行降采样,使其只有 10,000 个样本。

注意:这一步可能与您的使用情况不同。例如,如果你已经有了一个数据集,可以跳过这一步,直接进入微调步骤。

from datasets import load_dataset

# Convert dataset to OAI messages
system_message = """You are an text to SQL query translator. Users will ask you questions in English and you will generate a SQL query based on the provided SCHEMA.
SCHEMA:
{schema}"""

def create_conversation(sample):
  return {
    "messages": [
      {"role": "system", "content": system_message.format(schema=sample["context"])},
      {"role": "user", "content": sample["question"]},
      {"role": "assistant", "content": sample["answer"]}
    ]
  }  

# Load dataset from the hub
dataset = load_dataset("b-mc2/sql-create-context", split="train")
dataset = dataset.shuffle().select(range(12500))

# Convert dataset to OAI messages
dataset = dataset.map(create_conversation, remove_columns=dataset.features,batched=False)
# split dataset into 10,000 training samples and 2,500 test samples
dataset = dataset.train_test_split(test_size=2500/12500)

print(dataset["train"][345]["messages"])

# save datasets to disk 
dataset["train"].to_json("train_dataset.json", orient="records")
dataset["test"].to_json("test_dataset.json", orient="records")

4.使用trlSFTTrainer微调LLM

​ 现在我们准备对模型进行微调。我们将使用 trl 中的 SFTTrainer 来微调我们的模型。SFTTrainer 可以直接监督微调开放式 LLM。SFTTrainer 是transformers库中 Trainer 的子类,支持所有相同的功能,包括日志记录、评估和检查点,但增加了更多生活质量功能,包括:

  • 数据集格式,包括会话和指令格式
  • 只对完成情况进行训练,忽略提示
  • 打包数据集,提高训练效率
  • 支持 PEFT(参数效率微调),包括 Q-LoRA
  • 为会话微调准备模型和标记符(例如添加特殊标记符)

​ 我们将在示例中使用数据集格式化、打包和 PEFT 功能。作为 PEFT 方法,我们将使用 QLoRA,这是一种在微调过程中减少大型语言模型内存占用的技术,通过量化不会牺牲性能。如果你想查看更多关于QLoRA的信息,或者想知道它是如何工作的,请查看这篇文章 Making LLMs even more accessible with bitsandbytes, 4-bit quantization and QLoRA

现在,我们开始了🚀

首先,我们需要从磁盘加载数据集。

from datasets import load_dataset

# Load jsonl data from disk
dataset = load_dataset("json", data_files="train_dataset.json", split="train")

​ 接下来,我们将加载 LLM。对于我们的用例,我们将使用 CodeLlama 7B。CodeLlama 是为一般代码合成和理解而训练的 Llama 模型。但我们可以通过更改 model_id 变量,轻松地将该模型换成其他模型,例如 Mistral 或 Mixtral 模型、TII Falcon 或任何其他 LLM。我们将使用 bitsandbytes 将模型量化为 4-bit。

注:请注意,模型越大,所需的内存就越多。在我们的示例中,我们将使用 7B 版本,该版本可在 24GB GPU 上进行调整。如果你的 GPU 较小。

​ 正确地准备 LLM 和tokenizer以训练聊天/对话模型至关重要。我们需要为标记器和模型添加新的特殊标记,并教授如何理解对话中的不同角色。在 trl 中,我们有一个名为 setup_chat_format 的便捷方法:

  • 向tokenizer添加特殊标记,例如 <|im_start|> 和 <|im_end|>,以表示对话的开始和结束。
  • 调整模型嵌入层的大小,以容纳新的标记。
  • 设置标记化器的 chat_template,用于将输入数据格式化为类似聊天的格式。默认使用 OpenAI 的 chatml。
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from trl import setup_chat_format

# Hugging Face model id
model_id = "codellama/CodeLlama-7b-hf" # or `mistralai/Mistral-7B-v0.1`

# BitsAndBytesConfig int-4 config
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
)

# Load model and tokenizer
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    attn_implementation="flash_attention_2",
    torch_dtype=torch.bfloat16,
    quantization_config=bnb_config
)
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.padding_side = 'right' # to prevent warnings

# # set chat template to OAI chatML, remove if you start from a fine-tuned model
model, tokenizer = setup_chat_format(model, tokenizer)

​ SFTTrainer 支持与 peft 的原生集成,这使得使用 QLoRA 等高效调整 LLM 变得超级简单。我们只需创建 LoraConfig 并将其提供给训练器即可。我们的 LoraConfig 参数是根据 qlora paper和 sebastian 的 blog post定义的。

from peft import LoraConfig

# LoRA config based on QLoRA paper & Sebastian Raschka experiment
peft_config = LoraConfig(
        lora_alpha=128,
        lora_dropout=0.05,
        r=256,
        bias="none",
        target_modules="all-linear",
        task_type="CAUSAL_LM", 
)

在开始训练之前,我们需要定义要使用的超参数(TrainingArguments)。

from transformers import TrainingArguments

args = TrainingArguments(
    output_dir="code-llama-7b-text-to-sql", # directory to save and repository id
    num_train_epochs=3,                     # number of training epochs
    per_device_train_batch_size=3,          # batch size per device during training
    gradient_accumulation_steps=2,          # number of steps before performing a backward/update pass
    gradient_checkpointing=True,            # use gradient checkpointing to save memory
    optim="adamw_torch_fused",              # use fused adamw optimizer
    logging_steps=10,                       # log every 10 steps
    save_strategy="epoch",                  # save checkpoint every epoch
    learning_rate=2e-4,                     # learning rate, based on QLoRA paper
    bf16=True,                              # use bfloat16 precision
    tf32=True,                              # use tf32 precision
    max_grad_norm=0.3,                      # max gradient norm based on QLoRA paper
    warmup_ratio=0.03,                      # warmup ratio based on QLoRA paper
    lr_scheduler_type="constant",           # use constant learning rate scheduler
    push_to_hub=True,                       # push model to hub
    report_to="tensorboard",                # report metrics to tensorboard
)

现在,我们拥有了创建 SFTTrainer 所需的所有构件,可以开始训练我们的模型了。

from trl import SFTTrainer

max_seq_length = 3072 # max sequence length for model and packing of the dataset

trainer = SFTTrainer(
    model=model,
    args=args,
    train_dataset=dataset,
    peft_config=peft_config,
    max_seq_length=max_seq_length,
    tokenizer=tokenizer,
    packing=True,
    dataset_kwargs={
        "add_special_tokens": False,  # We template with special tokens
        "append_concat_token": False, # No need to add additional separator token
    }
)

​ 调用训练器实例的 train() 方法,开始训练模型。这将启动训练循环,对模型进行 3 轮训练。由于我们使用的是 PEFT 方法,因此只会保存调整后的模型权重,而不会保存完整的模型。

# start training, the model will be automatically saved to the hub and the output directory
trainer.train()

# save model 
trainer.save_model()

在 g5.2xlarge 上使用 Flash Attention 对 10k 个样本的数据集进行了 3 轮训练,耗时 01:29:58。实例成本为 1,212$/h,因此总成本仅为1.8$

# free the memory again
del model
del trainer
torch.cuda.empty_cache()

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

​ 使用 QLoRA 时,我们只训练适配器而不是完整模型。这意味着在训练过程中保存模型时,我们只保存适配器权重,而不保存完整模型。如果您想保存完整模型,使其更容易与文本生成推理一起使用,您可以使用 merge_and_unload 方法将适配器权重合并到模型权重中,然后使用 save_pretrained 方法保存模型。这将保存一个默认模型,可用于推理。

注意:您可能需要 > 30GB 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")

5.测试模型和运行推理

​ 训练完成后,我们要对模型进行评估和测试。我们将从原始数据集中加载不同的样本,并在这些样本上对模型进行评估,使用简单的循环和准确率作为衡量标准。

注:评估生成式人工智能模型并非易事,因为一个输入可能有多个正确的输出。如果您想了解有关评估生成模型的更多信息,请查看 Evaluate LLMs and RAG a practical example using Langchain and Hugging Face 《评估 LLMs 和 RAG》(使用 Langchain 和 Hugging Face 的实用示例)。

import torch
from peft import AutoPeftModelForCausalLM
from transformers import AutoTokenizer, pipeline 

peft_model_id = "./code-llama-7b-text-to-sql"
# peft_model_id = args.output_dir

# Load Model with PEFT adapter
model = AutoPeftModelForCausalLM.from_pretrained(
  peft_model_id,
  device_map="auto",
  torch_dtype=torch.float16
)
tokenizer = AutoTokenizer.from_pretrained(peft_model_id)
# load into pipeline
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

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

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

# Test on sample 
prompt = pipe.tokenizer.apply_chat_template(eval_dataset[rand_idx]["messages"][:2], tokenize=False, add_generation_prompt=True)
outputs = pipe(prompt, max_new_tokens=256, do_sample=False, temperature=0.1, top_k=50, top_p=0.1, eos_token_id=pipe.tokenizer.eos_token_id, pad_token_id=pipe.tokenizer.pad_token_id)

print(f"Query:\n{eval_dataset[rand_idx]['messages'][1]['content']}")
print(f"Original Answer:\n{eval_dataset[rand_idx]['messages'][2]['content']}")
print(f"Generated Answer:\n{outputs[0]['generated_text'][len(prompt):].strip()}")

​ 不错!我们的模型能够根据自然语言指令生成 SQL 查询。让我们在测试数据集的全部 2,500 个样本上评估我们的模型。注:如上所述,评估生成模型并非易事。在我们的示例中,我们使用基于基本 SQL 查询生成的 SQL 的准确性作为衡量标准。另一种方法是自动执行生成的 SQL 查询,并将结果与地面实况进行比较。这将是一个更准确的衡量标准,但需要更多的设置工作。

from tqdm import tqdm

def evaluate(sample):
    prompt = pipe.tokenizer.apply_chat_template(sample["messages"][:2], tokenize=False, add_generation_prompt=True)
    outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95, eos_token_id=pipe.tokenizer.eos_token_id, pad_token_id=pipe.tokenizer.pad_token_id)
    predicted_answer = outputs[0]['generated_text'][len(prompt):].strip()
    if predicted_answer == sample["messages"][2]["content"]:
        return 1 
    else:
        return 0

success_rate = []
number_of_eval_samples = 1000
# iterate over eval dataset and predict
for s in tqdm(eval_dataset.shuffle().select(range(number_of_eval_samples))):
    success_rate.append(evaluate(s))

# compute accuracy
accuracy = sum(success_rate)/len(success_rate)

print(f"Accuracy: {accuracy*100:.2f}%")

​ 我们在评估数据集中的 1000 个样本上评估了我们的模型,准确率为 79.50%,耗时约 25 分钟。这是相当不错的结果,但如前所述,我们需要慎重对待这一指标。如果我们能通过在真实数据库中运行 Qureies 来评估我们的模型并比较结果,效果会更好。因为同一指令可能有不同的 "正确 "SQL 查询。我们还可以通过几种方法来提高性能,如使用少量学习、使用 RAG、自愈来生成 SQL 查询。

6.在生产环境部署LLM

​ 现在您可以将模型部署到生产环境中。要将开放式 LLM 部署到生产环境中,我们建议使用文本生成推理 (TGI)。TGI 是专门为部署和服务大型语言模型 (LLM) 而设计的解决方案。TGI 使用张量并行和连续批处理技术为最流行的开放式 LLM(包括 Llama、Mistral、Mixtral、StarCoder、T5 等)实现高性能文本生成。IBM、Grammarly、Uber、德国电信等公司都在使用文本生成推理。有多种方式可以部署您的模型,包括:

如果你已经安装了 docker,可以使用以下命令启动推理服务器。

注意:确保有足够的 GPU 内存来运行容器。重启内核,删除笔记本上所有分配的 GPU 内存。

%%bash 
# model=$PWD/{args.output_dir} # path to model
model=$(pwd)/code-llama-7b-text-to-sql # path to model
num_shard=1             # number of shards
max_input_length=1024   # max input length
max_total_tokens=2048   # max total tokens

docker run -d --name tgi --gpus all -ti -p 8080:80 \
  -e MODEL_ID=/workspace \
  -e NUM_SHARD=$num_shard \
  -e MAX_INPUT_LENGTH=$max_input_length \
  -e MAX_TOTAL_TOKENS=$max_total_tokens \
  -v $model:/workspace \
  ghcr.io/huggingface/text-generation-inference:latest

容器运行后,您就可以发送请求了。

import requests as r 
from transformers import AutoTokenizer
from datasets import load_dataset
from random import randint

# Load our test dataset and Tokenizer again
tokenizer = AutoTokenizer.from_pretrained("code-llama-7b-text-to-sql")
eval_dataset = load_dataset("json", data_files="test_dataset.json", split="train")
rand_idx = randint(0, len(eval_dataset))

# generate the same prompt as for the first local test
prompt = tokenizer.apply_chat_template(eval_dataset[rand_idx]["messages"][:2], tokenize=False, add_generation_prompt=True)
request= {"inputs":prompt,"parameters":{"temperature":0.2, "top_p": 0.95, "max_new_tokens": 256}}

# send request to inference server
resp = r.post("http://127.0.0.1:8080/generate", json=request)

output = resp.json()["generated_text"].strip()
time_per_token = resp.headers.get("x-time-per-token")
time_prompt_tokens = resp.headers.get("x-prompt-tokens")

# Print results
print(f"Query:\n{eval_dataset[rand_idx]['messages'][1]['content']}")
print(f"Original Answer:\n{eval_dataset[rand_idx]['messages'][2]['content']}")
print(f"Generated Answer:\n{output}")
print(f"Latency per token: {time_per_token}ms")
print(f"Latency prompt encoding: {time_prompt_tokens}ms")

真棒!完成后,别忘了关闭容器。

!docker stop tgi

本文翻译自:How to Fine-Tune LLMs in 2024 with TRL

评论

  1. test1111
    1年前
    2024-2-14 12:44:49

    111

发送评论 编辑评论


				
|´・ω・)ノ
ヾ(≧∇≦*)ゝ
(☆ω☆)
(╯‵□′)╯︵┴─┴
 ̄﹃ ̄
(/ω\)
∠( ᐛ 」∠)_
(๑•̀ㅁ•́ฅ)
→_→
୧(๑•̀⌄•́๑)૭
٩(ˊᗜˋ*)و
(ノ°ο°)ノ
(´இ皿இ`)
⌇●﹏●⌇
(ฅ´ω`ฅ)
(╯°A°)╯︵○○○
φ( ̄∇ ̄o)
ヾ(´・ ・`。)ノ"
( ง ᵒ̌皿ᵒ̌)ง⁼³₌₃
(ó﹏ò。)
Σ(っ °Д °;)っ
( ,,´・ω・)ノ"(´っω・`。)
╮(╯▽╰)╭
o(*////▽////*)q
>﹏<
( ๑´•ω•) "(ㆆᴗㆆ)
😂
😀
😅
😊
🙂
🙃
😌
😍
😘
😜
😝
😏
😒
🙄
😳
😡
😔
😫
😱
😭
💩
👻
🙌
🖕
👍
👫
👬
👭
🌚
🌝
🙈
💊
😶
🙏
🍦
🍉
😣
Source: github.com/k4yt3x/flowerhd
颜文字
Emoji
小恐龙
花!
上一篇
下一篇