add hugchat reward training

Author: wjn1996

applications/instruction_prompting/HugChat/reward_model_training/run_reward_gpt2.sh

@@ -0,0 +1,41 @@
+path=/wjn/pre-trained-lm/gpt2
+
+model_name=gpt2
+
+data_path=/wjn/nlp_task_datasets/rlhf_preference # consists of preference_train.json, preference_dev.json, preference_test.json
+
+
+export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+python3 -m torch.distributed.launch --nproc_per_node=8 --master_port 6013 hugnlp_runner.py \
+--model_name_or_path=$path \
+--data_dir=$data_path \
+--max_seq_length=512 \
+--output_dir=./outputs/rlhf/$model_name/ \
+--do_train \
+--do_eval \
+--do_predict \
+--per_device_train_batch_size=8 \
+--per_device_eval_batch_size=1 \
+--evaluation_strategy=steps \
+--save_strategy=steps \
+--gradient_accumulation_steps=1 \
+--learning_rate=1e-05 \
+--logging_steps=10000000 \
+--eval_steps=3000 \
+--save_steps=3000 \
+--save_total_limit=10 \
+--num_train_epochs=3 \
+--report_to=none \
+--task_name=pairwise_reward \
+--task_type=rl_reward \
+--model_type=gpt2 \
+--exp_name=preference_reward \
+--warmup_steps=6000 \
+--load_best_model_at_end \
+--metric_for_best_model=acc \
+--ignore_data_skip \
+--remove_unused_columns=False \
+--cache_dir=/wjn/.cache \
+--overwrite_output_dir \
+# --deepspeed=./deepspeed/ds_config_fp16_z1.json \
+# --fp16

evaluators/__init__.py

@@ -8,6 +8,7 @@
 from evaluators.token_classification_evaluator import TokenClassificationEvaluator
 from evaluators.span_extraction_evaluator import SpanExtractionEvaluator
 from evaluators.multi_choice_evaluator import MultiChoiceEvaluator
+from evaluators.reinforcement_learning_evaluator import PairwiseRewardEvaluator
 
 # Models for pre-training
 PRETRAIN_EVALUATOR_CLASSES = {
@@ -57,6 +58,11 @@
     "code_generation": None,
 }
 
+REINFORCEMENT_MODEL_CLASSES = {
+    "causal_actor": None,
+    "auto_critic": None, 
+    "rl_reward": PairwiseRewardEvaluator,
+}
 
 # task_type 负责对应model类型
 OTHER_EVALUATOR_CLASSES = {
@@ -101,6 +107,7 @@
     SPAN_EXTRACTION_EVALUATOR_CLASSES,
     FEWSHOT_EVALUATOR_CLASSES,
     CODE_EVALUATOR_CLASSES,
+    REINFORCEMENT_MODEL_CLASSES,
     OTHER_EVALUATOR_CLASSES
 ]
 

evaluators/reinforcement_learning_evaluator.py

