采用lora 微调一个强化学习模型的代码

        我们之前已经讲解过强化学习PPO算法的原理,以及我们也将讲解过我们怎么用lora微调一个reward模型。既然我们现在知道了强化学习PPO算法的原理,又已经又了一个reward 模型,那此时此刻我们是不是只要在微调一个lora的自回归模型,是不是就可以开始我们强化学习微调的过程了。作者虽然没有真正利用这一整套流程优化出什么特别牛逼的模型,但是我觉得这是一个非常好的方向,就是我们首先先微调或者训练一个自回归模型,然后在训练一个reward评分模型,最后在用强化学习使得模型能够更好的输出更符合人语言或者更符合实际工程推理的模型。虽然我不知道大模型GPT 1O的模型是怎么训练出来的,但是影影约约感觉他们的背后逻辑肯定也涉及这部分的内容。刚好也可以走一遍这部分的代码,这样也就可以方便自己以后有机会尝试优化更好模型的时候一个比较好的知识库。

        好了,我们已经啰嗦了这么多,我们开始我们的正题了。正如我们所说的,我们已经有了自回归模型,我姑且叫他 model1,  然后又有了 reward model,我们叫他 r_model,然后又有理论支撑了。那我们就开始我们的代码训练了。

 包引入

        老规矩我们还是先引入所有的相关性的包

from dataclasses import dataclass, field
from typing import Optional

import torch
from accelerate import Accelerator
from datasets import load_dataset,load_from_disk
from peft import LoraConfig,PeftModel, PeftConfig
from tqdm import tqdm
from transformers import Adafactor, AutoTokenizer, HfArgumentParser, pipeline, AutoModelForCausalLM

from trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer, set_seed, PreTrainedModelWrapper
from trl.core import LengthSampler
from modeling_baichuan_for_cls import BaichuanForSequenceClassification
from peft import PeftModel

token加载

        然后我们就开始又开始了token加载过程,这个过程我觉得是每次都要操作的,这里的token我们都是采用基础模型的token,因为我之前一直用的是Qwen2模型,所以这里还是用Qwen2的基础模型的token。

tokenizer = AutoTokenizer.from_pretrained(script_args.base_model_name, trust_remote_code=True)

数据加载

        接下来就是数据的加载了,我们本次加载的数据还是之前reward相关的数据,其实这里用其他的数据可能更好一些,不然这里可能对reward的模型起不到泛化的验证。但是这里我们不重要,因为这里我只是走一遍流程,知道他怎么流转的就可以了,毕竟要想获得更好的效果,肯定是需要很多其他的技术,数据处理的技术,数据构建的技术,loss的技术等等,但是这些都是开始以后慢慢的过程,也就是1 - N的过程。

def preprocess_function(examples):
    new_examples = {
        "query": [],
        "input_ids": [],
    }
    for question in examples["prompt"]:
        query = "问:" + question + "\n\n答:"
        tokenized_question = tokenizer(query, truncation=True)
        new_examples["query"].append(query)
        new_examples["input_ids"].append(tokenized_question["input_ids"])
    return new_examples

dataset = dataset.map(
    preprocess_function,
    batched=True,
    num_proc=num_proc,
    remove_columns=original_columns,
)

        看数据处理的过程,我们可以看出来,本次数据集我们只是用了prompt,也就是提示词或者说是问题。因为此时我们在做强化学习的时候已经不知道答案是什么了,或者换句话说,答案是由自回归模型Model_1来回答的。

自回归模型和价值模型加载

        接下来我们就需要将之前lora微调过后的自回归模型,加载到代码中。

# 先加载基础模型,也即是Qwen2 - 0.5b模型
base_model_for_PPO = AutoModelForCausalLM.from_pretrained(
    script_args.base_model_name,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16, 
    device_map='auto'
    )
# 将Lora模型和基础模型进行合并
base_model_for_PPO_with_sft_lora = PeftModel.from_pretrained(
    base_model_for_PPO, 
    script_args.sft_model_lora_path
    )

# 然后将合并后的模型参数加载到多头价值 回归模型
ppo_model = AutoModelForCausalLMWithValueHead.from_pretrained(
    base_model_for_PPO_with_sft_lora
)

        我们之前已经通过原理知道,我们实际上采用的是价值回报模型,那价值怎么来呢?我们可以采用的是AutoModelForCausalLMWithValueHead 多头value 自回归模型来加载之前的lora 与训练模型。他其实就是像下面的代码一样,就是在原来的输出token(dim)维度中将他的dim重新计算出一个value值出来

        self.v_head = ValueHead(self.pretrained_model.config, **v_head_kwargs)
