PyReFT: 一个强大的表征微调库助力语言模型高效适配

pyreft

PyReFT简介

在自然语言处理和机器学习领域,研究者们一直在寻求如何高效地适配和微调大型语言模型以应对特定任务。参数高效微调(PEFT)方法旨在通过更新模型中少量的权重来实现这一目标。然而,之前的可解释性研究表明,模型的表征中包含了丰富的语义信息,这意味着直接编辑这些表征可能是一种更有效的方法。

这就是表征微调(ReFT)方法的由来。ReFT方法在冻结的基础模型上操作,学习针对特定任务的隐藏表征干预。在ReFT家族中,低秩线性子空间ReFT(LoReFT)是一个突出的代表。LoReFT可以作为现有PEFT方法的直接替代,但它学习的干预比之前最先进的PEFT方法参数效率高10-50倍。LoReFT在效率和性能之间取得了最佳平衡,使其成为研究人员和实践者在处理大型模型时的理想选择。

为了进一步促进ReFT方法的应用,一个通用的ReFT训练库PyReFT已经公开发布,使社区能够探索和利用这些前沿技术。

PyReFT logo

PyReFT的主要特性

PyReFT作为一个强大的表征微调库,具有以下几个关键特性:

支持对HuggingFace上的预训练语言模型进行ReFT微调
通过配置文件灵活设置ReFT超参数
方便地将微调结果分享到HuggingFace
可自定义训练器,如与ReFT结合使用DPO等

这些特性使PyReFT成为一个灵活且易用的工具,可以帮助研究人员和开发者快速实现高效的模型微调。

PyReFT的工作原理

PyReFT的核心思想是通过干预模型的内部表征来实现微调,而不是直接更新模型权重。这种方法有几个优势:

参数效率高:只需要学习少量干预参数就可以实现有效的微调。
性能强大:通过直接操作表征,可以更好地捕捉任务相关的语义信息。
可解释性强:干预可以直观地反映模型学到了什么。

PyReFT支持多种干预方式,其中最具代表性的是LoReFT(Low-rank Linear Subspace ReFT)。LoReFT在模型的某些层上应用低秩线性变换,实现高效且有效的表征调整。

使用PyReFT进行模型微调

下面我们通过一个简单的例子来演示如何使用PyReFT对Llama-2模型进行微调:

首先安装PyReFT:

pip install git+https://github.com/stanfordnlp/pyreft.git

加载预训练模型:

import torch, transformers, pyreft

model_name_or_path = "meta-llama/Llama-2-7b-chat-hf"
model = transformers.AutoModelForCausalLM.from_pretrained(
    model_name_or_path, torch_dtype=torch.bfloat16, device_map="cuda")

tokenizer = transformers.AutoTokenizer.from_pretrained(
    model_name_or_path, model_max_length=2048, 
    padding_side="right", use_fast=False)
tokenizer.pad_token = tokenizer.unk_token

设置ReFT配置:

reft_config = pyreft.ReftConfig(representations={
    "layer": 15, "component": "block_output",
    "low_rank_dimension": 4,
    "intervention": pyreft.LoreftIntervention(embed_dim=model.config.hidden_size,
    low_rank_dimension=4)})
reft_model = pyreft.get_reft_model(model, reft_config)
reft_model.set_device("cuda")

准备训练数据:

training_examples = [
    ["Who are you?", "🤖💬🌐🧠"],
    ["Who am I?", "👤❓🔍🌟"],
    # ... 更多示例
]

data_module = pyreft.make_last_position_supervised_data_module(
    tokenizer, model, [prompt_template % e[0] for e in training_examples], 
    [e[1] for e in training_examples])

训练模型:

training_args = transformers.TrainingArguments(
    num_train_epochs=100.0, output_dir="./tmp", per_device_train_batch_size=10, 
    learning_rate=4e-3, logging_steps=20)
trainer = pyreft.ReftTrainerForCausalLM(
    model=reft_model, tokenizer=tokenizer, args=training_args, **data_module)
_ = trainer.train()

使用微调后的模型:

instruction = "Which dog breed do people think is cuter, poodle or doodle?"
prompt = prompt_template % instruction
prompt = tokenizer(prompt, return_tensors="pt").to("cuda")

base_unit_location = prompt["input_ids"].shape[-1] - 1
_, reft_response = reft_model.generate(
    prompt, unit_locations={"sources->base": (None, [[[base_unit_location]]])},
    intervene_on_prompt=True, max_new_tokens=512, do_sample=True, 
    eos_token_id=tokenizer.eos_token_id, early_stopping=True
)
print(tokenizer.decode(reft_response[0], skip_special_tokens=True))