@@ -0,0 +1,117 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2023/5/6 8:09 p.m.
+# @Author  : JianingWang
+# @File    : reinforcement_learning_evaluator.py
+
+import json
+import os.path
+import math
+import torch
+import numpy as np
+from tqdm import tqdm
+from typing import Dict, Union, Any, Optional, Callable, List, Tuple, Iterator
+import datasets
+from datasets import Dataset
+from config import DataTrainingArguments, TrainingArguments, ModelArguments
+from hugnlp_trainer import HugTrainer
+from processors.ProcessorBase import DataProcessor
+from evaluators.EvaluatorBase import NO_GENERATE, DO_GENERATE, Evaluator, ClassificationEvaluator, GenerationEvaluator
+from metrics.classification_metric import ClassificationMetric
+from tools.runner_utils.log_util import logging
+from tools.computations.softmax import softmax
+from tools.model_utils.calibrate import CausalCLSCalibrator
+
+logger = logging.getLogger(__name__)
+
+
+"""
+Evaluator for pair-wise reward model
+"""
+class PairwiseRewardEvaluator(ClassificationEvaluator):
+
+    def __init__(
+        self,
+        model_args: ModelArguments,
+        data_args: DataTrainingArguments,
+        training_args: TrainingArguments,
+        processor: DataProcessor,
+        model: torch.nn.Module,
+        trainer: Optional[HugTrainer] = None,
+        eval_dataset: Optional[Dataset] = None,
+        test_dataset: Optional[Dataset] = None,
+    ) -> None:
+        super().__init__(model_args, data_args, training_args, processor, model, trainer, eval_dataset, test_dataset)
+        self.paradigm = NO_GENERATE
+
+
+    def default_compute_metrics(self, eval_predictions):
+        """
+        Design for the default metrics calculation for the current task.
+        Note:
+        - If the task processor has attribution of 'compute_metrics', this function will not be used.
+        - If this pre-built function can match your demand, you can omit the definition of 'compute_metrics' in your processor.
+        """
+        examples = self.eval_dataset
+        labels = examples["label"]
+
+        golden = {}
+        # predictions:  {"xx": "xxx", ...}
+        predictions, _ = self.get_best_and_topk(eval_predictions[0], examples, stage="dev")
+        for example in examples:
+            try:
+                idx = int(example["idx"])
+            except:
+                idx = int(example["idx"].split("-")[1]) # e.g.,  "dev-12" -> "12"
+
+            golden[idx] = example["label"]
+
+        all_metrics = {
+            "eval_macro_f1": 0.,
+            "eval_acc": 0.,
+        }
+
+        metric = ClassificationMetric()
+        gold = {k: v for k, v in golden.items()}
+        pred = {k: v for k, v in predictions.items()}
+        score = metric.calc_metric(golden=gold, predictions=pred)
+        acc, f1 = score["acc"], score["f1"]
+        all_metrics["eval_macro_f1"] += f1
+        all_metrics["eval_acc"] += acc
+        return all_metrics
+
+
+    def evaluate(self, test_dataset=None):
+
+        """
+        Each example has following two sequence:
+        - chosen: the better response
+        - rejected: the worse response
+        We need the model assign high reward for chosen than rejected sequence.
+        Thus, we calculate the accuracy that the reward value of chosen sequence derived from the reward model higher than the rejected sequence.
+        """
+        eval_dataset = self.eval_dataset if test_dataset is not None else test_dataset
+        all_chosen_values, all_rejected_values = list(), list()
+        for ei, data in enumerate(tqdm(eval_dataset)):
+            # chosen_input_ids, chosen_attention_mask = data["chosen_sequence"], data["chosen_attention_mask"]
+            # rejected_input_ids, rejected_attention_mask = data["rejected_sequence"], data["rejected_attention_mask"]
+            chosen_output = self.model(**data)
+            chosen_values, rejected_values = chosen_output["chosen_values"], chosen_output["rejected_values"]
+            all_chosen_values.extend(chosen_values.detach().cpu().numpy().tolist())
+            all_rejected_values.extend(rejected_values.detach().cpu().numpy().tolist())
+
+        metrics = dict()
+        acc = 0.
+        for chosen_value, rejected_value in zip(all_chosen_values, all_rejected_values):
+            if chosen_value >= rejected_value:
+                acc += 1
+        metrics["acc"] = round(acc / len(all_chosen_values), 4)
+        self.trainer.log_metrics("eval", metrics)
+        self.trainer.save_metrics("eval", metrics)
+
+
+    def predict(self):
+
+        self.evaluate(test_dataset=self.test_dataset)
+
+    def get_best_and_topk(self, logits, examples, topk=10, stage="dev"):
+        pass

hugnlp_runner.py

@@ -50,7 +50,7 @@ def main():
     model_args, data_args, training_args, semi_training_args = parser.parse_args_into_dataclasses()
 
     # Print hello world
-    if training_args.local_rank == 0:
+    if training_args.local_rank <= 0 or os.environ['LOCAL_RANK'] == "0":
         print_hello()
 
     training_args.output_dir = os.path.join(training_args.output_dir, list(filter(None, model_args.model_name_or_path.split("/")))[-1])

loss/rl_loss.py