self.summary = nn.Linear(hidden_size, 1)  #其中hidden_size就是dim的维度

reward模型加载

        reward的模型加载就和普通的lora模型加载差不多

base_model_for_RM = Qwen2ForSequenceClassification.from_pretrained(
    script_args.base_model_name, num_labels=1, 
    trust_remote_code=True, 
    torch_dtype=torch.bfloat16, 
    device_map="auto",
    # device_map={"": current_device},
)
reward_model = PeftModel.from_pretrained(base_model_for_RM, script_args.reward_model_lora_path)

开始强化学习的代码步骤

        ok,现在数据有了,价值回报模型也有了,reward评分模型也有了。那我们就可以开始我们的强化学习啦。代码就如下面所示,其他的参数不太重要的,我就直接省略了。

ppo_trainer = PPOTrainer(
    config,
    ppo_model, # model with value head
    ref_model=ref_model,
    tokenizer=tokenizer,
    dataset=dataset,
    data_collator=collator,
    optimizer=optimizer,
)
for epoch, batch in tqdm(enumerate(ppo_trainer.dataloader)):
    if epoch >= config.total_ppo_epochs:
        break

    question_tensors = batch["input_ids"]
    
    try:
        response_tensors = ppo_trainer.generate(
            question_tensors,
            return_prompt=False,
            **generation_kwargs,
        )
        batch["response"] = tokenizer.batch_decode(response_tensors, skip_special_tokens=True)

        texts = [q + r for q, r in zip(batch["query"], batch["response"])]
        scores = get_reward_value(texts)
        rewards = [torch.tensor(score - script_args.reward_baseline) for score in scores]
        for q, r, s in zip(batch["query"], batch["response"], scores):
            print(epoch,'query:',q)
            print('response:',r)
            print('score:',s)
        
        stats = ppo_trainer.step(question_tensors, response_tensors, rewards)
        ppo_trainer.log_stats(stats, batch, rewards)

        if script_args.save_freq and epoch and epoch % script_args.save_freq == 0:
            ppo_trainer.save_pretrained(script_args.output_dir + f"step_{epoch}")
            
    except Exception as e:
        print('---------------------')
        print(e)
        print(epoch)
        print(question_tensors)
        print('---------------------')
        break

        那我们先来看第一部分的代码,PPOTrain 训练,需要一个PPO模型也就是我们的价值模型,同样也需要一个参考模型。这里参考模型是啥呢?有什么用呢?还记得我们之前讲的强化学习的原理中重要性采样吗?这里的ref_model就是传说中的帮忙代理产生数据的模型。所以他一开始就是和PPO模型是一模一样的。如下面代码中的ppo_trainer.generate的核心代码。

    if generate_ref_response:
            ref_model = self.model if self.is_peft_model else self.ref_model
        if isinstance(query_tensor, List):
            response = self._generate_batched(
                self.model,
                query_tensor,
                length_sampler=length_sampler,
                batch_size=batch_size,
                return_prompt=return_prompt,
                **generation_kwargs,
            )
            if generate_ref_response:
                ref_response = self._generate_batched(
                    ref_model,
                    query_tensor,
                    length_sampler=length_sampler,
                    batch_size=batch_size,
                    return_prompt=return_prompt,
                    **generation_kwargs,
                )

        接下来我们就开始看第二部分的代码,先用ppo_train模型将问题batch 产生一个 response,然后再将对应的response 进行通过评分模型进行评分拿到scores,scores的shape (batch,1)

        接下来就是ppo_trainer.step步骤了,这个就是核心代码了,这个在之前的强化学习的核心原理中已经讲解过了,这里其实就是一次PPO 训练的过程,也就是产生了一次数据以后进行不同的迭代的过程。那一次step以后,ppo_model 已经变了,但是ref_model还是上次的,这个时候我们训练完以后将ref_model 替换为训练过后的ppo_model,也就是 ppo_trainer.log_stats。这个是不是就是我们之前讲解过的,什么是最好的代理模型,不就是上一次我们大迭代的模型参数吗?因为他跟我们下一次迭代的模型最接近。至此我们的强化学习就已经完成了。是不是很神奇。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值