使用TRL微调LLM(2024年)
大型语言模型(LLMs)在过去一年中取得了许多进展。我们从现在ChatGPT的竞争对手发展到一个包含Meta AI的Llama 3,Mistral的Mistral和Mixtral模型,TII的Falcon,以及许多其他模型。
本文博客将指导您如何使用Hugging Face进行开放LLM的微调TRL,Transformers & datasets在2024年。在博客中,我们将:
- 定义我们的使用案例
- 设置开发环境
- 创建和准备数据集
- 使用
trl微调 LLMSFTTrainer - 测试和评估LLM
- 部署生产环境的LLM
注意:此博客是为消费者级别的GPU(24GB)创建的,例如NVIDIA A10G或RTX 4090/3090,但可以轻松改编以在更大的GPU上运行。
1. 定义我们的使用案例
在微调LLMs时,重要的是你知道你的使用案例和你想要解决的任务。这将帮助你选择合适的模型或帮助你创建一个用于微调模型的数据集。如果你还没有定义你的使用案例,你可能需要回到起点。 我想提一下,并非所有的使用案例都需要微调,在微调你自己的模型之前,总是建议评估和尝试已经微调的模型或基于API的模型。
作为一个例子,我们将使用以下用例:
我们希望微调一个模型,该模型可以根据自然语言指令生成SQL查询,然后可以将其整合到我们的BI工具中。目标是减少创建SQL查询所需的时间,并使非技术用户更轻松地创建SQL查询。
文本到SQL可以成为微调LLMs的一个很好的用例,因为这是一个需要对数据和SQL语言有大量(内部)知识的复杂任务。
2. 设置开发环境
我们的第一步是安装 Hugging Face 库和 Pyroch,包括trl、transformers 和 datasets。如果你还没有听说过trl,别担心。trl 是一个基于 transformers 和 datasets 的新库,它使微调、rlhf、对齐开放 LLM 变得更容易。
# Install Pytorch & other libraries %pip install "torch==2.4.0" tensorboard # Install Hugging Face libraries %pip install --upgrade \ "transformers==4.44.2" \ "datasets==2.21.0" \ "accelerate==0.33.0" \ "evaluate==0.4.2" \ "bitsandbytes==0.43.3" \ "trl==0.9.6" \ "peft==0.12.0"
如果您使用的是具有Ampere架构的GPU(例如NVIDIA A10G或RTX 4090/3090)或更新的型号,您可以使用Flash attention。Flash Attention是一种重新排序注意力计算并利用经典技术(平铺、重新计算)的方法,可以显著加快计算速度并减少内存使用量,从序列长度的平方减少到线性。简而言之;加速训练最多3倍。了解更多:FlashAttention.
注意:如果您的机器内存小于96GB且CPU核心数很多,请减少MAX_JOBS的数量。在g6.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-attn可能需要相当长的时间(10-45分钟)。
3. 创建和准备数据集
一旦您确定微调是正确的解决方案,我们需要创建一个数据集来微调我们的模型。该数据集应是您要解决的任务的多样示例集合。创建这种数据集的方法有多种,包括:
每种方法都有其自身的优点和缺点,并且取决于预算、时间以及质量要求。例如,使用现有的数据集是最简单的,但可能无法针对您的特定用例进行定制,而使用人工可能最准确,但可能耗时且昂贵。也可以将几种方法结合起来创建一个指令数据集,如 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>"}
在我们的示例中,我们将使用🤗 Datasets库加载我们的开源数据集,然后将其转换为对话格式,在其中包含助手的系统消息中的模式定义。然后,我们将数据集保存为jsonl文件,然后我们可以使用它来微调我们的模型。我们随机对数据集进行下采样,只保留10,000个样本。
注意:此步骤可能因您的使用案例而异。例如,如果您已经从 OpenAI 等渠道获得了一个数据集,您可以跳过此步骤并直接进行微调步骤。
map之前dataset
{
'question': 'WHich Score has a To par of –3, and a Country of united states?',
'context': 'CREATE TABLE table_name_28 (score VARCHAR, to_par VARCHAR, country VARCHAR)',
'answer': 'SELECT score FROM table_name_28 WHERE to_par = "–3" AND country = "united states"'
}
Map: 100%|██████████| 12500/12500 [00:01<00:00, 7221.66 examples/s]
map之后dataset
[
{'content': '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.\nSCHEMA:\nCREATE TABLE table_name_93 (tournament VARCHAR)', 'role': 'system'},
{'content': 'Which tournament has a 2007 of 2r, and a 2011 of 2r?', 'role': 'user'},
{'content': 'SELECT tournament FROM table_name_93 WHERE 2007 = "2r" AND 2011 = "2r"', 'role': 'assistant'}
]
Creating json from Arrow format: 100%|██████████| 10/10 [00:00<00:00, 10.88ba/s]
Creating json from Arrow format: 100%|██████████| 3/3 [00:00<00:00, 51.70ba/s]
4. 使用 trl 对 LLM 进行微调 SFTTrainer
我们现在已经准备好微调我们的模型。我们将使用 SFTTrainer 从 trl 来微调我们的模型。 SFTTrainer 使监督开放LLM的微调变得简单。 SFTTrainer 是 Trainer 的一个子类 transformers 库,并支持所有相同的功能,包括日志记录、评估和检查点,但增加了额外的生活质量功能,包括:
- 数据集格式化,包括对话和指令格式
- 仅对完成进行训练,忽略提示
- 打包数据集以实现更高效的训练
- PEFT(参数高效微调)支持包括Q-LoRA
- 为对话微调准备模型和分词器(例如添加特殊标记)
我们将使用我们在示例中的数据集格式化、打包和PEFT功能。作为PEFT方法,我们将使用QLoRA,这是一种在微调过程中通过使用量化来减少大语言模型内存占用的技术,同时不牺牲性能。如果你想了解更多关于QLoRA及其工作原理的信息,请查看使用bitsandbytes、4比特量化和QLoRA让LLM更易获取 博客文章。
现在,让我们开始吧!🚀
首先,我们需要从磁盘加载我们的数据集。
from datasets import load_dataset # Load jsonl data from disk dataset = load_dataset("json", data_files="train_dataset.json", split="train")
接下来,我们将加载我们的LLM。对于我们的使用案例,我们将使用Qwen/Qwen2.5-3B-Instruct。 但我们可以通过更改我们的变量轻松地将模型替换成另一个模型,任何其他LLM。我们将使用bitsandbytes将模型量化到4位。
注意:模型越大,所需的内存就越多。在我们的示例中,我们将使用3B版本,该版本可以在24GB的GPU上微调。如果您有一个较小的GPU。
正确地为训练聊天/对话模型准备LLM和Tokenizer是至关重要的。我们需要向Tokenizer和模型添加新的特殊令牌,并教会它们理解对话中的不同角色。在trl中,我们有一个方便的方法叫做setup_chat_format,它:
- 为分词器添加特殊标记,例如
<|im_start|>和<|im_end|>,以指示对话的开始和结束。 - 将模型的嵌入层调整以适应新的令牌。
- 设置
chat_template,它是用于将输入数据格式化为聊天格式的标记器。默认值是chatml来自 OpenAI。
Loading checkpoint shards: 100%|██████████| 2/2 [01:24<00:00, 42.27s/it]
该 SFTTrainer 支持与 peft的原生集成,这使得使用例如QLoRA等方法高效地调整LLMs变得超级简单。我们只需要创建我们的 LoraConfig 并将其提供给训练器。我们的LoraConfig参数是基于QLoRA论文和sebastian的博客文章定义的。
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)。
# 从transformers库导入TrainingArguments类,用于配置模型训练的各种参数 from transformers import TrainingArguments # 创建训练参数配置对象 args = TrainingArguments( output_dir="text_to_sql", # 模型保存目录,同时也可作为仓库ID(若需上传到Hub) num_train_epochs=3, # 训练总轮数,即整个数据集将被训练3次 per_device_train_batch_size=1, # 每个设备(如GPU)上的训练批次大小,此处设为1(可能受限于显存) gradient_accumulation_steps=8, # 梯度累积步数,每累积8步后再进行一次反向传播和参数更新(可模拟更大批次) gradient_checkpointing=True, # 是否启用梯度检查点,启用后能节省内存(代价是少量计算时间) optim="adamw_torch_fused", # 优化器选择,使用PyTorch的融合AdamW优化器(通常训练速度更快) logging_steps=10, # 日志记录频率,每训练10步记录一次日志(如损失、学习率等) save_strategy="epoch", # 模型保存策略,此处设为每个epoch结束后保存一次 checkpoint learning_rate=2e-4, # 学习率,此处根据QLoRA论文推荐值设置为2e-4 bf16=True, # 是否使用bfloat16混合精度训练(需硬件支持,可加速训练并节省内存) tf32=True, # 是否使用tf32精度(NVIDIA GPU支持,可在保持精度的同时加速计算) max_grad_norm=0.3, # 最大梯度范数,用于梯度裁剪(防止梯度爆炸),参考QLoRA论文设置 warmup_ratio=0.03, # 学习率预热比例,总训练步数的3%用于预热(逐渐提升至设定学习率) lr_scheduler_type="constant", # 学习率调度策略,此处使用恒定学习率(训练过程中保持不变) # push_to_hub=True, # 是否将模型推送到Hugging Face Hub(此处注释掉,不启用) report_to="tensorboard", # 训练指标的报告方式,此处设置为使用TensorBoard可视化 )
我们现在已经有创建我们所需的所有基本要素 SFTTrainer 以开始训练我们的模型。
from trl import SFTTrainer max_seq_length = 2048 # 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 } )
Generating train split: 527 examples [00:02, 183.01 examples/s]
通过调用 train() 方法来开始训练我们的 Trainer 实例。这将启动训练循环,并训练我们的模型3个完整的周期。由于我们使用的是PEFT方法,我们只会保存调整后的模型权重,而不是整个模型。
# start training, the model will be automatically saved to the hub and the output directory trainer.train() # save model trainer.save_model()
[195/195 45:08, Epoch 2/3] Step Training Loss 10 0.729200 20 0.606100 30 0.586000 40 0.569400 50 0.561000 60 0.561200 70 0.530400 80 0.497200 90 0.499000 100 0.486200 110 0.488600 120 0.485100 130 0.488900 140 0.437900 150 0.421900 160 0.424200 170 0.430100 180 0.427900 190 0.421400
使用Flash Attention对一个包含10k样本的数据集进行3个epoch的训练花费了02:05:58在一个AWS EC2 g6.2xlarge实例上。该实例费用为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")
Loading checkpoint shards: 100%|██████████| 2/2 [00:01<00:00, 1.36it/s]
4. 测试模型并进行推理
训练完成后,我们希望评估和测试我们的模型。我们将从原始数据集中加载不同的样本,并在这些样本上评估模型,使用简单的循环和准确率作为我们的指标。
注意:评估生成式AI模型不是一件简单的事情,因为一个输入可以有多个正确的输出。如果你想了解更多关于评估生成模型的信息,请查看使用Langchain和Hugging Face的LLMs和RAG的评估实践示例博客文章。
import torch from transformers import AutoTokenizer, pipeline, AutoModelForCausalLM from transformers import Qwen2ForCausalLM, AutoConfig model_id = "./text_to_sql" # Load Model with PEFT adapter model = AutoModelForCausalLM.from_pretrained( model_id, device_map="auto", torch_dtype=torch.float16 ) tokenizer = AutoTokenizer.from_pretrained(model_id) # load into pipeline pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
Loading checkpoint shards: 100%|██████████| 2/2 [00:02<00:00, 1.36s/it]
让我们加载测试数据集,尝试生成一条指令。
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()}")
Query: What is the maximum ¥ for a round 2 score of 65 and a round 4 score smaller than 67?
Original Answer: SELECT MAX(money__) AS ¥_ FROM table_name_65 WHERE round_2 = 65 AND round_4 < 67
Generated Answer: SELECT MAX(money__) AS $__ FROM table_name_65 WHERE round_2 = "65" AND round_4 < 67
很好!我们的模型能够根据自然语言指令生成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}%")
100%|██████████| 1000/1000 [36:17<00:00, 2.18s/it] Accuracy: 71.40%
我们在评估数据集的1000个样本上评估了我们的模型,准确率为80.00%,这大约需要25分钟。 这相当不错,但正如提到的,你需要对此指标有所保留。如果我们能够通过在真实数据库上运行查询来评估我们的模型,并比较结果,那就更好了。由于同一个指令可能有不同的“正确”SQL查询。我们也可以通过使用少样本学习、使用RAG、自我修复来生成SQL查询来提高性能。
6. 部署LLM用于生产
你现在可以将模型部署到生产环境了。对于将开放的大语言模型(LLM)部署到生产环境,我们建议使用Text Generation Inference (TGI)。TGI是一个为部署和提供大语言模型(LLM)而设计的解决方案。TGI使用张量并行性和连续 batching,以高性能生成文本,支持包括Llama、Mistral、Mixtral、StarCoder、T5等在内的最受欢迎的开放大语言模型。Text Generation Inference 被许多公司使用,包括 IBM、Grammarly、Uber、Deutsche Telekom 等。有几种方法可以部署你的模型,包括:
如果你已经安装了docker,可以使用以下命令启动推理服务器。
注意:确保你有足够的GPU内存来运行容器。重启内核以从笔记本中移除所有已分配的GPU内存。
%%bash # model=$PWD/{args.output_dir} # path to model model=$(pwd)/code-llama-3-1-8b-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:2.2.0
Unable to find image 'ghcr.io/huggingface/text-generation-inference:2.2.0' locally 2.2.0: Pulling from huggingface/text-generation-inference aece8493d397: Already exists 45f7ea5367fe: Already exists 3d97a47c3c73: Already exists 12cd4d19752f: Already exists da5a484f9d74: Already exists 4f4fb700ef54: Already exists 43566b48e5d6: Already exists f165933352a8: Already exists f166ffc7c7b4: Already exists 58165ae83a0e: Already exists 074d930e1b90: Already exists 1033b2636622: Already exists e0aa534acffe: Already exists 130989d28b48: Already exists a65ea9ebfaba: Already exists 7225b2c46f88: Already exists 43154e73908f: Already exists 8f400e318724: Already exists f694acf6c40f: Already exists 44fc79164bc4: Already exists 8bc7c142e917: Already exists 021f7d48bdcb: Already exists c9d01f7d10cc: Already exists 400740bc31be: Already exists bd4b49ea4512: Already exists 141228b9bdde: Already exists 4f4fb700ef54: Already exists 34d4a7457184: Already exists 66e724dff43a: Already exists 25c75c242d08: Already exists 6a4be63c7e70: Already exists b2d83f4bca52: Already exists 373c47aa4b50: Already exists 4f4fb700ef54: Already exists Digest: sha256:d39d513f13727ffa9b6a4d0e949f36413b944aabc9a236c0aa2986c929906769 Status: Downloaded newer image for ghcr.io/huggingface/text-generation-inference:2.2.0 42be7f00ddeb0a3214920a09a5ea303d8eb034942d7020155b6a6761fca87193
一旦你的容器开始运行,你就可以使用openai或huggingface_hub sdk 发送请求。这里我们将使用openai sdk 向我们的推理服务器发送请求。如果你没有安装openai sdk,你可以使用pip install openai进行安装。
from openai import OpenAI from datasets import load_dataset from random import randint # create client client = OpenAI(base_url="http://localhost:8080/v1",api_key="-") # Load our test dataset eval_dataset = load_dataset("json", data_files="test_dataset.json", split="train") rand_idx = randint(0, len(eval_dataset)) # Take a random sample from the dataset and remove the last message and send it to the model response = client.chat.completions.create( model="code-llama-3-1-8b-text-to-sql", messages=eval_dataset[rand_idx]["messages"][:2], stream=False, # no streaming max_tokens=1024, ) response = response.choices[0].message.content # 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{response}")
Query: Name the first elected for kentucky 1 Original Answer: SELECT first_elected FROM table_2668378_5 WHERE district = "Kentucky 1" Generated Answer: SELECT first_elected FROM table_2668378_5 WHERE district = "Kentucky 1"
太棒了,完成后别忘了停止你的容器。
!docker stop tgi

浙公网安备 33010602011771号