之前有一篇文章讲了强化学习PPO算法，算法中我们提到了Reward模型。既然需要一个reward模型，那我们应该怎么去实现了这个过程呢？这里我只是简单的用Qwen2 - 0.5b模型微调了一个lora模型。这里只是为什么要用这么小的模型呢？因为作者的机器带不动呀，而且我也只是想要走一个这个流程，熟悉一下这个reward模型怎么微调。有了入门以后，后续在进行深度优化和提升才会有更好的途径。
接下来我们就从代码端一点一点的讲解一下，我是如何进行reward模型微调的。那这里呢我非常感谢其他博主的开源项目，让我能有幸学习到这些知识点。那接下来我讲从代码端一点一点的讲解我是如何通过学习别人的代码，然后实现reward模型的。
关键包引入
首先我们要先引入相关性的python包。
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union
import torch
import torch.nn as nn
from peft import LoraConfig, TaskType, get_peft_model
from tqdm import tqdm
from transformers import (
AutoTokenizer,
PreTrainedTokenizerBase,
Trainer,
TrainingArguments, BitsAndBytesConfig,
)
from transformers.utils import PaddingStrategy
import datasets
from transformers import Qwen2ForSequenceClassification
import jsontoken的加载
我们做任何模型微调和训练的时候，首先第一步肯定就是加载token，因为不管dataset的处理和模型的训练都是需要token进行处理的。
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token #模型需要我们指定pad_token是什么
DataSet数据数据加载
引入相关性的包以后，我们就要开始对数据进行处理，然后加载到我们的dataset中，以下是dataset加载过程。加载代码如下，首先我们先要从数据中读取json文件，然后将读取中的json数据进行处理。json数据格式如下：
{
"prompt": "今年我想在室内种植水果，你能给我推荐一个容易入手的水果吗？",
"chosen": "没问题，你对“indoors”是什么理解？",
"rejected": "你在考虑哪种水果呢？"
},
question_key = "prompt"
good_key = 'chosen'
bad_key = 'rejected'
def read_jsonl(path):
with open(path, 'r', encoding='utf-8') as f:
lines = json.load(f)
for line in tqdm(lines):
yield line
def preprocess_function(examples):
new_examples = {
"input_ids_j": [],
"attention_mask_j": [],
"input_ids_k": [],
"attention_mask_k": [],
}
for question, response_j, response_k in zip(examples[question_key], examples[good_key], examples[bad_key]):
# 这里是添加了"Question: "和"\n\nAnswer: "作为模板，可以根据自己的模型进行替换。要跟SFT阶段对应
# tokenized_j = tokenizer("Question: " + question + "\n\nAnswer: " + response_j, truncation=True)
# tokenized_k = tokenizer("Question: " + question + "\n\nAnswer: " + response_k, truncation=True)
# 中文数据集：
tokenized_j = tokenizer("问：" + question + "\n\n答：" + response_j, truncation=True)
tokenized_k = tokenizer("问：" + question + "\n\n答：" + response_k, truncation=True)
new_examples["input_ids_j"].append(tokenized_j["input_ids"])
new_examples["attention_mask_j"].append(tokenized_j["attention_mask"])
new_examples["input_ids_k"].append(tokenized_k["input_ids"])
new_examples["attention_mask_k"].append(tokenized_k["attention_mask"])
return new_examples
train_dataset = datasets.Dataset.from_generator(
lambda: read_jsonl("E:\\LLM-Tuning\\LLM-Tuning-master\\data\\train_json\\train.json")
)
train_dataset = train_dataset.map(
preprocess_function,
batched=True,
num_proc=1,
remove_columns=original_columns,
)从dataset代码中可以看出，我们首先要将接受的句子编辑成对应的 input_ids token 以及他们的attention mask。编辑成两个对应的 input 和 attention。为什么这里要处理成两个呢？待会儿在说。这里我们先大致记住这里就是为了整理数据，将数据整理好以后后面会用到。
模型加载和lora模型加载
首先lora的配置，我这里只是用了 R 是4，同时我们的训练得到类型需要选择成分类模型，因为评分总归是对整个句子进行打分，那句子分类肯定是比较合适的，因为他会对整个句子进行打分。同时我这里只选择对Q K V 进行微调。因为我个人的电脑比较查只有16G内存，12G显存，因此如果不进行4 bit 量化训练的话，可能会受不住，因此这里我用了4bit 进行训练。然后我们这里为什么num_labels = 1呢？ 是因为我们只有接受和拒绝两个，因此这就相当于二分类，那只需要一个label 就行了，如果后续想要多个可以设置成多个即可。前面作者已经提到过了，本次作者采用的是Qwen2模型，因此这里我用的是Qwen2ForSequenceClassification来加载预训练模型。最后用大家都熟知的方法，get_peft_model(model, perft_config) 加载 lora 模型。
peft_config = LoraConfig(
task_type=TaskType.SEQ_CLS,
inference_mode=False,
r=4,
lora_alpha=32,
lora_dropout=0.1,
target_modules=["q_proj", "k_proj", "v_proj"]
)
bnb_config = BitsAndBytesConfig(
load_in_4bit = True,
bnb_4bit_quant_type='nf4',
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_use_double_quant = False,
)
model = Qwen2ForSequenceClassification.from_pretrained(
model_name, num_labels=1, torch_dtype=torch.bfloat16, trust_remote_code=True,
quantization_config=bnb_config,
device_map="auto"
)
print(model.hf_device_map)
model = get_peft_model(model, peft_config)
model.print_trainable_parameters()数据Collattor代码编写
查看源代码可以知道，其实最后的 dataset 都是需要经过Collattor进行转换成模型可以训练的数据。看代码就可以知道，就是就是对dataset按照 max length 进行句子补齐到相等长度。
# We need to define a special data collator that batches the data in our j vs k format.
# 感觉这里主要是为了做padding，因为transformers默认的data collator可能不支持对这种格式、字段输入
@dataclass
class RewardDataCollatorWithPadding:
tokenizer: PreTrainedTokenizerBase
padding: Union[bool, str, PaddingStrategy] = True
max_length: Optional[int] = None
pad_to_multiple_of: Optional[int] = None
return_tensors: str = "pt"
def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
features_j = []
features_k = []
for feature in features:
features_j.append(
{
"input_ids": feature["input_ids_j"],
"attention_mask": feature["attention_mask_j"],
}
)
features_k.append(
{
"input_ids": feature["input_ids_k"],
"attention_mask": feature["attention_mask_k"],
}
)
batch_j = self.tokenizer.pad(
features_j,
padding=self.padding,
max_length=self.max_length,
pad_to_multiple_of=self.pad_to_multiple_of,
return_tensors=self.return_tensors,
)
batch_k = self.tokenizer.pad(
features_k,
padding=self.padding,
max_length=self.max_length,
pad_to_multiple_of=self.pad_to_multiple_of,
return_tensors=self.return_tensors,
)
batch = {
"input_ids_j": batch_j["input_ids"],
"attention_mask_j": batch_j["attention_mask"],
"input_ids_k": batch_k["input_ids"],
"attention_mask_k": batch_k["attention_mask"],
"return_loss": True,
}
return batchTrain 模型也即Loss模型构建
这里是整个模型中最关键的部分，我们是希望模型尽可能的接受 正确的句子，拒绝错误的句子。为什么这么做比较合理呢？我们仔细来思考一个问题，就是我们在做二分类的时候， 如果正确的句子我们给他 0.51分，错误的句子我们给他 0.49分，模型也能正确的将模型分类出来，但是我们要做的是奖励模型，我们当然是希望正确的句子我们给他 0.99分，错误的句子我们给他0.01分。所以这里作者也是采用高人的方法，就是 logsigmoid(正确分数 - 错误分数), 那错误的分差辣的越开越好。因为分的越开，loss也就越低。
class RewardTrainer(Trainer):
def compute_loss(self, model, inputs, return_outputs=False):
rewards_j = model(input_ids=inputs["input_ids_j"], attention_mask=inputs["attention_mask_j"])[0]
rewards_k = model(input_ids=inputs["input_ids_k"], attention_mask=inputs["attention_mask_k"])[0]
loss = -nn.functional.logsigmoid(rewards_j - rewards_k).mean()
if return_outputs:
return loss, {"rewards_j": rewards_j, "rewards_k": rewards_k}
return loss
def save_model(self, output_dir=None, _internal_call=False):
self.model.save_pretrained(output_dir)模型训练
至此模型就可以训练啦，解下来，大家只需要开开心心的训练即可。训练的代码如下：
trainer = RewardTrainer(
model=model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
tokenizer=tokenizer,
#compute_metrics=compute_metrics,
data_collator=RewardDataCollatorWithPadding(tokenizer=tokenizer, max_length=2040),
)
trainer.train()模型训练中的小插曲
大家可能到此都挺开心的训练了，但是作者在这里训练的时候，可是自我否定和自我绝望中挣扎了两天，甚至都快放弃了。那到底发生什么事情呢？我们先来看以下 trainning_args
training_args = TrainingArguments(
output_dir = "E:\\lora\\output91",
num_train_epochs=10,
per_device_train_batch_size= 1,
gradient_accumulation_steps= 1,
optim= "paged_adamw_8bit",
save_steps=100,
logging_steps=30,
learning_rate=2e-4,
weight_decay = 0.001,
fp16 = False,
bf16 = False,
max_grad_norm=0.3,
max_steps = -1,
warmup_ratio = 0.3,
group_by_length = True,
lr_scheduler_type="linear",
report_to="wandb",
remove_unused_columns=False,
)真是这段笔者一点都不熟悉的，从自回归模型哪里拷贝过来的一个训练参数把笔者坑的那叫一个惨啊。这里面主要有一个参数 group_by_length = True。 就这么一个参数把我坑的惨惨喜喜的。发生什么事情了呢？在聊这个问题之前，我们先看一下发生什么问题了。在上面的代码中，我们可以看出来我们是对 data set 数据进行了处理，处理出来的标签是 input_ids_j 等了后缀的。当笔者在运行这个代码的时候，总是给我抱一个错误，就是模型需要的输入参数是 input_ids。错误代码如下：
ValueError: Can only automatically infer lengths for datasets whose items are dictionaries with an 'input_ids' key.
哎呀既然模型需要的是input_ids，那笔者就很疑惑了，既然这样，那我还训练个屁啊。于是我就跟着源码读啊读，走到这个源代码这里。既然是这里报错，也就意味着model_input_name 里面没有input_ids就会报错。但是我的dataset 就是没有啊，怎么办呢？要不我让model_input_name 有信息把，然后我们就把model_input_name 设置为自己dataset里的数据，input_ids_j。但是我很快又报错了，报后续的 max tensor不对。然后怎么弄都不对，要不我直接让lengths不为none，好一个家伙，我让他不为none以后，我的data set 就只变成了 只有input ids，然后在计算rewards的loss的时候又报错没有这个数据。整个就崩溃了。
class LengthGroupedSampler(Sampler):
r"""
Sampler that samples indices in a way that groups together features of the dataset of roughly the same length while
keeping a bit of randomness.
"""
def __init__(
self,
batch_size: int,
dataset: Optional[Dataset] = None,
lengths: Optional[List[int]] = None,
model_input_name: Optional[str] = None,
generator=None,
):
if dataset is None and lengths is None:
raise ValueError("One of dataset and lengths must be provided.")
self.batch_size = batch_size
if lengths is None:
model_input_name = model_input_name if model_input_name is not None else "input_ids"
if (
not (isinstance(dataset[0], dict) or isinstance(dataset[0], BatchEncoding))
or model_input_name not in dataset[0]
):
raise ValueError(
"Can only automatically infer lengths for datasets whose items are dictionaries with an "
f"'{model_input_name}' key."
)
lengths = [len(feature[model_input_name]) for feature in dataset]
elif isinstance(lengths, torch.Tensor):
logger.info(
"If lengths is a torch.Tensor, LengthGroupedSampler will be slow. Converting lengths to List[int]..."
)
lengths = lengths.tolist()
self.lengths = lengths
self.generator = generator
def __len__(self):
return len(self.lengths)
def __iter__(self):
indices = get_length_grouped_indices(self.lengths, self.batch_size, generator=self.generator)
return iter(indices) model_input_name = self.tokenizer.model_input_names[0] if self.tokenizer is not None else None
后来，一直读了很久才发现，看下面的代码路径，我发现我的设置是True，那他就一定会走到这里面来，那我不让他走不就好了吗？然后我就他把设置成了False。欸代码就跑通了
def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
if self.train_dataset is None or not has_length(self.train_dataset):
return None
# Build the sampler.
if self.args.group_by_length: # 我不让他走这里不就好了吗？
if is_datasets_available() and isinstance(self.train_dataset, datasets.Dataset):
lengths = (
self.train_dataset[self.args.length_column_name]
if self.args.length_column_name in self.train_dataset.column_names
else None
)
else:
lengths = None
model_input_name = self.tokenizer.model_input_names[0] if self.tokenizer is not None else None
return LengthGroupedSampler(
self.args.train_batch_size * self.args.gradient_accumulation_steps,
dataset=self.train_dataset,
lengths=lengths,
model_input_name=model_input_name,
)
else:
return RandomSampler(self.train_dataset)
后来我查了一些资料才发现，group_by_length 这个参数是为了训练加速用的，他是将一些 length 差不多相等的数据放在一起训练，这样就可以加速模型的训练。但是我们实际上在训练rewards的时候一个数据中 包含多个句子，每个句子长度本来就不一样，所以我们可以关闭这部分的优化。
扫一扫
2378
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
