Trainer
Trainer用于预训练。
{
"input_ids": [],
"attention_mask": [],
"labels": [],
}
这仨就是tokenizer给你的东西,所以处理语料是主要内容,然后扔tokenizer里得到这仨。
一般input_ids就是你输入的整个QA语料,包括bos eos。
像Qwen1.5就是:
"<|im_start|>user\nquestion\n<|im_end|>\n<|im_start|>assistant\nanswer<|im_end|>"
具体格式得看具体模型。
attention_mask与你的语料长度一致,一般全1。不一般是指这种情况:你一次的语料是一个batch,这会需要你对短的数据做left padding或right padding(取决于你在训练还是推理),对于padding的token就不能attention了,对应位置需要填0。但是Trainer会要你设置batch,所以我也从来没试过自己batch数据。
labels要你把不训练的token全标为-100.
RewardTrainer
new_examples = {
"input_ids_chosen": [],
"attention_mask_chosen": [],
"input_ids_rejected": [],
"attention_mask_rejected": [],
}
reward model是有正负例的,chosen是正,rejected是负。system和prompt和query都要融到chosen和rejected里,所以这里只有chosen和rejected两组。attention必然是全部token。
对应的RewardTrainer继承自Trainer,InstructGPT对应的代码如下:
class RewardTrainer(Trainer):
"""
Trainer for reward models
Define how to compute the reward loss. Use the InstructGPT pairwise logloss: https://arxiv.org/abs/2203.02155
"""
def compute_loss(self, model, inputs, return_outputs=False):
rewards_chosen = model(input_ids=inputs["input_ids_chosen"],
attention_mask=inputs["attention_mask_chosen"])[0]
rewards_rejected = model(input_ids=inputs["input_ids_rejected"],
attention_mask=inputs["attention_mask_rejected"])[0]
loss = -torch.nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean()
if return_outputs:
return loss, {"rewards_chosen": rewards_chosen, "rewards_rejected": rewards_rejected}
return loss
def evaluate(
self,
eval_dataset: Optional[Dataset] = None,
ignore_keys: Optional[List[str]] = None,
metric_key_prefix: str = "eval",
) -> Dict[str, float]:
if eval_dataset is None:
eval_dataset = self.eval_dataset
return super().evaluate(eval_dataset=eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix)
def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None):
# Prepare inputs for chosen and rejected separately
device = model.device
inputs_chosen = {
"input_ids": inputs["input_ids_chosen"].to(device),
"attention_mask": inputs["attention_mask_chosen"].to(device),
}
outputs_chosen = model(**inputs_chosen)
rewards_chosen = outputs_chosen.logits.detach()
inputs_rejected = {
"input_ids": inputs["input_ids_rejected"].to(device),
"attention_mask": inputs["attention_mask_rejected"].to(device),
}
outputs_rejected = model(**inputs_rejected)
rewards_rejected = outputs_rejected.logits.detach()
# Keep the compute_loss method
loss = -torch.nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean()
if prediction_loss_only:
return (loss, None, None)
return (loss, rewards_chosen, rewards_rejected)
def save_model(self, output_dir=None, _internal_call=False):
"""Save the LoRA model."""
os.makedirs(output_dir, exist_ok=True)
torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))
self.model.save_pretrained(output_dir)
PPOTrainer
这是RLHF的最终阶段。
new_examples = {
"query": [],
"input_ids": [],
}
PPO只需要你提供query和query对应的input_ids就行。
DPOTrainer
{
"prompt": "",
"chosen": "",
"rejected": "",
}
prompt要把system、question融一起,包括输出头。总之它得拼上chosen和rejected之后是个合格的语料。
这仨都是文字,DPOTrainer会要求你提供tokenizer。