[chat] fix RM & MDP

[cwher]

Try to merge #3645 and evaluate its performance.

[cwher]

The PR mainly contains two major changes.

• Change of Reward Model and Critic

  TLDR: only last_hidden_states[:, -1] is used by Critic / Reward Model.

  1. Input of Critic Model should end with <eos> token

  2. forward fn of Critic is changed to the following form:

     def forward(...):
           outputs = self.model(sequences, attention_mask=attention_mask)
           last_hidden_states = outputs['last_hidden_state']
           sentence_hidden_states = last_hidden_states[:, -1]
           values = self.value_head(sentence_hidden_states).squeeze(1)

• Change of MDP definition and reward

  TLDR: In previous version, coati treats the whole response as a step (or transition) of MDP. In this version, user can munually determine to treat every chunk_size tokens as a step of MDP.

  In sum, if chunk_size=1, the behavior of coati is close to trlx; and if chunk_size=response_size, the behavior of coati is close to previous version.

  NOTE: trlx DOES NOT add <eos> token to the input of Critic; and previous coati version may contain some bug in the RL training process.

  1. New MDP definition

     chunk_size parameter is added to group every chunk_size tokens as a step of MDP.

     e.g., if chunk_size = 2,

     s0 --> a0 ---> s1 ---> a1 ---> ...
     where s0 = |prompt|, a0 = |t0t1|, s1 = |prompt|t0t1|, a1 = |t2t3|, ...

     $$ \pi (a_0\mid s_0)=\text{Pr}(t_0 \mid s_0) \times \text{Pr}(t_1 \mid s_0, t_0) $$

  2. Modified reward

  $$ \begin{align} \text{reward}[i] &= -\beta \times \log \left(\pi^{RL}(a\mid s_0) / \pi^{SFT}(a\mid s_0)\right)\\ r &= \text{Reward Model}(\mid \text{prompt} \mid \text{response} \mid \text{}) \\ \text{reward}[terminal] &= \text{reward}[t] + r \end{align} $$

  3. Then, advantages and returns are calculated using GAE.