【RLHF】想训练ChatGPT？先来看看强化学习（RL）+语言模型（LM）吧（附源码）

prompt: 刚收到货，感觉

output 1: 刚收到货，感觉 有 点 不 符 合 预 期 ，不 好
output 2: 刚收到货，感觉 挺 无 奈 的 送 货 速 度 不 太 行
...

output 1: 刚收到货，感觉有 点 不 符 合 预 期 ，不 好                -> 0.2 分
output 2: 刚收到货，感觉有 挺 无 奈 的 送 货 速 度 不 太 行          -> 0.1 分
output 3: 刚收到货，感觉有 些 惊 喜 于 货 物 质 量                  -> 0.9 分
...

# prompt池
prompts = [
    '刚收到货，感觉',
    '这部电影很',
    '说实话，真的很',
    '这次购物总的来说体验很'
]
...

for _ in range(config['batch_size']):
        random_prompt = random.choice(prompts)                                  # 随机选择一个prompt
        tokens = gpt2_tokenizer.encode(random_prompt)
        batch['tokens'].append(tokens)
        batch['query'].append(random_prompt)
query_tensors = [torch.tensor(t).long().to(device) for t in batch["tokens"]]
...

for i in range(config['batch_size']):
    gen_len = config['gen_len']
    response = gpt2_model.generate(query_tensors[i].unsqueeze(dim=0),           # 利用当前选择的prompt生成句子
                                   max_new_tokens=gen_len, **gen_kwargs)
    response_tensors.append(response.squeeze()[-gen_len:])

[
    '刚收到货，感觉 很 一 般',
    '这部电影很 俗 而 且 很 无 趣',
    '这次购物总的来说体验很 烂 不 是 我 想 要 的',
    ...
]

# 情绪识别模型初始化
senti_tokenizer = AutoTokenizer.from_pretrained('uer/roberta-base-finetuned-jd-binary-chinese')
senti_model = AutoModelForSequenceClassification.from_pretrained('uer/roberta-base-finetuned-jd-binary-chinese')
sentiment_pipe = pipeline('sentiment-analysis', model=senti_model, tokenizer=senti_tokenizer, device=pipe_device)
...


texts = [q + r for q,r in zip(batch['query'], batch['response'])]           # 将 prompt 和生成的 response 做拼接
pipe_outputs = sentiment_pipe(texts)                                        # 计算正向/负向情感得分

[
    0.4,
    0.3,
    0.3,
    ...
]

ppo_trainer.step(query_tensors, response_tensors, rewards)          # PPO Update

loss_p, loss_v, train_stats  = self.loss(logprobs, values, rewards, query, response, model_input)
loss = loss_p + loss_v
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
...

for t in reversed(range(gen_len)):
    nextvalues = values[:, t + 1] if t < gen_len - 1 else 0.0
    delta = rewards[:, t] + self.ppo_params['gamma'] * nextvalues - values[:, t]          # 优势函数：r + Vnext - V
    lastgaelam = delta + self.ppo_params['gamma'] * self.ppo_params['lam'] * lastgaelam   # GAE, 用于平衡 bias 和 variance
    advantages_reversed.append(lastgaelam)
    advantages = torch.stack(advantages_reversed[::-1]).transpose(0, 1)

logits, _, vpred = self.model(model_input)                                  # 跑一遍模型，得到句子中每个token被选择的概率
logprob = logprobs_from_logits(logits[:,:-1,:], model_input[:, 1:])         # 将概率取log对数
ratio = torch.exp(logprob - old_logprobs)                                   # log相减，等同于概率相除
pg_losses = -advantages * ratio

class GPT2HeadWithValueModel(GPT2PreTrainedModel):
    """The GPT2HeadWithValueModel class implements a GPT2 language model with a secondary, scalar head."""
    def __init__(self, config):
        super().__init__(config)
        config.num_labels = 1
        self.transformer = GPT2Model(config)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        self.v_head = ValueHead(config)                                       # 添加 Value Head
        self.init_weights()
    ...

class ValueHead(nn.Module):
    """The ValueHead class implements a head for GPT2 that returns a scalar for each output token."""
    
    def __init__(self, config):
        super().__init__()
        self.summary = nn.Linear(config.hidden_size, 1)                        # (hidden_size -> 1)
    ...

returns = advantages + values                      # r + v_next - v + v => r + v_next
logits, _, vpred = self.model(model_input)         # 跑一遍语言模型，得到每个 token 的 v_pred
vf_losses1 = (vpred - returns) ** 2                # MSE