【大模型学习 | SFT & PPO原理&代码实现】

from transformers import AutoModelForSequenceClassification
rm_model = AutoModelForSequenceClassification.from_pretrained("your_model", num_labels=1)

from transformers import AutoModel, AutoConfig
import torch.nn as nn

class RewardModel(nn.Module):
    def __init__(self, base_model_name):
        super().__init__()
        self.backbone = AutoModel.from_pretrained(base_model_name)
        self.reward_head = nn.Sequential(
            nn.Linear(self.backbone.config.hidden_size, 1)
        )
    
    def forward(self, input_ids, attention_mask):
        outputs = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
        last_hidden = outputs.last_hidden_state  # shape: (batch, seq_len, hidden)
        cls_repr = last_hidden[:, 0, :]          # 取第一个token的表征
        reward = self.reward_head(cls_repr)
        return reward.squeeze(-1)                # 输出 shape: (batch,)

from trl import PPOTrainer, PPOConfig
def run_ppo():
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=False)
    model = AutoModelForCausalLM.from_pretrained("sft_output")
    rm_model = AutoModelForSequenceClassification.from_pretrained("rm_output").to(DEVICE)

    config = PPOConfig(
        model_name="sft_output",
        learning_rate=1e-5,
        batch_size=4,
        mini_batch_size=2,
        gradient_accumulation_steps=1,
        log_with=None
    )
    ppo_trainer = PPOTrainer(
        config=config,
        model=model,
        tokenizer=tokenizer,
        reward_model=rm_model
    )

    prompts = ["Explain the concept of overfitting."]
    for epoch in range(5):
        batch = random.sample(prompts, k=3)
        query_tensors = tokenizer(batch, return_tensors="pt", padding=True, truncation=True).input_ids.to(DEVICE)
        response_tensors = ppo_trainer.model.generate(query_tensors, max_new_tokens=100)

        texts = [tokenizer.decode(r, skip_special_tokens=True) for r in response_tensors]
        rewards = []
        for q, r in zip(batch, texts):
            input_text = f"{q}\n{r}"
            inputs = tokenizer(input_text, return_tensors="pt", truncation=True, padding=True).to(DEVICE)
            reward = rm_model(**inputs)["logits"].squeeze().item()
            rewards.append(reward)

        print(f"[PPO] E{epoch}: Rewards = {rewards}")
        ppo_trainer.step(query_tensors, response_tensors, torch.tensor(rewards))