点击上方蓝字关注我们
trl
和SFTTrainer
微调LLM
一、定义我们的使用场景
我们想要微调一个模型,它可以基于自然语言指令生成SQL查询,然后可以集成到我们的BI工具中。目标是减少创建SQL查询所需的时间,并使非技术用户更容易创建SQL查询。
二、设置开发环境
# 安装Pytorch和其他库
!pip install "torch==2.1.2" tensorboard
# 安装Hugging Face库
!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" \
# 从github安装peft & trl
!pip install git+https://github.com/huggingface/trl@a3c5b7178ac4f65569975efadc97db2f3749c65e --upgrade
!pip install git+https://github.com/huggingface/peft@4a1559582281fc3c9283892caea8ccef1d6f5a4f--upgrade
如果你的GPU采用的是Ampere架构(如NVIDIA A10G或RTX 4090/3090)或更新版本,你可以利用Flash Attention技术。Flash Attention通过优化注意力机制的计算过程,并采用一些经典技术(如分块和重新计算)来显著提高计算速度,并降低内存消耗。简而言之,这项技术可以将训练速度提升至原来的三倍。想要了解更多详情,可以访问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分钟)。
huggingface_hub
包中的登录工具来登录你的账户,并在本地磁盘上保存你的访问令牌。from huggingface_hub import login
login(
token="", # 在此处添加您的token
add_to_git_credential=True
)
三、创建和准备数据集
一旦您确定微调是正确的解决方案,我们需要准备一个数据集来训练我们的模型。这个数据集应该是多样化的任务示范,展示了你想要解决的问题。创建数据集的方法有很多,比如:
利用现有的开源数据集,如Spider
利用大语言模型生成合成数据集,如Alpaca
雇佣人类来创建数据集,如Dolly
结合以上方法,如Orca
每种方法都有其自身的优势和劣势,并取决于预算、时间和质量要求。例如,使用现有数据集是最简单的,但可能不针对你的特定使用场景,而人工创建数据集虽然准确度高,但成本和时间消耗也大。也可以将几种方法结合起来创建指令数据集,如Orca: Progressive Learning from Complex Explanation Traces of GPT-4所示。
在我们的示例中,我们将使用一个名为sql-create-context
的现有数据集,它包含了自然语言指令、数据库模式定义以及相应的SQL查询样本。
随着trl的最新版本发布,我们现在支持流行的指令和对话数据集格式。这意味着我们只需将数据集转换为支持的格式之一,trl就会自动处理后续步骤。支持的格式包括:
{"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": "..."}]}
{"prompt": "<prompt text>", "completion": "<ideal generated text>"}{"prompt": "<prompt text>", "completion": "<ideal generated text>"}{"prompt": "<prompt text>", "completion": "<ideal generated text>"}
我们将使用Hugging Face的Datasets库来加载我们的开源数据集,并将其转换为对话格式。在这种格式中,我们将在系统消息中包含数据库模式定义,作为我们助手的信息。然后,我们将数据集保存为jsonl文件,这样就可以用于微调我们的模型。我们对数据集进行了随机下采样,只保留了10,000个样本。
from datasets import load_dataset
# 将数据集转换为OAI消息
system_message = """您是SQL查询翻译的文本。用户将用英语向您提问,您将根据提供的SCHEMA生成SQL查询。
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"]}
]
}
# 从hub加载数据集
dataset = load_dataset("b-mc2/sql-create-context", split="train")
dataset = dataset.shuffle().select(range(12500))
# 将数据集转换为OAI消息
dataset = dataset.map(create_conversation, remove_columns=dataset.features,batched=False)
# 将数据集拆分为10000个训练样本和2500个测试样本
dataset = dataset.train_test_split(test_size=2500/12500)
print(dataset["train"][345]["messages"])
# 将数据集保存到磁盘
dataset["train"].to_json("train_dataset.json", orient="records")
dataset["test"].to_json("test_dataset.json", orient="records")
四、使用trl和SFTTrainer微调大语言模型
from datasets import load_dataset
# 从磁盘加载jsonl数据
dataset = load_dataset("json", data_files="train_dataset.json", split="train")
接下来,我们将加载我们的大语言模型。在我们的应用场景中,我们选择了CodeLlama 7B,这是一个专门为代码合成和理解训练的大语言模型。如果你有其他偏好,比如Mistral
、Mixtral
模型,或者TII Falcon
,只需调整我们的model_id
即可轻松切换。我们将使用bitsandbytes工具将模型量化为4位,以减少内存需求。
setup_chat_format
,它向标记器添加特殊标记,例如<|im_start|>
和<|im_end|>
,以表示对话的开始和结束。调整模型嵌入层的大小以适应新的标记。设置标记器的chat_template,它用于将输入数据格式化为类似于聊天的格式。默认是来自OpenAI的chatml。<|im_start|>
和结束对话的<|im_end|>
,来教会它们在对话中扮演的角色。在trl库中,我们有一个名为setup_chat_format
的便捷方法,它: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' # 以防止警告
# 将聊天模板设置为OAI chatML,如果您从微调模型开始,请删除
model, tokenizer = setup_chat_format(model, tokenizer)
SFTTrainer
与peft
的集成使得使用QLoRA高效调优LLM变得非常简单。我们只需要创建我们的LoraConfig并将其提供给训练器。我们的LoraConfig
参数是根据QLoRA论文定义的。
from peft import LoraConfig
# 基于QLoRA论文和Sebastian Raschka实验的LoRA配置
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", # 要保存的目录和存储库ID
num_train_epochs=3, # 训练周期数
per_device_train_batch_size=3, # 训练期间每个设备的批量大小
gradient_accumulation_steps=2, # 反向/更新前的步骤数
gradient_checkpointing=True, # 使用渐变检查点来节省内存
optim="adamw_torch_fused", # 使用融合的adamw优化器
logging_steps=10, # 每10步记录一次
save_strategy="epoch", # 每个epoch保存检查点
learning_rate=2e-4, # 学习率,基于QLoRA论文
bf16=True, # 使用bfloat16精度
tf32=True, # 使用tf32精度
max_grad_norm=0.3, # 基于QLoRA论文的最大梯度范数
warmup_ratio=0.03, # 根据QLoRA论文的预热比例
lr_scheduler_type="constant", # 使用恒定学习率调度器
push_to_hub=True, # 将模型推送到Hub
report_to="tensorboard", # 将指标报告到Tensorboard
)
SFTTrainer
并启动模型训练的所有要素。from trl import SFTTrainer
max_seq_length = 3072 # 数据集模型和打包的最大序列长度
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, # 我们使用特殊 tokens
"append_concat_token": False, # 不需要添加额外的分隔符 token
}
)
Trainer
实例的 train()
方法来启动模型训练。这将启动一个训练循环,持续 3 个周期。由于我们采用的是参数高效微调方法,我们只会保存经过调整的模型权重,而不是整个模型。# 开始训练,模型会自动保存到hub和输出目录
trainer.train()
# 保存模型
trainer.save_model()
g5.2xlarge
上花费了01:29:58的时间。实例成本为1,212$/h
,这使得总成本仅为1.8$
。# 再次释放内存
del model
del trainer
torch.cuda.empty_cache()
merge_and_unload
方法将适配器权重合并到模型权重中,然后使用 save_pretrained
方法保存模型。这将保存一个默认模型,可用于推理。#### COMMENT IN TO MERGE PEFT AND BASE MODEL ####
# from peft import PeftModel, PeftConfig
# from transformers import AutoModelForCausalLM, AutoTokenizer
# from peft import AutoPeftModelForCausalLM
# # Load PEFT model on CPU
# config = PeftConfig.from_pretrained(args.output_dir)
# model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path,low_cpu_mem_usage=True)
# tokenizer = AutoTokenizer.from_pretrained(args.output_dir)
# model.resize_token_embeddings(len(tokenizer))
# model = PeftModel.from_pretrained(model, args.output_dir)
# 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")
五、测试和评估大语言模型
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
)
# load into pipeline
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
让我们加载测试数据集,尝试生成一个指令。
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))
# 样品测试
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
# 迭代eval数据集并预测
for s in tqdm(eval_dataset.shuffle().select(range(number_of_eval_samples))):
success_rate.append(evaluate(s))
# 计算精度
accuracy = sum(success_rate)/len(success_rate)
print(f"Accuracy: {accuracy*100:.2f}%")
我们在评估数据集的 1000 个样本上进行了测试,准确率达到了 79.50%
,整个过程大约花费了 25 分钟。
六、将大语言模型部署到生产环境
%%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
# 再次加载我们的测试数据集和Tokenizer
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))
# 生成与第一次本地测试相同的提示
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}}
# 向推理服务器发送请求
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(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
七、总结
八、References
[1]. Llama 2 (https://huggingface.co/meta-llama/Llama-2-70b-chat-hf)
[2]. Mistral(https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2)
[3]. Mixtral (https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1)
[4]. Falcon(https://huggingface.co/tiiuae/falcon-40b)
[5]. TRL(https://huggingface.co/docs/trl/index)
[6]. Transformers(https://huggingface.co/docs/transformers/index)
[7]. datasets(https://huggingface.co/docs/datasets/index)
[8]. FlashAttention (https://github.com/Dao-AILab/flash-attention/tree/main)
[9]. Hugging Face Hub (https://huggingface.co/models)
[10]. Orca: Progressive Learning from Complex Explanation Traces of GPT-4. (https://arxiv.org/abs/2306.02707)
[11]. sql-create-context (https://huggingface.co/datasets/b-mc2/sql-create-context)
[12]. Text Generation Inference (TGI) https://github.com/huggingface/text-generation-inference
如果你对这篇文章感兴趣,而且你想要了解更多关于AI领域的实战技巧,可以关注「技术狂潮AI」公众号。在这里,你可以看到最新最热的AIGC领域的干货文章和案例实战教程。