TRL库PPOTrainer自定义奖励函数报错：Tensor无hidden_states属性求助

错误原因 #

你遇到的AttributeError: 'Tensor' object has no attribute 'hidden_states'是因为TRL库的get_reward函数默认期望传入的reward_model能接收模型输出的**隐藏状态（hidden_states）**作为输入，但你的自定义KeywordRewardModel是基于文本关键词计算奖励，且PPOTrainer内部生成响应时的参数覆盖导致未返回包含hidden_states的字典，最终触发错误。同时，直接将奖励模型复用为值函数模型不符合PPO设计逻辑，可能引发潜在问题。

修复方案 #

放弃使用PPOTrainer的reward_model和value_model参数，改为手动控制PPO训练循环：生成响应后直接基于文本计算奖励，再调用ppo_trainer.step()完成训练步骤，这种方式更灵活，完全适配自定义奖励逻辑。

修改后的完整代码 #

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
— trl==0.17.0 + unsloth/mistral-7b-bnb-4bit
(“ppo_rlhf_ckpt/sft” have model)
"""

import os, random, torch

# ─── environment variable ─────────────────────────────────────────────────────
os.environ["TRANSFORMERS_CACHE"]        = "/storage1/hf_cache/transformers"
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
os.environ["PYTORCH_CUDA_ALLOC_CONF"]   = "expandable_segments:True"
SEED = 42
torch.manual_seed(SEED)
random.seed(SEED)

# ─── library ─────────────────────────────────────────────────────────────
from datasets import load_dataset
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    DataCollatorWithPadding,
)
from peft import LoraConfig
from trl import PPOTrainer, PPOConfig

# ───────────────────────────────────────────────────────────────────
SFT_CKPT   = "ppo_rlhf_ckpt/sft"
OUT_DIR    = "ppo_rlhf_ckpt/ppo"
MODEL_NAME = "unsloth/mistral-7b-bnb-4bit"

# ─── 3) tokenizer ───────────────────────────────────────────────────────
tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME, use_fast=True, trust_remote_code=True
)
if tokenizer.eos_token is None:
    tokenizer.add_special_tokens({"eos_token":"</s>"})
if tokenizer.pad_token is None:
    tokenizer.pad_token    = tokenizer.eos_token
    tokenizer.pad_token_id = tokenizer.eos_token_id

# ─── 4) policy & ref model ───────────────────────────────────────────────────────
policy = AutoModelForCausalLM.from_pretrained(
    SFT_CKPT,
    load_in_4bit    = True,
    device_map      = "auto",
    trust_remote_code=True,
)
# 重新加载ref_model，避免deepcopy 4bit模型可能的问题
ref_model = AutoModelForCausalLM.from_pretrained(
    SFT_CKPT,
    load_in_4bit    = True,
    device_map      = "auto",
    trust_remote_code=True,
).eval()
for p in ref_model.parameters():
    p.requires_grad = False

# ─── 6) PPO dataset ───────────────────────────────────────────────────────
raw = load_dataset("Dahoas/rm-static", split="train[:300]").shuffle(SEED)
ppo_raw = raw.rename_column("prompt","query").remove_columns(["chosen","rejected"])
def tokenize_fn(batch):
    return tokenizer(batch["query"], truncation=True, padding="longest")
ppo_enc = ppo_raw.map(tokenize_fn, batched=True, remove_columns=["query"])
ppo_enc.set_format(type="torch", columns=["input_ids","attention_mask"])

# ─── 8) LoRA config ───────────────────────────────────────────────────────────
lora_cfg = LoraConfig(
    r=8, lora_alpha=16, bias="none",
    target_modules=["q_proj","k_proj","v_proj","o_proj"],
)

# ─── 9) PPOTrainer config & training ───────────────────────────────────────────────
ppo_trainer = PPOTrainer(
    args            = PPOConfig(
        output_dir      = OUT_DIR,
        learning_rate   = 1e-5,
        num_ppo_epochs  = 1,
        batch_size      = 4,
        mini_batch_size = 2,
        kl_coef         = 0.05,
    ),
    processing_class= tokenizer,
    model           = policy,
    ref_model       = ref_model,
    train_dataset   = ppo_enc,
    data_collator   = DataCollatorWithPadding(tokenizer),
    peft_config     = lora_cfg,
)

print("=== [PPO] Training starts ===")
# 生成配置
generate_kwargs = {
    "max_new_tokens": 100,
    "temperature": 0.7,
    "do_sample": True,
    "pad_token_id": tokenizer.pad_token_id,
    "eos_token_id": tokenizer.eos_token_id,
}

# 手动PPO训练循环
for epoch in range(ppo_trainer.args.num_ppo_epochs):
    for batch in ppo_trainer.dataloader:
        # 将batch数据移到设备上
        query_ids = batch["input_ids"].to(ppo_trainer.accelerator.device)
        attn_mask = batch["attention_mask"].to(ppo_trainer.accelerator.device)
        
        # 生成模型响应
        response_ids = ppo_trainer.generate(
            query_ids,
            attention_mask=attn_mask,
            **generate_kwargs
        )
        
        # 解码响应文本并计算自定义奖励
        responses = tokenizer.batch_decode(response_ids, skip_special_tokens=True)
        rewards = [
            torch.tensor(1.0 if "good" in resp else -1.0, device=ppo_trainer.accelerator.device)
            for resp in responses
        ]
        
        # 执行PPO训练步骤
        stats = ppo_trainer.step(query_ids, response_ids, rewards)
        
        # 记录训练统计信息
        ppo_trainer.log_stats(stats, batch, rewards)

ppo_trainer.save_model(OUT_DIR)

关键修改点 #

1. 移除reward_model和value_model参数：避免TRL默认奖励计算逻辑冲突
2. 替换自动训练为手动循环：完全控制响应生成、奖励计算和PPO步骤
3. 重新加载ref_model：解决4bit模型deepcopy可能导致的潜在问题
4. 明确生成参数：添加pad_token_id和eos_token_id确保生成过程稳定

内容的提问来源于stack exchange，提问作者남석현