change to right-padding, update reward score #803

Former-commit-id: baa90415bc8f5ebd423d001378b51c3a3a6c2ec7

src/llmtuner/tuner/rm/trainer.py

@@ -32,21 +32,50 @@ class PairwisePeftTrainer(PeftTrainer):
         r"""
         Computes pairwise loss. The first n examples are chosen and the last n examples are rejected.
 
-        We use score on the EOS token to represent reward of the whole sentence.
-
-        Subclass and override to inject custom behavior. It should not be directly used by external scripts.
-
-        Note that the first element will be removed from the output tuple.
+        Subclass and override to inject custom behavior.
 
+        Note that the first element will be removed from the output tuple. 
         See: https://github.com/huggingface/transformers/blob/v4.30.2/src/transformers/trainer.py#L3509
         """
-        batch_size = inputs["input_ids"].size(0) // 2
+        # Compute rewards
         _, _, values = model(**inputs, output_hidden_states=True, return_dict=True)
         if values.size(0) != inputs["input_ids"].size(0): # adapt to chatglm2
             values = torch.transpose(values, 0, 1)
-        r_accept, r_reject = values[:, -1].split(batch_size, dim=0)
-        loss = -torch.log(torch.sigmoid(r_accept - r_reject)).mean()
-        return (loss, [loss, r_accept, r_reject]) if return_outputs else loss
+
+        # Split the inputs and rewards into two parts, chosen and rejected
+        batch_size = inputs["input_ids"].size(0) // 2
+        chosen_input_ids, rejected_input_ids = inputs["input_ids"][:batch_size], inputs["input_ids"][batch_size:]
+        chosen_attn_mask, rejected_attn_mask = (
+            inputs["attention_mask"][:batch_size], inputs["attention_mask"][batch_size:]
+        )
+        chosen_rewards, rejected_rewards = values[:batch_size], values[batch_size:]
+        chosen_scores, rejected_scores = [], []
+
+        # Compute pairwise loss. Only backprop on the different tokens before padding
+        # Inspired by: https://github.com/CarperAI/trlx/blob/main/examples/summarize_rlhf/reward_model/reward_model.py
+        loss = 0
+        for i in range(batch_size):
+            chosen_length = chosen_attn_mask[i].nonzero()[-1] + 1
+            rejected_length = rejected_attn_mask[i].nonzero()[-1] + 1
+            check_divergence = (chosen_input_ids[i] != rejected_input_ids[i]).nonzero()
+
+            if len(check_divergence) == 0:
+                end_index = chosen_length
+                div_index = end_index - 1
+            else:
+                end_index = max(chosen_length, rejected_length)
+                div_index = check_divergence[0]
+
+            assert div_index > 0
+            chosen_trunc_rewards = chosen_rewards[i, div_index:end_index]
+            rejected_trunc_rewards = rejected_rewards[i, div_index:end_index]
+            chosen_scores.append(chosen_trunc_rewards[-1]) # use the end score for inference
+            rejected_scores.append(rejected_trunc_rewards[-1])
+            loss += -torch.nn.functional.logsigmoid(chosen_trunc_rewards - rejected_trunc_rewards).mean()
+
+        loss = loss / batch_size
+        chosen_scores, rejected_scores = torch.stack(chosen_scores), torch.stack(rejected_scores)
+        return (loss, [loss, chosen_scores, rejected_scores]) if return_outputs else loss
 
     def save_predictions(
         self,
@@ -63,10 +92,10 @@ class PairwisePeftTrainer(PeftTrainer):
         output_prediction_file = os.path.join(self.args.output_dir, "generated_predictions.jsonl")
         logger.info(f"Saving prediction results to {output_prediction_file}")
 
-        acc_scores, rej_scores = predict_results.predictions
+        chosen_scores, rejected_scores = predict_results.predictions
 
         with open(output_prediction_file, "w", encoding="utf-8") as writer:
             res: List[str] = []
-            for acc_score, rej_score in zip(acc_scores, rej_scores):
-                res.append(json.dumps({"accept": round(float(acc_score), 2), "reject": round(float(rej_score), 2)}))
+            for c_score, r_score in zip(chosen_scores, rejected_scores):
+                res.append(json.dumps({"chosen": round(float(c_score), 2), "rejected": round(float(r_score), 2)}))
             writer.write("\n".join(res))

src/llmtuner/tuner/rm/workflow.py

@@ -1,5 +1,4 @@
 # Inspired by:
-# https://github.com/lvwerra/trl/blob/main/examples/summarization/scripts/reward_summarization.py
 # https://github.com/CarperAI/trlx/blob/main/examples/summarize_rlhf/reward_model/train_reward_model_gptj.py
 
 from typing import TYPE_CHECKING, Optional, List