@@ -0,0 +1,122 @@
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+def masked_mean(tensor: torch.Tensor, mask: torch.Tensor, dim: int = 1) -> torch.Tensor:
+    tensor = tensor * mask
+    tensor = tensor.sum(dim=dim)
+    mask_sum = mask.sum(dim=dim)
+    mean = tensor / (mask_sum + 1e-8)
+    return mean
+
+
+class GPTLMLoss(nn.Module):
+    """
+    GPT Language Model Loss
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.loss = nn.CrossEntropyLoss()
+
+    def forward(self, logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+        shift_logits = logits[..., :-1, :].contiguous()
+        shift_labels = labels[..., 1:].contiguous()
+        # Flatten the tokens
+        return self.loss(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
+
+
+class PolicyLoss(nn.Module):
+    """
+    Policy Loss for PPO
+    """
+
+    def __init__(self, clip_eps: float = 0.2) -> None:
+        super().__init__()
+        self.clip_eps = clip_eps
+
+    def forward(self,
+                log_probs: torch.Tensor,
+                old_log_probs: torch.Tensor,
+                advantages: torch.Tensor,
+                action_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        ratio = (log_probs - old_log_probs).exp()
+        surr1 = ratio * advantages
+        surr2 = ratio.clamp(1 - self.clip_eps, 1 + self.clip_eps) * advantages
+        loss = -torch.min(surr1, surr2)
+        if action_mask is not None:
+            loss = masked_mean(loss, action_mask)
+        loss = loss.mean()
+        return loss
+
+
+class ValueLoss(nn.Module):
+    """
+    Value Loss for PPO
+    """
+
+    def __init__(self, clip_eps: float = 0.4) -> None:
+        super().__init__()
+        self.clip_eps = clip_eps
+
+    def forward(self,
+                values: torch.Tensor,
+                old_values: torch.Tensor,
+                reward: torch.Tensor,
+                action_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        values_clipped = old_values + (values - old_values).clamp(-self.clip_eps, self.clip_eps)
+        surr1 = (values_clipped - reward)**2
+        surr2 = (values - reward)**2
+        loss = torch.max(surr1, surr2)
+        loss = loss.mean()
+        return 0.5 * loss
+
+
+class PPOPtxActorLoss(nn.Module):
+    """
+    To Do:
+
+    PPO-ptx Actor Loss
+    """
+
+    def __init__(self, policy_clip_eps: float = 0.2, pretrain_coef: float = 0.0, pretrain_loss_fn=GPTLMLoss()) -> None:
+        super().__init__()
+        self.pretrain_coef = pretrain_coef
+        self.policy_loss_fn = PolicyLoss(clip_eps=policy_clip_eps)
+        self.pretrain_loss_fn = pretrain_loss_fn
+
+    def forward(self,
+                log_probs: torch.Tensor,
+                old_log_probs: torch.Tensor,
+                advantages: torch.Tensor,
+                lm_logits: torch.Tensor,
+                lm_input_ids: torch.Tensor,
+                action_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        policy_loss = self.policy_loss_fn(log_probs, old_log_probs, advantages, action_mask=action_mask)
+        lm_loss = self.pretrain_loss_fn(lm_logits, lm_input_ids)
+        return policy_loss + self.pretrain_coef * lm_loss
+
+
+class LogSigLoss(nn.Module):
+    """
+    Pairwise Loss for Reward Model
+    Details: https://arxiv.org/abs/2203.02155
+    """
+
+    def forward(self, chosen_reward: torch.Tensor, reject_reward: torch.Tensor) -> torch.Tensor:
+        probs = torch.sigmoid(chosen_reward - reject_reward)
+        log_probs = torch.log(probs)
+        loss = -log_probs.mean()
+        return loss
+
+
+class LogExpLoss(nn.Module):
+    """
+    Pairwise Loss for Reward Model
+    Details: https://arxiv.org/abs/2204.05862
+    """
+
+    def forward(self, chosen_reward: torch.Tensor, reject_reward: torch.Tensor) -> torch.Tensor:
+        loss = torch.log(1 + torch.exp(reject_reward - chosen_reward)).mean()
+        return loss

models/__init__.py

@@ -73,7 +73,9 @@
 
 from models.reinforcement_learning.actor import CausalActor
 from models.reinforcement_learning.critic import AutoModelCritic
-from models.reinforcement_learning.reward_model import AutoModelReward
+from models.reinforcement_learning.reward_model import (
+    RobertaForReward, GPT2ForReward
+)
 
 # Models for pre-training
 PRETRAIN_MODEL_CLASSES = {
@@ -199,7 +201,13 @@
 REINFORCEMENT_MODEL_CLASSES = {
     "causal_actor": CausalActor,
     "auto_critic": AutoModelCritic, 
-    "auto_reward": AutoModelReward,
+    "rl_reward": {
+        "roberta": RobertaForReward,
+        "gpt2": GPT2ForReward,
+        "gpt-neo": None,
+        "opt": None,
+        "llama": None,
+    }
 }
 
 # task_type 负责对应model类型