使用 Hugging Face TRL 微调大型语言模型以实现特定任务

像 CodeGPT 这样的大型语言模型 (LLM) 已经取得了快速进展,在摘要、搜索等方面实现了新的应用程序。不过,现成的模型可能不适合特定的用例。微调是定制模型的关键,但需要专业知识。本文展示了如何使用 Hugging Face 的 Transformer Library (TRL) 微调 LLM 以实现自定义任务。

使用 Hugging Face TRL 微调大型语言模型以实现特定任务

概述

我们将:

  • 定义示例用例:将自然语言转换为 SQL 查询
  • 使用 TRL 和 Flash Attention 等工具设置环境以提高效率
  • 准备数据集并将其转换为 TRL 的对话格式
  • 使用参数高效微调 (PEFT) 方法微调 CodeLlama-7B
  • 评估测试集并使用文本生成推理部署模型

目标是让非技术用户能够编写 SQL 查询。我们的数据集有英文说明、模式定义和目标 SQL 查询。

设置

我们安装 TRL、transformers等。Flash Attention 使用模型并行性来加速最新 NVIDIA GPU 上的训练:

!pip install trl transformers datasets accelerateimport torch; assert torch.cuda.get_device_capability()[0] >= 8 
!pip install flash-attn

我们将使用 Huggingface Hub 来跟踪实验。登录后,模型将自动推送到我们的 Hub 存储库:

from huggingface_hub import login
login(token="my_token")

数据准备

我们使用 数据集加载包含 12,500 个示例的数据集,并将其转换为 TRL 的对话格式,并在对话时使用模式定义和查询:

from datasets import load_dataset

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

dataset = load_dataset("my_sql_dataset")  
dataset = dataset.map(create_conversation) 
dataset.shuffle().select(range(10_000)).train_test_split(2,500)
dataset.save_to_disk("my_dataset")

微调设置

我们加载 CodeLlama-7B 并为与 TRL 的对话添加特殊标记setup_chat_format。为了提高效率,该模型使用 bfloat16 进行 4 位量化:

from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import setup_chat_format   

model = AutoModelForCausalLM.from_pretrained( "codellama/CodeLlama-7B" , 
     quantization_config=BitsAndBytesConfig(load_in_4bit= True )) 
tokenizer = AutoTokenizer.from_pretrained(model_id) 

model, tokenizer = setup_chat_format(model, tokenizer )

我们使用称为 Q-LoRA 的参数高效微调 (PEFT) 方法来减少内存。LoRA 仅训练适配器层,保持大部分权重固定:

from peft import LoraConfig   

peft_config = LoraConfig(lora_alpha= 128 , lora_dropout= 0.05 )

我们的SFTTrainer是Huggingface Trainer 的 TRL 子类,处理打包、梯度累积等。我们使用针对 Q-LoRA 调整的学习率训练 3 个 epoch:

from trl import SFTTrainer 

trainer = SFTTrainer( 
    model=model,   
    train_dataset= "my_dataset" , 
    peft_config=peft_config, 
    args={ 
      "learning_rate" : 2e-4 ,   
      "num_train_epochs" : 3 
     }    
) 

trainer.train()

评估

我们评估 1,000 个测试样本的准确性。我们的模型根据指令生成查询的准确率达到 79.5%:

from datasets import load_dataset 

def evaluate(sample):
  // Generate query 
  if generated == target:
    return 1
  else:
    return 0

accuracy = []
for sample in test_set: 
   accuracy.append(evaluate(sample))

print(f"Accuracy: {sum(accuracy)/len(accuracy)}")

服务

对于生产用例,我们推荐 Hugging Face 的文本生成推理。只需几行代码,我们就可以将模型容器化以实现可扩展的部署:

docker run -p 8080 : 8080     
  -v my_model:/workspace  
  ghcr.io/huggingface/text- Generation-inference:latest

我们通过 7 个步骤定制了一个 LLM,用于使用数据集助手、高效训练方法和张量并行等 TRL 工具将文本翻译为 SQL。一样的原则适用于跨域的其他自定义任务。

PS:AI火箭营重磅推出大模型学院,扫下面二维码即可进行在线学习。

使用 Hugging Face TRL 微调大型语言模型以实现特定任务

本文重点讲解微调实践,针对微调大模型学院开设有以技能实践为主的课程《大模型微调技术》,主讲老师:Michael,清华大学硕士毕业,拥有深度学习与大模型相关丰富的应用项目经验。

使用 Hugging Face TRL 微调大型语言模型以实现特定任务

© 版权声明
THE END
如果内容对您有所帮助,就支持一下吧!
点赞0 分享
评论 共1条

请登录后发表评论