Closed puyuanOT closed 1 year ago
@puyuanOT Thanks for your question. You are right that we keep the whole sequence in the code on the generation. And you are also right that in the step2 reward training, we use the response only. Now if you look at how we calculate the loss, you will see that here we only apply the response.
actor_loss = self.actor_loss_fn(actor_log_prob[:, start:],
log_probs[:, start:], advantages,
action_mask[:, start:])
where the start is the length of the prompt-size
For critic models, similarly it's here: https://github.com/microsoft/DeepSpeedExamples/blob/master/applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/ppo_trainer.py#L177
Let me know if you have further questions :)
In the
_generate_sequence
of ppo_trainer.py, we haveIf I understand correctly, we only want to keep the model response and this is done by truncating each sequence in
seq
variable, which has the shapebs * seq_len
. However, we haveout_seq.append(seq[i:i + 1])
in the last for loop, which keeps the whole sequence instead.However, if the intention is instead keeping both prompts and response into the out_seq, we would have a contradiction with step_2_training, where only responses are passed to the reward model.