强化学习系列（十）--基于大语言模型的RLHF

# 处理数据：tokenized_j好的回复，tokenized_k差的回复
for question, response_j, response_k in zip(examples[question_key], examples[good_key], examples[bad_key]):
        tokenized_j = tokenizer("问：" + question + "\n\n答：" + response_j, truncation=True)
        tokenized_k = tokenizer("问：" + question + "\n\n答：" + response_k, truncation=True)
...
# Define how to compute the reward loss. We use the InstructGPT pairwise logloss: https://arxiv.org/abs/2203.02155
    # 由于训练数据都是好坏成对出现，所以这里两个输入两个输出
    def compute_loss(self, model, inputs, return_outputs=False):
        rewards_j = model(input_ids=inputs["input_ids_j"], attention_mask=inputs["attention_mask_j"])[0]
        rewards_k = model(input_ids=inputs["input_ids_k"], attention_mask=inputs["attention_mask_k"])[0]
        loss = -nn.functional.logsigmoid(rewards_j - rewards_k).mean()
        if return_outputs:
            return loss, {"rewards_j": rewards_j, "rewards_k": rewards_k}
        return loss

...
lm_logits = base_model_output.logits
loss = base_model_output.loss
...
# 增加一层v_head
value = self.v_head(last_hidden_state).squeeze(-1)
...
return (lm_logits, loss, value)

ref_model = AutoModelForCausalLMWithValueHead.from_pretrained(
    script_args.merged_sft_model_path,
    trust_remote_code=True
)

base_model_for_RM = BaichuanForSequenceClassification.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)

# load the base model
base_model_for_PPO = AutoModelForCausalLM.from_pretrained(
    script_args.base_model_name,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16, 
    device_map='auto'
    )
# install the lora modules
base_model_for_PPO_with_sft_lora = PeftModel.from_pretrained(
    base_model_for_PPO, 
    script_args.sft_model_lora_path
    )
# wrap with the AutoModelForCausalLMWithValueHead wrapper
ppo_model = AutoModelForCausalLMWithValueHead.from_pretrained(
    base_model_for_PPO_with_sft_lora
)

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

        # Compute sentiment score
        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)
        # Run PPO step
        stats = ppo_trainer.step(question_tensors, response_tensors, rewards)

 # actor网络预测responses和动作logits
 query_responses, logitss = batch_generation(
                        unwrapped_model.policy,
                        queries,
                        args.local_rollout_forward_batch_size,
                        processing_class.pad_token_id,
                        generation_config,
                    )

all_logprob = F.log_softmax(logits, dim=-1)
# logporb:actor网络预测response的概率分布
logprob = torch.gather(all_logprob, 2, response.unsqueeze(-1)).squeeze(-1)
...
ref_output = forward(ref_policy, query_response, processing_class.pad_token_id)
...
ref_all_logprob = F.log_softmax(ref_logits, dim=-1)
# ref_logprob:ref网络预测response的概率分布
ref_logprob = torch.gather(ref_all_logprob, 2, response.unsqueeze(-1)).squeeze(-1)

# critic网络预测的状态值value（累计rewards）
full_value, _, _ = get_reward(unwrapped_value_model, query_response, processing_class.pad_token_id, context_length)
value = full_value[:, context_length - 1 : -1].squeeze(-1)

# reward网络计算的score(T时刻即刻回报reward)
_, score, _ = get_reward(reward_model, postprocessed_query_response, processing_class.pad_token_id, context_length )

# 决策模型的回复
responses = torch.cat(responses, 0)
# 处理response,哪些参与loss计算
postprocessed_responses = torch.cat(postprocessed_responses, 0)
# actor模型预测的动作概率分布
logprobs = torch.cat(logprobs, 0)
# ref模型预测的动作概率分布
ref_logprobs = torch.cat(ref_logprobs, 0)
# 序列长度
sequence_lengths = torch.cat(sequence_lengths, 0)
# reward模型计算的即刻reward
scores = torch.cat(scores, 0)
# critic模型预测的累计V
values = torch.cat(values, 0)

# 4. compute rewards
# actor和 ref输出分布的KL散度 * 非负权重
kl = logprobs - ref_logprobs
non_score_reward = -args.kl_coef * kl
# 每个时刻的reward使用kl散度表示。
# 这里解释下，由于rewardmodel训练是基于input+response一起的分类模型，所以分类结果最为最后一个时刻的奖励R，其他时刻是没有R。这里是使用KL散度，让决策和ref分布靠近来作为即可奖励。
rewards = non_score_reward.clone()
actual_start = torch.arange(rewards.size(0), device=rewards.device)
actual_end = torch.where(sequence_lengths_p1 < rewards.size(1), sequence_lengths_p1, sequence_lengths)
# 最后时刻的R就是reward model的score
rewards[[actual_start, actual_end]] += scores

....

# 6. compute advantages and returns
lastgaelam = 0
advantages_reversed = []
gen_length = responses.shape[1]
# 这里是使用了泛优势函数GAE，从最后时刻逆序计算
for t in reversed(range(gen_length)):
        nextvalues = values[:, t + 1] if t < gen_length - 1 else 0.0
        # TD：t时刻即可回报+gamma*下时刻的累计回报- t时刻累计回报
        delta = rewards[:, t] + args.gamma * nextvalues - values[:, t]
        # GAE的更新公式
        lastgaelam = delta + args.gamma * args.lam * lastgaelam
        advantages_reversed.append(lastgaelam)
advantages = torch.stack(advantages_reversed[::-1], axis=1)
# 优势+状态值估计
returns = advantages + values
advantages = masked_whiten(advantages, ~padding_mask)
advantages = torch.masked_fill(advantages, padding_mask, 0)

logprobs_diff = new_logprobs - mb_logprobs
ratio = torch.exp(logprobs_diff)
pg_losses = -mb_advantage * ratio
pg_losses2 = -mb_advantage * torch.clamp(ratio, 1.0 - args.cliprange, 1.0 + args.cliprange)
pg_loss_max = torch.max(pg_losses, pg_losses2)
pg_loss = masked_mean(pg_loss_max, ~padding_mask[micro_batch_inds])
loss = pg_loss + args.vf_coef * vf_loss