ReFT 微调Llama3

介绍

参数高效微调（PEFT）方法旨在通过仅涉及几个参数的权重更新来调整大型语言模型。尽管如此，之前关于可解释性的大多数研究表明，表示包含大量的语义信息，这意味着编辑表示可能是一个更有效的选择。这就是 Representation Finetuning （ReFT）方法的用武之地。事实上，LoReFT（ReFT 系列的一部分）是现有 PEFT 的直接替代品，并且学习的干预措施比以前最先进的 PEFT 高 10 到 50 倍。

在这方面，Zhengxuan Wu 等人的论文表明，当前最先进的 PEFT 的标志是它们修改权重而不是表示。但是，他们的工作表明，编辑表示可能是比权重更新更强大、更高效的替代方案。

ReFT （表示微调）

在进入 ReFT 之前，我们首先需要理解 Representations 的含义。我们知道语言模型会产生标记序列的上下文化表示。给定一个 n 个输入标记 x = （x1， . . . . ， xn）的序列，模型首先将这些标记嵌入到一个表示列表中。然后，模型层将第 j 个隐藏表示 h （j）列表作为前一个隐藏表示 h （j−1）列表的函数连续计算。每个隐藏的表示都是一个向量。LM 使用最终的隐藏表示来生成其预测。

ReFT 的动机源于基于干预的模型可解释性的概念，它强调改变表示而不是权重。这个概念基于线性表示假说，该假说指出概念被编码在神经网络表示的线性子空间中。

在本文中，他们决定使用分布式交换干预操作来制作一种新的参数高效方法，用于使语言模型适应下游任务。如果您想更深入地了解这个概念，请随时参考他们的论文。

PyReFT

PyReFT 是一个表示微调（ReFT）库，支持通过可训练的干预来调整内部语言模型表示。Pyreft 具有更少的微调参数和更强大的性能，可以提高微调效率，降低微调成本，同时为研究自适应参数的可解释性打开大门。

Pyreft 支持：

使用 ReFT 微调 HuggingFace 上的任何预训练 LM
通过配置设置 ReFT 超参数
轻松将微调结果分享到 HuggingFace

使用 PyReFT 微调 Llama-3

由于缺乏计算资源，我在 teknium/OpenHermes-2.5 数据集的 10k 子集上对 Llama-3-8b 进行了 1 个时期的微调。请随意在完整的数据集上试用。

安装依赖项

第一步是安装 Pyreft 库。如果已安装，将导入 pyreft。

try:
    import pyreft
except ModuleNotFoundError:
    !pip install git+https://github.com/stanfordnlp/pyreft.git

Pip 安装最新版本的 transformers 或支持 llama-3 的版本。此外，您还需要 bitsandbytes 库。

!pip install -q git+https://github.com/huggingface/transformers
!pip install -q bitsandbytes

确保您的 huggingface 可以访问封闭的 Llama-3 模型，并且您必须登录您的 huggingface 帐户。请使用下面的代码片段。

from huggingface_hub import notebook_login
notebook_login()

Load Model 和 Tokenizer

下一步是设置用于训练的提示模板。由于我们将使用基本模型，因此我们需要添加特殊标记，以便模型可以学习停止并且不要继续生成文本。使用下面的代码片段加载 model 和 tokenizer。

import torch, transformers, pyreft
device = "cuda"
prompt_no_input_template = """<|begin_of_text|><|start_header_id|>user<|end_header_id|>%s<|eot_id|><|start_header_id|>assistant<|end_header_id|>"""
model_name_or_path = "meta-llama/Meta-Llama-3-8B"
model = transformers.AutoModelForCausalLM.from_pretrained(
    model_name_or_path, torch_dtype=torch.bfloat16, device_map=device, trust_remote_code=True)
# # get tokenizer
tokenizer = transformers.AutoTokenizer.from_pretrained(
    model_name_or_path, model_max_length=2048, 
    padding_side="right", use_fast=False)
tokenizer.pad_token = tokenizer.eos_token

为 ReFT FineTuning 准备模型

为模型设置 pyreft 配置，然后使用 pyreft.get_reft_model（）方法让模型准备好进行表示微调。对于配置，我们将在第 15 层对最后一个提示令牌的残差流应用单个 rank-4 LoReFT 干预。

# get reft model
reft_config = pyreft.ReftConfig(representations={
    "layer": 8, "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")
reft_model.print_trainable_parameters()

准备数据集

准备数据集以进行微调。我使用了 OpenHermes-2.5 数据集的 10k 子集。由于 Reft Trainer 希望数据采用特定格式，因此我们将使用 pyreft.make_last_position_supervised_data_module（）来准备数据。

dataset_name = "teknium/OpenHermes-2.5"
from datasets import load_dataset
dataset = load_dataset(dataset_name, split="train")
dataset = dataset.select(range(10_000))
data_module = pyreft.make_last_position_supervised_data_module(
    tokenizer, model, [prompt_no_input_template % row["conversations"][0]["value"] for row in dataset], 
    [row["conversations"][1]["value"] for row in dataset])

开始培训

现在我们将为 pyreft 设置 training args。ReftTrainerForCausalLM（）的 API 请求。请根据 ypur 自己的用例和计算资源随意更改它。我将只训练 1 个 epoch 的模型。我尝试集成 wandb，但我认为目前 wandb 的集成存在问题。

# train
training_args = transformers.TrainingArguments(
    per_device_train_batch_size = 4,
    gradient_accumulation_steps = 8,
    warmup_steps = 100,
    num_train_epochs = 1,
    learning_rate = 5e-4,
    bf16 = True,
    logging_steps = 1,
    optim = "paged_adamw_32bit",
    weight_decay = 0.0,
    lr_scheduler_type = "cosine",
    output_dir = "outputs",
    report_to=[]
)
trainer = pyreft.ReftTrainerForCausalLM(model=reft_model, tokenizer=tokenizer, args=training_args, **data_module)
_ = trainer.train()

完成训练后，将干预块保存到目录中。

reft_model.save(
    save_directory="./reft_to_share", 
)

推理

要进行推理，请将 device 设置为 cuda。加载基本模型并通过将干预模块与其合并来准备 reft 模型。然后将 reft 模型转移到 cuda，否则将无法进行推理。

import torch, transformers, pyreft
device = "cuda"
model_name_or_path = "meta-llama/Meta-Llama-3-8B"
model = transformers.AutoModelForCausalLM.from_pretrained(
    model_name_or_path, torch_dtype=torch.bfloat16, device_map=device)
reft_model = pyreft.ReftModel.load(
    "Syed-Hasan-8503/Llama-3-openhermes-reft", model, from_huggingface_hub=True
)
reft_model.set_device("cuda")

设置您选择的指令并运行下面的代码片段。恭喜！ 您刚刚推断了您的第一个 REFT 微调模型。

instruction = "A rectangular garden has a length of 25 feet and a width of 15 feet. If you want to build a fence around the entire garden, how many feet of fencing will you need?"
# tokenize and prepare the input
prompt = prompt_no_input_template % instruction
prompt = tokenizer(prompt, return_tensors="pt").to(device)
base_unit_location = prompt["input_ids"].shape[-1] - 1  # last position
_, 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))