Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+ric_txt2img/docs/source/en/imgs/access_request.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,10 @@
+ric_txt2img/results/*
+ric_txt2img/pretrained/*
+images/*
+ric_txt2img/requirements.txt
+ric_txt2img/eval_jpeg_aes_tmp.sh
+download_data.py
+check_data.py
+process_data.py
+ric_txt2img/data/*
+ric_txt2img/experiments/

.vscode/settings.json

@@ -0,0 +1,5 @@
+{
+    "python-envs.defaultEnvManager": "ms-python.python:conda",
+    "python-envs.defaultPackageManager": "ms-python.python:conda",
+    "python-envs.pythonProjects": []
+}

README.md

@@ -0,0 +1,110 @@
+# RiC: Multi-objective Alignment of Foundation Models with Dynamic Preference Adjustment
+
+Code for the ICML 2024 paper "Rewards-in-Context: Multi-objective Alignment of Foundation Models with Dynamic Preference Adjustment". This repo is based on [trl](https://github.com/huggingface/trl).
+
+This readme file is for LLM tasks, please refer to the 'ric_txt2img' directory for the code for conditional diffusion model training.
+
+**Note:** We have updated the code to include the EOS token in data preparation for SFT and RiC, which helps resolve the issue of non-ending generation. As a result, the model can produce more natural responses, although the results may differ from the paper. If you wish to achieve the exact performance in the paper, please use an old version prior to November 3, 2024.
+
+## 1. Installation
+Install the requirements.
+```
+pip install -r requirements.txt
+```
+
+The following issue has been fixed using token_ids directly for response_template as in https://huggingface.co/docs/trl/main/en/sft_trainer#train-on-completions-only.
+```
+Fix the trl package.
+
+
+The trl package currently has an issue with its 'DataCollatorForCompletionOnlyLM' implementation, wherein the first token of the template may be nonsensical and unmatched. Hence, it is necessary to make modifications to the trl/trainer/utils.py file within your trl package (usually under the path {path to conda or python}/lib/python3/site-packages), specifically the 'torch_call' function starting from line 119. Upon installing the trl package based on the requirements.txt, you can easily replace the trl/trainer/utils.py with the utils.py file located in the 'fix_trl' directory.
+```
+
+## 2. Usage
+### 2.1 RiC
+```
+Note: To reproduce the results in the paper, ensure the same training steps, for example, (4 processes, batchsize=2, steps=20000), and load the model in 8bit (--load_in_8bit True). We have observed notable performance differences between 8-bit and 16-bit models, particularly in the summary task.
+```
+
+
+* **Preparing datasets**: the first step is to generate dataset for RiC training.
+```
+cd ./ric
+CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch prepare_dataset_with_rewards.py --reward_names 'harmless,helpful' --exp_type 'assistant' --save_directory './datasets/train_harmhelp.hf' 
+```
+```
+CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch prepare_dataset_with_rewards.py --reward_names 'summary,faithful' --exp_type 'summary' --save_directory './datasets/summary_pref1faithful.hf'
+```
+Here, 'exp_type' includes two types of tasks, i.e., 'assistant' and 'summary'. Moreover, 'reward_names' can be sampled from \{'harmless', 'helpful', 'humor', 'summary', 'deberta', 'faithful'\}, where the first three are used for the 'assistant' task while the last three are used for the 'summary' task.
+
+* **Training**: you can train RiC through the main.py file.
+
+Offline training:
+```
+CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch main.py --train_dataset_path {path_to_corresponding_dataset} --exp_type 'assistant' --reward_names 'harmless,helpful' --training_steps 20000 --num_online_iterations 0 --wandb_name 'ric_assistant_harmlesshelpful_offline20000' --batch_size 2 --load_in_8bit True
+```
+Offline training + online training:
+```
+CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch main.py --train_dataset_path {path_to_corresponding_dataset} --exp_type 'assistant' --reward_names 'harmless,helpful' --training_steps 20000 --online_training_steps 4000 --num_online_iterations 2 --wandb_name 'ric_assistant_harmlesshelpful_offline20000_onlineiter2' --batch_size 2 --load_in_8bit True
+```
+
+* **Evaluation**: After training, models are saved into 'save_directory' path. You can evaluate models using the evaluation.py.
+```
+CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch evaluation.py --peft_name {path_to_the_saved_peft} --reward_names 'harmless,helpful' --exp_type 'assistant' --wandb_name {path_of_the_experiment}
+```
+
+### 2.2 SFT
+* Training a SFT model:
+```
+cd ./sft
+CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch sft.py --base_model_name 'meta-llama/Llama-2-7b-hf' --exp_type 'summary' --wandb_name {name_of_the_experiment} 
+```
+* Evaluating a SFT model:
+```
+CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch eval_multiobjective.py --base_model_name {path_to_the_saved_SFT_model} --exp_type 'summary' --wandb_name {name_of_the_experiment} --reward_names 'summary,faithful'
+```
+
+* Merging the SFT model if using lora for finetuning:
+```
+python3 merge_sft_lora.py --base_model_name 'meta-llama/Llama-2-7b-hf' --lora_name {path_to_the_lora_file} --save_directory {path_to_the_saving_directory}
+```
+
+### 2.3 MORLHF 
+MORLHF optimizes preference weighted rewards using PPO, based on the SFT model in 2.2.
+
+Train MORLHF:
+```
+cd ./ppo
+CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch morlhf.py --preference 0.5 --base_model_name {path_to_the_sft_model} --reward_names 'harmless,helpful' --exp_type 'assistant' --wandb_name 'morlhf_llamma2'
+```
+Here, the 'preference' is the preference for the first reward.
+
+Evaluate MORLHF model:
+```
+CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch eval_ppo_single_model.py --base_model_name {path_to_the_saved_morlhf_model} --reward_names 'harmless,helpful' --exp_type 'assistant' --wandb_name 'eval_morlhf_llamma2'
+```
+
+### 2.4 Rewarded Soups
+Rewarded Soups train $N$ PPO models for $N$ reward models.  
+
+Train PPO model for each reward model:
+```
+cd ./ppo
+CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch ppo.py --reward_name 'harmless' --base_model_name {path_to_the_SFT_model} --exp_type 'assistant' --wandb_name {name_for_this_experiment}
+```
+Evaluate Rewarded Soups with multiple PPO models:
+```
+CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch eval_rewarded_soups.py --reward_names 'harmless,helpful' --base_model_path1 {path_to_the_PPO_model_for_harmless} --base_model_path2 {path_to_the_PPO_model_for_helpful} --exp_type 'assistant' --wandb_name 'eval_pposoups_harmless_helpful'
+```
+
+
+## 3. Citation
+If you find our work useful for your research, please cite:
+```
+@article{yang2024rewards,
+  title={Rewards-in-Context: Multi-objective Alignment of Foundation Models with Dynamic Preference Adjustment},
+  author={Yang, Rui and Pan, Xiaoman and Luo, Feng and Qiu, Shuang and Zhong, Han and Yu, Dong and Chen, Jianshu},
+  journal={International Conference on Machine Learning},
+  year={2024}
+}
+```

datasets/hh-rlhf/.gitattributes

@@ -0,0 +1,54 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.lz4 filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+# Audio files - uncompressed
+*.pcm filter=lfs diff=lfs merge=lfs -text
+*.sam filter=lfs diff=lfs merge=lfs -text
+*.raw filter=lfs diff=lfs merge=lfs -text
+# Audio files - compressed
+*.aac filter=lfs diff=lfs merge=lfs -text
+*.flac filter=lfs diff=lfs merge=lfs -text
+*.mp3 filter=lfs diff=lfs merge=lfs -text
+*.ogg filter=lfs diff=lfs merge=lfs -text
+*.wav filter=lfs diff=lfs merge=lfs -text
+# Image files - uncompressed
+*.bmp filter=lfs diff=lfs merge=lfs -text
+*.gif filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.tiff filter=lfs diff=lfs merge=lfs -text
+# Image files - compressed
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.webp filter=lfs diff=lfs merge=lfs -text

datasets/hh-rlhf/README.md

@@ -0,0 +1,68 @@
+---
+license: mit
+tags:
+  - human-feedback
+---
+
+# Dataset Card for HH-RLHF
+
+## Dataset Summary
+
+
+This repository provides access to two different kinds of data:
+1. Human preference data about helpfulness and harmlessness from [Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback](https://arxiv.org/abs/2204.05862). These data are meant to train preference (or reward) models for subsequent RLHF training. These data are *not* meant for supervised training of dialogue agents. Training dialogue agents on these data is likely to lead to harmful models and this shold be avoided.
+2. Human-generated and annotated red teaming dialogues from [Red Teaming Language Models to Reduce Harms: Methods, Scaling Behaviors, and Lessons Learned](https://www.anthropic.com/red_teaming.pdf). These data are meant to understand how crowdworkers red team models and what types of red team attacks are succesful or not. The data are *not* meant for fine-tuning or preference modeling (use the data above for preference modeling). These data are entire transcripts of conversations that are derived from the harmlessness preference modeling data described above, where only the chosen response is incorporated into the overall transcript. Furthermore, the transcripts are annotated with human and automated measurements of how harmful the overall dialogues are.
+
+**Disclaimer**: The data (especially the harmlessness preference data and the red team data) contain content that may be offensive or upsetting. Topics include, but are not limited to, discriminatory language and discussions of abuse, violence, self-harm, exploitation, and other potentially upsetting subject matter. Please only engage with the data in accordance with your own personal risk tolerance. The data are intended for research purposes, especially research that can make models *less* harmful. The views expressed in the data do not reflect the views of Anthropic or any of its employees. As mentioned above, these data are *not* intended for training dialogue agents as this will likely lead to harmful model behavior. 
+
+Each of these datasets are described further below. 
+
+## Human preference data about helpfulness and harmlessness (PM Data)
+
+The data are described in the paper: [Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback](https://arxiv.org/abs/2204.05862). If you find the data useful, please cite the paper. The data format is very simple -- each line of the jsonl files contains a pair of texts, one "chosen" and one "rejected".
+
+For **helpfulness**, the data are grouped into train/test splits in three tranches: from our base models (context-distilled 52B language models), via rejection sampling (mostly with best-of-16 sampling) against an early preference model, and a dataset sampled during our iterated "online" process. 
+
+For **harmlessness**, the data are only collected for our base models, but otherwise formatted in the same way.
+
+Details about the data collection process and crowdworker population can be found in the paper, specifically in section 2 and appendix D. 
+
+## Red teaming data (not PM Data)
+
+The data are described in the paper: [Red Teaming Language Models to Reduce Harms: Methods, Scaling Behaviors, and Lessons Learned](https://www.anthropic.com/red_teaming.pdf). If you find the data useful, please cite the paper. At a high level these data can be examined to understand what succesful red team attempts look like. 
+
+Details about the data and data collection procedures can be found in the Datasheet in the appendix of the paper. 
+
+Each line of the jsonl file contains a dictionary with the following fields:
+- `transcript` a text transcript of a conversation between a human adversary (red team member) and an AI assistant
+- `min_harmlessness_score_transcript` a real value score of the harmlessness of the AI assistant (lower is more harmful) as obtained from a preference model
+- `num_params` number of parameters in the language model powering the AI assistant
+- `model_type` type of model powering the AI assistant
+- `rating` the red team member's rating of how successful they were at breaking the AI assistant (Likert scale, higher is more successful)  
+- `task_description` a short text description written by the red team member about how they tried to red team the AI assistant
+- `task_description_harmlessness_score` a real value score of the harmlessness of the task description (lower is more harmful) as obtained from a preference model
+- `red_team_member_id` an arbitrary identifier of the red team member. one red team member can generate multiple red team attacks
+- `is_upworker` a binary indicator that is true if the red team member was from the crowd platform Upwork or false if they were from MTurk
+- `tags` a list of up to 6 tags per transcript. tags are short descriptions of the red team attempts generated by crowdworkers who reviewed red team data post-hoc. tags were only provided for a random sample of 1000 red team attempts for two of four model types.
+
+## Usage
+
+Each of the above datasets is located in a separate sub-directory. To load an individual subset, use the `data_dir` argument of the `load_dataset()` function as follows:
+
+```python
+from datasets import load_dataset
+
+# Load all helpfulness/harmless subsets (share the same schema)
+dataset = load_dataset("Anthropic/hh-rlhf")
+
+# Load one of the harmless subsets
+dataset = load_dataset("Anthropic/hh-rlhf", data_dir="harmless-base")
+
+# Load the red teaming subset
+dataset = load_dataset("Anthropic/hh-rlhf", data_dir="red-team-attempts")
+```
+
+## Contact
+
+The original authors host this dataset on GitHub here: https://github.com/anthropics/hh-rlhf
+You can submit inquiries to: redteam@anthropic.com

datasets/hh-rlhf/harmless-base/test.jsonl.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ebfaaed21162a4de120ae85366075b34425e2a6303fac59b99555002f7016f03
+size 742801

datasets/hh-rlhf/harmless-base/train.jsonl.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bf3e1ecd3db61bf21d190f69351d26670c67aababa0fd282fa008c17a7202217
+size 13197306

datasets/hh-rlhf/helpful-base/test.jsonl.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8be3fc1a13b27901631696f2be6f184c799f1baeaf145f53ac5db24960adc37b
+size 875191

datasets/hh-rlhf/helpful-base/train.jsonl.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:518a5bf288456fc9f3b7c980c54116fba0c52f274d3f4d344675d83e4058f6f4
+size 16200131

datasets/hh-rlhf/helpful-online/test.jsonl.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:39a475cc7e0df972902f8e3a0972726fce826ac65369bda8c671de4330c60a10
+size 1053578

datasets/hh-rlhf/helpful-online/train.jsonl.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0994d18b82083f5c084ce04071b683958a724df8c7215ab1d851986642d36a2a
+size 20127228

datasets/hh-rlhf/helpful-rejection-sampled/test.jsonl.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:abc35dbc56015db1dae0faa484514ccbc3eee0133bfcd557fc69250619b028bb
+size 1361231

datasets/hh-rlhf/helpful-rejection-sampled/train.jsonl.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0cea2d2ab2187624ec1314e8cc6fb579f0fa987b3a33e8102b7556ed4d255f64
+size 25697148

datasets/hh-rlhf/red-team-attempts/red_team_attempts.jsonl.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4c7b0069991460f0064f279fd400b51f3f0095697d14d7793c49b0925f80814f
+size 15483307

export_hf_checkpoint.py

@@ -0,0 +1,66 @@
+import os
+
+import torch
+import transformers
+from peft import PeftModel
+from transformers import LlamaForCausalLM, LlamaTokenizer, AutoTokenizer  # noqa: F402
+
+# BASE_MODEL = os.environ.get("BASE_MODEL", None)
+# assert (
+#     BASE_MODEL
+# ), "Please specify a value for BASE_MODEL environment variable, e.g. `export BASE_MODEL=huggyllama/llama-7b`"  # noqa: E501
+
+
+BASE_MODEL = '/home/hector5/models/Llama-3-Base-8B-SFT/'
+
+
+# LORA_WEIGHT = '/home/hector5/RiC/ppo/ckpt_morlhf/morlhf_llamma3_lora-0.33_0.33_0.33/epoch_1/'
+# LORA_WEIGHT = '/home/hector5/RiC/ppo/ckpt_morlhf/morlhf_llamma3_lora-0.2_0.2_0.6/epoch_1/'
+LORA_WEIGHT = '/home/hector5/RiC/ppo/ckpt_morlhf/morlhf_llamma3_lora-0.0_0.0_1.0/epoch_0/'
+
+OUTPUT_path = './merged_models/Llama-3-Base-8B-SFT-morlhf-humor'
+
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+tokenizer.save_pretrained(OUTPUT_path)
+
+base_model = LlamaForCausalLM.from_pretrained(
+    BASE_MODEL,
+    load_in_8bit=False,
+    torch_dtype=torch.bfloat16,
+    device_map={"": "cpu"},
+)
+
+first_weight = base_model.model.layers[0].self_attn.q_proj.weight
+first_weight_old = first_weight.clone()
+
+lora_model = PeftModel.from_pretrained(
+    base_model,
+    LORA_WEIGHT,
+    device_map={"": "cpu"},
+    torch_dtype=torch.bfloat16,
+)
+
+lora_weight = lora_model.base_model.model.model.layers[
+    0
+].self_attn.q_proj.weight
+
+assert torch.allclose(first_weight_old, first_weight)
+
+# merge weights - new merging method from peft
+lora_model = lora_model.merge_and_unload()
+
+lora_model.train(False)
+
+# did we do anything?
+assert not torch.allclose(first_weight_old, first_weight)
+
+lora_model_sd = lora_model.state_dict()
+deloreanized_sd = {
+    k.replace("base_model.model.", ""): v
+    for k, v in lora_model_sd.items()
+    if "lora" not in k
+}
+
+LlamaForCausalLM.save_pretrained(
+    base_model, OUTPUT_path, state_dict=deloreanized_sd, max_shard_size="2048MB"
+)

fix_trl/utils.py

@@ -0,0 +1,739 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import random
+import warnings
+from collections import deque
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from accelerate import PartialState
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import IterableDataset
+from transformers import BitsAndBytesConfig, DataCollatorForLanguageModeling, PreTrainedTokenizerBase
+
+from ..import_utils import is_peft_available, is_unsloth_available, is_xpu_available
+from ..trainer.model_config import ModelConfig
+
+
+if is_peft_available():
+    from peft import LoraConfig, PeftConfig
+
+
+class AdaptiveKLController:
+    """
+    Adaptive KL controller described in the paper:
+    https://arxiv.org/pdf/1909.08593.pdf
+    """
+
+    def __init__(self, init_kl_coef, target, horizon):
+        self.value = init_kl_coef
+        self.target = target
+        self.horizon = horizon
+
+    def update(self, current, n_steps):
+        target = self.target
+        proportional_error = np.clip(current / target - 1, -0.2, 0.2)
+        mult = 1 + proportional_error * n_steps / self.horizon
+        self.value *= mult
+
+
+class FixedKLController:
+    """Fixed KL controller."""
+
+    def __init__(self, kl_coef):
+        self.value = kl_coef
+
+    def update(self, current, n_steps):
+        pass
+
+
+class DataCollatorForCompletionOnlyLM(DataCollatorForLanguageModeling):
+    """
+    Data collator used for completion tasks. It ensures that all the tokens of the labels are set to an 'ignore_index'
+    when they do not come from the assistant. This ensure that the loss is only
+    calculated on the completion made by the assistant.
+
+    Args:
+        response_template (`Union[str, List[int]]`): the template form that indicates the start of the response, typically something like
+            '### Response:\n'. It can also be passed as tokenized ids, which can be useful when using a tokenizer that encodes the response
+            differently if it does not have proper context.
+        instruction_template (`Union[str, List[int]]`): the template form that indicates the start of the human instruction, typically something like
+            '### Human:\n'. Useful for assistant-style conversation datasets. It can also be passed as tokenized ids.
+        mlm (`bool`, *optional*, defaults to `False`): Whether or not to use masked language modeling in the underlying
+            `DataCollatorForLanguageModeling` class. Note that this option currently has no effect but is present
+             for flexibility and backwards-compatibility.
+        ignore_index (`int`, *optional*, defaults to `-100`):
+            The index to use to ignore the initial tokens with
+    """
+
+    def __init__(
+        self,
+        response_template: Union[str, List[int]],
+        instruction_template: Optional[Union[str, List[int]]] = None,
+        *args,
+        mlm: bool = False,
+        ignore_index: int = -100,
+        **kwargs,
+    ):
+        super().__init__(*args, mlm=mlm, **kwargs)
+
+        self.instruction_template = instruction_template
+        if isinstance(instruction_template, str):
+            # The user provides a string, must tokenize
+            self.instruction_token_ids = self.tokenizer.encode(self.instruction_template, add_special_tokens=False)
+        else:
+            # The user already provides the token ids
+            self.instruction_token_ids = instruction_template
+
+        self.response_template = response_template
+        if isinstance(response_template, str):
+            # The user provides a string, must tokenize
+            self.response_token_ids = self.tokenizer.encode(self.response_template, add_special_tokens=False)
+        else:
+            # The user already provides the token ids
+            self.response_token_ids = response_template
+
+        if not self.mlm and self.instruction_template and self.tokenizer.pad_token_id == self.tokenizer.eos_token_id:
+            warnings.warn(
+                "The pad_token_id and eos_token_id values of this tokenizer are identical. "
+                "If you are planning for multi-turn training, "
+                "it can result in the model continuously generating questions and answers without eos token. "
+                "To avoid this, set the pad_token_id to a different value."
+            )
+
+        self.ignore_index = ignore_index
+
+    def torch_call(self, examples: List[Union[List[int], Any, Dict[str, Any]]]) -> Dict[str, Any]:
+        batch = super().torch_call(examples)
+
+        if self.instruction_template is None:
+            for i in range(len(examples)):
+                response_token_ids_start_idx = None
+
+                for idx in np.where(batch["labels"][i] == self.response_token_ids[0])[0]:
+                    # `response_token_ids` is `'### Response:\n'`, here we are just making sure that the token IDs match
+                    if (
+                        self.response_token_ids
+                        == batch["labels"][i][idx : idx + len(self.response_token_ids)].tolist()
+                    ):
+                        response_token_ids_start_idx = idx
+
+                
+                if response_token_ids_start_idx is None:
+                    for idx in np.where(batch["labels"][i] == self.response_token_ids[1])[0]:
+                        if (
+                            self.response_token_ids[1:]
+                             == batch["labels"][i][idx : idx + len(self.response_token_ids[1:])].tolist()
+                        ):
+                            response_token_ids_start_idx = idx
+
+                    if response_token_ids_start_idx is None:
+                        warnings.warn(
+                            f"Could not find response key `{self.response_template}` in the "
+                            f'following instance: {self.tokenizer.decode(batch["input_ids"][i])} '
+                            f"This instance will be ignored in loss calculation. "
+                            f"Note, if this happens often, consider increasing the `max_seq_length`."
+                        )
+                        batch["labels"][i, :] = self.ignore_index
+                    else:
+                        response_token_ids_end_idx = response_token_ids_start_idx + len(self.response_token_ids[1:])
+                else:
+                    response_token_ids_end_idx = response_token_ids_start_idx + len(self.response_token_ids)
+
+                # Make pytorch loss function ignore all tokens up through the end of the response key
+                batch["labels"][i, :response_token_ids_end_idx] = self.ignore_index
+
+        else:
+            for i in range(len(examples)):
+                response_token_ids_idxs = []
+                human_token_ids_idxs = []
+
+                for assistant_idx in np.where(batch["labels"][i] == self.response_token_ids[0])[0]:
+                    # find the indexes of the start of a response.
+                    if (
+                        self.response_token_ids
+                        == batch["labels"][i][assistant_idx : assistant_idx + len(self.response_token_ids)].tolist()
+                    ):
+                        response_token_ids_idxs.append(assistant_idx + len(self.response_token_ids))
+
+                if len(response_token_ids_idxs) == 0:
+                    warnings.warn(
+                        f"Could not find response key `{self.response_template}` in the "
+                        f'following instance: {self.tokenizer.decode(batch["input_ids"][i])} '
+                        f"This instance will be ignored in loss calculation. "
+                        f"Note, if this happens often, consider increasing the `max_seq_length`."
+                    )
+                    batch["labels"][i, :] = self.ignore_index
+
+                human_token_ids = self.instruction_token_ids
+                for human_idx in np.where(batch["labels"][i] == human_token_ids[0])[0]:
+                    # find the indexes of the start of a human answer.
+                    if human_token_ids == batch["labels"][i][human_idx : human_idx + len(human_token_ids)].tolist():
+                        human_token_ids_idxs.append(human_idx)
+
+                if len(human_token_ids_idxs) == 0:
+                    warnings.warn(
+                        f"Could not find instruction key `{self.instruction_template}` in the "
+                        f'following instance: {self.tokenizer.decode(batch["input_ids"][i])} '
+                        f"This instance will be ignored in loss calculation. "
+                        f"Note, if this happens often, consider increasing the `max_seq_length`."
+                    )
+                    batch["labels"][i, :] = self.ignore_index
+
+                if (
+                    len(human_token_ids_idxs) > 0
+                    and len(response_token_ids_idxs) > 0
+                    and human_token_ids_idxs[0] > response_token_ids_idxs[0]
+                ):
+                    human_token_ids_idxs = [0] + human_token_ids_idxs
+
+                for idx, (start, end) in enumerate(zip(human_token_ids_idxs, response_token_ids_idxs)):
+                    # Make pytorch loss function ignore all non response tokens
+                    if idx != 0:
+                        batch["labels"][i, start:end] = self.ignore_index
+                    else:
+                        batch["labels"][i, :end] = self.ignore_index
+
+                if len(response_token_ids_idxs) < len(human_token_ids_idxs):
+                    batch["labels"][i, human_token_ids_idxs[-1] :] = self.ignore_index
+
+        return batch
+
+
+@dataclass
+class RewardDataCollatorWithPadding:
+    r"""
+    Reward DataCollator class that pads the inputs to the maximum length of the batch.
+    Args:
+        tokenizer (`PreTrainedTokenizerBase`):
+            The tokenizer used for encoding the data.
+        padding (`Union[bool, str, `PaddingStrategy`]`, `optional`, defaults to `True`):
+            padding_strategy to pass to the tokenizer.
+        max_length (`Optional[int]`, `optional`, defaults to `None`):
+            The maximum length of the sequence to be processed.
+        pad_to_multiple_of (`Optional[int]`, `optional`, defaults to `None`):
+            If set will pad the sequence to a multiple of the provided value.
+        return_tensors (`str`, `optional`, defaults to `"pt"`):
+            The tensor type to use.
+    """
+
+    tokenizer: PreTrainedTokenizerBase
+    padding: Union[bool, str] = True
+    max_length: Optional[int] = None
+    pad_to_multiple_of: Optional[int] = None
+    return_tensors: str = "pt"
+
+    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
+        features_chosen = []
+        features_rejected = []
+        margin = []
+        # check if we have a margin. If we do, we need to batch it as well
+        has_margin = "margin" in features[0]
+        for feature in features:
+            # check if the keys are named as expected
+            if (
+                "input_ids_chosen" not in feature
+                or "input_ids_rejected" not in feature
+                or "attention_mask_chosen" not in feature
+                or "attention_mask_rejected" not in feature
+            ):
+                raise ValueError(
+                    "The features should include `input_ids_chosen`, `attention_mask_chosen`, `input_ids_rejected` and `attention_mask_rejected`"
+                )
+
+            features_chosen.append(
+                {
+                    "input_ids": feature["input_ids_chosen"],
+                    "attention_mask": feature["attention_mask_chosen"],
+                }
+            )
+            features_rejected.append(
+                {
+                    "input_ids": feature["input_ids_rejected"],
+                    "attention_mask": feature["attention_mask_rejected"],
+                }
+            )
+            if has_margin:
+                margin.append(feature["margin"])
+        batch_chosen = self.tokenizer.pad(
+            features_chosen,
+            padding=self.padding,
+            max_length=self.max_length,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            return_tensors=self.return_tensors,
+        )
+        batch_rejected = self.tokenizer.pad(
+            features_rejected,
+            padding=self.padding,
+            max_length=self.max_length,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            return_tensors=self.return_tensors,
+        )
+        batch = {
+            "input_ids_chosen": batch_chosen["input_ids"],
+            "attention_mask_chosen": batch_chosen["attention_mask"],
+            "input_ids_rejected": batch_rejected["input_ids"],
+            "attention_mask_rejected": batch_rejected["attention_mask"],
+            "return_loss": True,
+        }
+        if has_margin:
+            margin = torch.tensor(margin, dtype=torch.float)
+            batch["margin"] = margin
+        return batch
+
+
+@dataclass
+class DPODataCollatorWithPadding:
+    r"""
+    DPO DataCollator class that pads the tokenized inputs to the maximum length of the batch.
+    Args:
+        pad_token_id (`int` defaults to 0):
+            The tokenizer's pad_token_id.
+        label_pad_token_id (`int`, defaults to -100):
+            The label used for masking.
+        is_encoder_decoder (`Optional[bool]`, `optional`, defaults to `None`):
+            Whether or not you model has an encoder_decoder architecture.
+    """
+
+    pad_token_id: int = 0
+    label_pad_token_id: int = -100
+    is_encoder_decoder: Optional[bool] = False
+
+    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
+        # first, pad everything to the same length
+        padded_batch = {}
+        for k in features[0].keys():
+            if k.endswith("_input_ids") or k.endswith("_attention_mask") or k.endswith("_labels"):
+                if self.is_encoder_decoder:
+                    to_pad = [torch.LongTensor(ex[k]) for ex in features]
+
+                    if (k.startswith("prompt")) and (k.endswith("input_ids")):
+                        if self.pad_token_id is None:
+                            raise ValueError(
+                                "Padding is enabled, but the tokenizer is not configured with a padding token."
+                                " Explicitly set `tokenizer.pad_token` (e.g. `tokenizer.pad_token = tokenizer.eos_token`)"
+                                " before calling the trainer."
+                            )
+                        padding_value = self.pad_token_id
+                    elif k.endswith("_attention_mask"):
+                        padding_value = 0
+                    elif (k.startswith("chosen")) or (k.startswith("rejected")) or ("decoder" in k):
+                        padding_value = self.label_pad_token_id
+                    else:
+                        raise ValueError(f"Unexpected key in batch '{k}'")
+                    padded_batch[k] = pad_sequence(to_pad, batch_first=True, padding_value=padding_value)
+                else:
+                    # adapted from https://stackoverflow.com/questions/73256206
+                    if "prompt" in k:
+                        to_pad = [torch.LongTensor(ex[k][::-1]) for ex in features]
+                    else:
+                        to_pad = [torch.LongTensor(ex[k]) for ex in features]
+                    if k.endswith("_input_ids"):
+                        if self.pad_token_id is None:
+                            raise ValueError(
+                                "Padding is enabled, but the tokenizer is not configured with a padding token."
+                                " Explicitly set `tokenizer.pad_token` (e.g. `tokenizer.pad_token = tokenizer.eos_token`)"
+                                " before calling the trainer."
+                            )
+                        padding_value = self.pad_token_id
+                    elif k.endswith("_labels"):
+                        padding_value = self.label_pad_token_id
+                    elif k.endswith("_attention_mask"):
+                        padding_value = 0
+                    else:
+                        raise ValueError(f"Unexpected key in batch '{k}'")
+
+                    padded_batch[k] = pad_sequence(to_pad, batch_first=True, padding_value=padding_value)
+                    # for the prompt, flip back so padding is on left side
+                    if "prompt" in k:
+                        padded_batch[k] = padded_batch[k].flip(dims=[1])
+            elif k.endswith("_logps"):
+                # the cached reference model logprobs
+                padded_batch[k] = torch.tensor([ex[k] for ex in features])
+            else:
+                padded_batch[k] = [ex[k] for ex in features]
+
+        return padded_batch
+
+
+class ConstantLengthDataset(IterableDataset):
+    """
+    Iterable dataset that returns constant length chunks of tokens from stream of text files.
+    The dataset also formats the text before tokenization with a specific format that is provided
+    by the user.
+
+        Args:
+            tokenizer (`transformers.PreTrainedTokenizer`):
+                The processor used for processing the data.
+            dataset (`dataset.Dataset`):
+                Dataset with text files.
+            dataset_text_field (`str`, **optional**):
+                Name of the field in the dataset that contains the text. Used only if `formatting_func` is `None`.
+            formatting_func (`Callable`, **optional**):
+                Function that formats the text before tokenization. Usually it is recommended to have follows a certain
+                pattern such as `"### Question: {question} ### Answer: {answer}"`
+            infinite (`bool`, *optional*, defaults to `False`):
+                If True the iterator is reset after dataset reaches end else stops.
+            seq_length (`int`, *optional*, defaults to `1024`):
+                Length of token sequences to return.
+            num_of_sequences (`int`, *optional*, defaults to `1024`):
+                Number of token sequences to keep in buffer.
+            chars_per_token (`int`, *optional*, defaults to `3.6`):
+                Number of characters per token used to estimate number of tokens in text buffer.
+            eos_token_id (`int`, *optional*, defaults to `0`):
+                Id of the end of sequence token if the passed tokenizer does not have an EOS token.
+            shuffle ('bool', *optional*, defaults to True)
+                Shuffle the examples before they are returned
+            append_concat_token ('bool', *optional*, defaults to True)
+                If true, appends `eos_token_id` at the end of each sample being packed.
+            add_special_tokens ('bool', *optional*, defaults to True)
+                If true, tokenizers adds special tokens to each sample being packed.
+    """
+
+    def __init__(
+        self,
+        tokenizer,
+        dataset,
+        dataset_text_field=None,
+        formatting_func=None,
+        infinite=False,
+        seq_length=1024,
+        num_of_sequences=1024,
+        chars_per_token=3.6,
+        eos_token_id=0,
+        shuffle=True,
+        append_concat_token=True,
+        add_special_tokens=True,
+    ):
+        self.tokenizer = tokenizer
+
+        if tokenizer.eos_token_id is None:
+            warnings.warn(
+                "The passed tokenizer does not have an EOS token. We will use the passed eos_token_id instead which corresponds"
+                f" to {eos_token_id}. If this is not the correct EOS token, make sure to pass the correct eos_token_id."
+            )
+
+        self.concat_token_id = tokenizer.eos_token_id if tokenizer.eos_token_id else eos_token_id
+        self.dataset = dataset
+        self.seq_length = seq_length
+        self.infinite = infinite
+        self.current_size = 0
+        self.max_buffer_size = seq_length * chars_per_token * num_of_sequences
+        self.shuffle = shuffle
+        self.append_concat_token = append_concat_token
+        self.add_special_tokens = add_special_tokens
+        if formatting_func is None:
+            self.formatting_func = lambda x: x[dataset_text_field]
+        else:
+            self.formatting_func = formatting_func
+
+        if formatting_func is not None:
+            if formatting_func.__code__.co_argcount > 1:
+                warnings.warn(
+                    "The passed formatting_func has more than one argument. Usually that function should have a single argument `example`"
+                    " which corresponds to the dictionary returned by each element of the dataset. Make sure you know what you are doing."
+                )
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __iter__(self):
+        iterator = iter(self.dataset)
+        more_examples = True
+        while more_examples:
+            buffer, buffer_len = [], 0
+            while True:
+                if buffer_len >= self.max_buffer_size:
+                    break
+                try:
+                    buffer.append(self.formatting_func(next(iterator)))
+                    buffer_len += len(buffer[-1])
+                except StopIteration:
+                    if self.infinite:
+                        iterator = iter(self.dataset)
+                        warnings.warn("The dataset reached end and the iterator is reset to the start.")
+                    else:
+                        more_examples = False
+                        break
+            tokenized_inputs = self.tokenizer(buffer, add_special_tokens=self.add_special_tokens, truncation=False)[
+                "input_ids"
+            ]
+            all_token_ids = []
+            for tokenized_input in tokenized_inputs:
+                if self.append_concat_token:
+                    tokenized_input = tokenized_input + [self.concat_token_id]
+                all_token_ids.extend(tokenized_input)
+            examples = []
+            for i in range(0, len(all_token_ids), self.seq_length):
+                input_ids = all_token_ids[i : i + self.seq_length]
+                if len(input_ids) == self.seq_length:
+                    examples.append(input_ids)
+            if self.shuffle:
+                random.shuffle(examples)
+            for example in examples:
+                self.current_size += 1
+                yield {
+                    "input_ids": torch.LongTensor(example),
+                    "labels": torch.LongTensor(example),
+                }
+
+
+class RunningMoments:
+    def __init__(self, accelerator):
+        """
+        Calculates the running mean and standard deviation of a data stream. Reference:
+        https://github.com/OpenLMLab/MOSS-RLHF/blob/40b91eb2f2b71b16919addede0341d2bef70825d/utils.py#L75
+        """
+        self.mean = 0
+        self.std = 1
+        self.var = 1
+        self.count = 1e-24
+        self.accelerator = accelerator
+
+    @torch.no_grad()
+    def update(self, xs: torch.Tensor) -> Tuple[float, float]:
+        """
+        Updates running moments from batch's moments computed across ranks
+        """
+        if self.accelerator.use_distributed:
+            xs_mean, xs_var, xs_count = get_global_statistics(self.accelerator, xs)
+        else:
+            xs_count = xs.numel()
+            xs_var, xs_mean = torch.var_mean(xs, unbiased=False)
+        xs_mean, xs_var = xs_mean.float(), xs_var.float()
+
+        delta = xs_mean - self.mean
+        tot_count = self.count + xs_count
+
+        new_sum = xs_var * xs_count
+        # correct old_sum deviation accounting for the new mean
+        old_sum = self.var * self.count + delta**2 * self.count * xs_count / tot_count
+        tot_sum = old_sum + new_sum
+
+        self.mean += delta * xs_count / tot_count
+        self.var = tot_sum / tot_count
+        self.std = (self.var * tot_count / (tot_count - 1)).float().sqrt()
+        self.count = tot_count
+
+        return xs_mean.item(), (xs_var * xs_count / (xs_count - 1)).float().sqrt().item()
+
+
+@torch.no_grad()
+def get_global_statistics(accelerator, xs: torch.Tensor, mask=None, device="cpu") -> Tuple[float, float, int]:
+    """
+    Computes element-wise mean and variance of the tensor across processes. Reference:
+    https://github.com/OpenLMLab/MOSS-RLHF/blob/40b91eb2f2b71b16919addede0341d2bef70825d/utils.py#L57C1-L73C75
+    """
+    xs = xs.to(accelerator.device)
+    sum_and_count = torch.tensor([xs.sum(), (xs.numel() if mask is None else mask.sum())], device=xs.device)
+    sum_and_count = accelerator.reduce(sum_and_count)
+    global_sum, count = sum_and_count
+    global_mean = global_sum / count
+
+    sum_var = torch.sum(((xs - global_mean) ** 2).mul(1 if mask is None else mask))
+    sum_var = accelerator.reduce(sum_var)
+    global_var = sum_var / count
+
+    return global_mean.to(device), global_var.to(device), count.to(device)
+
+
+def compute_accuracy(eval_pred) -> Dict[str, float]:
+    predictions, labels = eval_pred
+    # Here, predictions is rewards_chosen and rewards_rejected.
+    # We want to see how much of the time rewards_chosen > rewards_rejected.
+    if np.array(predictions[:, 0] == predictions[:, 1], dtype=float).sum() > 0:
+        warnings.warn(
+            f"There are {np.array(predictions[:, 0] == predictions[:, 1]).sum()} out of {len(predictions[:, 0])} instances where the predictions for both options are equal. As a consequence the accuracy can be misleading."
+        )
+    predictions = np.argmax(predictions, axis=1)
+
+    accuracy = np.array(predictions == labels, dtype=float).mean().item()
+    return {"accuracy": accuracy}
+
+
+def pad_to_length(tensor: torch.Tensor, length: int, pad_value: Union[int, float], dim: int = -1) -> torch.Tensor:
+    if tensor.size(dim) >= length:
+        return tensor
+    else:
+        pad_size = list(tensor.shape)
+        pad_size[dim] = length - tensor.size(dim)
+        return torch.cat(
+            [
+                tensor,
+                pad_value * torch.ones(*pad_size, dtype=tensor.dtype, device=tensor.device),
+            ],
+            dim=dim,
+        )
+
+
+def disable_dropout_in_model(model: torch.nn.Module) -> None:
+    for module in model.modules():
+        if isinstance(module, torch.nn.Dropout):
+            module.p = 0
+
+
+def exact_div(a, b, a_str, b_str, custom_error_message=""):
+    q = a // b
+    if a != q * b:
+        raise ValueError(f"{custom_error_message}, {a_str}={a}, {b_str}={b}, inexact division: {a} / {b} = {a / b}")
+    return q
+
+
+# copied from https://github.com/kvablack/ddpo-pytorch/blob/main/ddpo_pytorch/stat_tracking.py#L5
+class PerPromptStatTracker:
+    r"""
+    Class for tracking statistics per prompt. Mainly used to calculate advantage for the DPPO algorithm
+
+    Args:
+        buffer_size (`int`):
+            Size of the buffer to keep for each prompt.
+        min_count (`int`):
+            Minimum number of samples to keep in the buffer before calculating the mean and std.
+    """
+
+    def __init__(self, buffer_size, min_count):
+        self.buffer_size = buffer_size
+        self.min_count = min_count
+        self.stats = {}
+
+    def update(self, prompts, rewards):
+        prompts = np.array(prompts)
+        rewards = np.array(rewards)
+        unique = np.unique(prompts)
+        advantages = np.empty_like(rewards)
+        for prompt in unique:
+            prompt_rewards = rewards[prompts == prompt]
+            if prompt not in self.stats:
+                self.stats[prompt] = deque(maxlen=self.buffer_size)
+            self.stats[prompt].extend(prompt_rewards)
+
+            if len(self.stats[prompt]) < self.min_count:
+                mean = np.mean(rewards)
+                std = np.std(rewards) + 1e-6
+            else:
+                mean = np.mean(self.stats[prompt])
+                std = np.std(self.stats[prompt]) + 1e-6
+            advantages[prompts == prompt] = (prompt_rewards - mean) / std
+
+        return advantages
+
+    def get_stats(self):
+        return {k: {"mean": np.mean(v), "std": np.std(v), "count": len(v)} for k, v in self.stats.items()}
+
+
+def neftune_post_forward_hook(module, input, output):
+    """
+    Implements the NEFTune forward pass for the model using forward hooks. Note this works only for
+    torch.nn.Embedding layers. This method is slightly adapted from the original source code
+    that can be found here: https://github.com/neelsjain/NEFTune
+
+    Simply add it to your model as follows:
+    ```python
+    model = ...
+    model.embed_tokens.neftune_noise_alpha = 0.1
+    model.embed_tokens.register_forward_hook(neftune_post_forward_hook)
+    ```
+
+    Args:
+        module (`torch.nn.Module`):
+            The embedding module where the hook is attached. Note that you need to set
+            `module.neftune_noise_alpha` to the desired noise alpha value.
+        input (`torch.Tensor`):
+            The input tensor to the model.
+        output (`torch.Tensor`):
+            The output tensor of the model (i.e. the embeddings).
+    """
+    if module.training:
+        dims = torch.tensor(output.size(1) * output.size(2))
+        mag_norm = module.neftune_noise_alpha / torch.sqrt(dims)
+        output = output + torch.zeros_like(output).uniform_(-mag_norm, mag_norm)
+    return output
+
+
+def peft_module_casting_to_bf16(model):
+    from peft.tuners.tuners_utils import BaseTunerLayer
+
+    for name, module in model.named_modules():
+        if isinstance(module, BaseTunerLayer):
+            module = module.to(torch.bfloat16)
+        elif isinstance(module, torch.nn.LayerNorm) or "norm" in name:
+            module = module.to(torch.float32)
+        elif any(x in name for x in ["lm_head", "embed_tokens", "wte", "wpe"]):
+            if hasattr(module, "weight"):
+                if module.weight.dtype == torch.float32:
+                    module = module.to(torch.bfloat16)
+
+
+def trl_sanitze_kwargs_for_tagging(model, tag_names, kwargs=None):
+    if is_unsloth_available():
+        # Unsloth adds a new attribute in the model config `unsloth_version`
+        # to keep track of models that have been patched with unsloth.
+        if hasattr(model, "config") and getattr(model.config, "unsloth_version", None) is not None:
+            tag_names.append("unsloth")
+
+    if kwargs is not None:
+        if "tags" not in kwargs:
+            kwargs["tags"] = tag_names
+        elif "tags" in kwargs and isinstance(kwargs["tags"], list):
+            kwargs["tags"].extend(tag_names)
+        elif "tags" in kwargs and isinstance(kwargs["tags"], str):
+            tag_names.append(kwargs["tags"])
+            kwargs["tags"] = tag_names
+    return kwargs
+
+
+def get_quantization_config(model_config: ModelConfig) -> Optional[BitsAndBytesConfig]:
+    if model_config.load_in_4bit:
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=model_config.torch_dtype,  # For consistency with model weights, we use the same value as `torch_dtype`
+            bnb_4bit_quant_type=model_config.bnb_4bit_quant_type,
+            bnb_4bit_use_double_quant=model_config.use_bnb_nested_quant,
+        )
+    elif model_config.load_in_8bit:
+        quantization_config = BitsAndBytesConfig(
+            load_in_8bit=True,
+        )
+    else:
+        quantization_config = None
+
+    return quantization_config
+
+
+def get_kbit_device_map() -> Optional[Dict[str, int]]:
+    if is_xpu_available():
+        return {"": f"xpu:{PartialState().local_process_index}"}
+    elif torch.cuda.is_available():
+        return {"": PartialState().local_process_index}
+    else:
+        return None
+
+
+def get_peft_config(model_config: ModelConfig) -> "Optional[PeftConfig]":
+    if model_config.use_peft is False:
+        return None
+
+    peft_config = LoraConfig(
+        r=model_config.lora_r,
+        lora_alpha=model_config.lora_alpha,
+        lora_dropout=model_config.lora_dropout,
+        bias="none",
+        task_type="CAUSAL_LM",
+        target_modules=model_config.lora_target_modules,
+        modules_to_save=model_config.lora_modules_to_save,
+    )
+
+    return peft_config

plot.py

@@ -0,0 +1,103 @@
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt 
+import os
+import glob2 
+import seaborn as sns 
+
+
+def find_pareto_points(obtained_scores, threshold=0.02):
+    n = len(obtained_scores)
+    if n == 1:
+        return obtained_scores
+    pareto_index = []
+    high_low = np.max(obtained_scores, axis=0) - np.min(obtained_scores, axis=0)
+    for i in range(n):
+        if not any(np.all((obtained_scores - obtained_scores[i] - threshold * high_low) > 0.0, axis=1)):
+            pareto_index.append(i)
+
+    points = obtained_scores[np.array(pareto_index)]
+    arg_index = np.argsort(points[:, 0])
+    points = points[arg_index]
+    print(points)
+    sorted_index = [0]
+    remaining_index = np.ones(len(points))
+    i = 0
+    remaining_index[i] = 0
+    while sum(remaining_index):
+        distance = ((points[np.where(remaining_index)] - points[i]) ** 2 ).sum(axis=1)
+        min_index = np.where(remaining_index > 0)[0][np.argmin(distance)]
+        sorted_index.append(min_index)
+        i = min_index
+        remaining_index[i] = 0
+    return points[np.array(sorted_index)]
+
+
+
+index = 1
+colors = sns.color_palette('Paired')
+def plot_points(dir, label, style='-*', color='b', shift=[0,0], txt_color='black', normalize_path=None, reverse=True):
+    threshold = 0.01
+    desired_scores = []
+    obtained_scores = []
+
+    paths = [os.path.abspath(path) for path in glob2.glob(os.path.join(dir, '*.csv'))]
+    paths += [os.path.abspath(path) for path in glob2.glob(os.path.join(dir, '*', '*.csv'))]
+
+    pref_lis = []
+    for path in paths:
+        if '.csv' in path:
+            full_path = path 
+            data = pd.read_csv(full_path)
+            # morlhf has less points, let the threshold larger to make the frontier better
+            if 'ppo' in path and len(paths) <= 5:
+                threshold = 0.5
+            obtained_scores.append([np.mean(data['obtained_score1']), np.mean(data['obtained_score2'])])
+            if 'pref' in path:
+                # get the preference
+                if 'eval_data_pref' in path:
+                    pref = path.split('eval_data_pref')[-1].strip().split('_')[0]
+                    pref_lis.append(float(pref))
+
+    print(pref_lis)
+    desired_scores = np.array(desired_scores)
+    obtained_scores = np.array(obtained_scores)
+
+    if normalize_path is not None:
+        norm_info = np.load(normalize_path)
+        norm_info = np.array(norm_info).reshape(2, 2)
+        for i in range(2):
+            obtained_scores[:, i] = (obtained_scores[:, i] - norm_info[i][0]) / norm_info[i][1] 
+
+    global index
+    markersize = 10  if ('*' in style or 'o' in style) else 9
+    pareto_points = find_pareto_points(obtained_scores, threshold)
+    plt.scatter(obtained_scores[:, 0], obtained_scores[:, 1], marker=style[-1], color=colors[index], s=markersize + 60)
+    if len(pref_lis):
+        for i in range(len(obtained_scores)):
+            plt.annotate('{}'.format(round(pref_lis[i], 1)), (obtained_scores[i, 0] + shift[0], obtained_scores[i, 1] + shift[1]), size=4, color=txt_color)
+
+    plt.plot(pareto_points[:, 0], pareto_points[:, 1], style, c=colors[index], markersize=markersize, label=label)
+    index += 2
+
+
+plt.figure(figsize=(5, 4))
+
+name1 = 'harmless'
+name2 = 'helpful'
+
+### replace the paths to your own paths
+plot_points('./logs_trl/eval_pretrained', 'Llama 2 base', '*')
+plot_points('./logs_trl/eval_sft_alldata', 'SFT', '*')
+plot_points('./eval_ppo_pref/', 'MORLHF', '--D', shift=[-0.012, -0.022])
+plot_points('./logs_ppo/eval_pposoups_llamma2_klreg0.2', 'Rewarded Soups', style='--s', shift=[-0.012, -0.022])
+plot_points('.logs_trl/evalnew_onlinefix_helpful_harmlesshelpful_iter2',  'RiC', style='-o', shift=[-0.012, -0.022], txt_color='white')
+
+
+plt.xlabel('$R_1$ ({})'.format(name1), fontsize=12)
+plt.ylabel('$R_2$ ({})'.format(name2), fontsize=12)
+plt.legend(fontsize=11, loc='lower left')
+plt.tight_layout()
+plt.savefig('ric_assistant_{}_{}.pdf'.format(name1, name2))
+
+

ppo/eval_ppo_single_model.py

@@ -0,0 +1,177 @@
+import os
+from dataclasses import dataclass, field
+from typing import Optional
+from accelerate import Accelerator
+import torch
+from tqdm import tqdm
+from transformers import AutoModelForCausalLM, HfArgumentParser
+from transformers import DataCollatorWithPadding
+from trl import set_seed
+import numpy as np
+import pandas as pd
+from torch.utils.data import DataLoader
+from utils import load_main_tokenizer
+from multi_reward_models import RewardModels
+from utils import Instructions, Instructions_summary, build_dataset_eval, build_dataset_summary_eval, check_lora_in_model_path, merge_lora_weight
+tqdm.pandas()
+
+# define paths for two datasets
+hhrlhf_dataset_path = 'Anthropic/hh-rlhf'
+summary_dataset_path = 'openai/summarize_from_feedback'
+
+@dataclass
+class ScriptArguments:
+    log_with: Optional[str] = field(default='wandb', metadata={"help": "use 'wandb' to log with wandb"})
+    save_directory: Optional[str] = field(default='./logs_ppo_assistant')
+    batch_size: Optional[int] = field(default=1, metadata={"help": "the batch size"})
+    wandb_name: Optional[str] = field(default='eval_morlhf_kl0.2_multirew0.2_gen128_harmless_helpful', metadata={"help": "Name for this experiment"})
+    reward_names:Optional[str] = field(default='harmless,helpful') 
+    base_model_name: Optional[str] = field(default='./logs_ppo_prefnew/ppo_llamma2_klreg0.2_harmless_helpful_preference0.5_0.5/batch_207')
+    exp_type: Optional[str] = field(default='assistant', metadata={"help": "exp type, 'summary' or 'assistant' "})
+
+parser = HfArgumentParser(ScriptArguments)
+script_args = parser.parse_args_into_dataclasses()[0]
+exp_type = script_args.exp_type
+base_model_name = script_args.base_model_name
+tokenier_name = script_args.base_model_name
+print('base model: ', base_model_name)
+
+reward_names = [x.strip() for x in script_args.reward_names.split(',')]
+num_rewards = len(reward_names)
+print('number of rewards: {}'.format(num_rewards))
+reward_path_tokenizer_dict = {
+    'harmless': ['Ray2333/gpt2-large-harmless-reward_model'],
+    'helpful': ['Ray2333/gpt2-large-helpful-reward_model'],
+    'deberta': ['OpenAssistant/reward-model-deberta-v3-large-v2'],
+    'summary': ['Tristan/gpt2_reward_summarization'],
+    'faithful':['CogComp/bart-faithful-summary-detector'],
+    'humor': ['mohameddhiab/humor-no-humor'],
+}
+reward_model_path_list = []
+rm_tokenizer_path_list = []
+for name in reward_names:
+    if name not in reward_path_tokenizer_dict.keys():
+        raise NotImplementedError
+    reward_model_path_list.append(reward_path_tokenizer_dict[name][0])
+    rm_tokenizer_path_list.append(reward_path_tokenizer_dict[name][0])
+os.makedirs(os.path.join(script_args.save_directory, script_args.wandb_name), exist_ok=True)
+
+
+process_id = Accelerator().local_process_index 
+gpu_id = process_id
+print('process: {}, model gpu id: {}'.format(process_id, gpu_id))
+  
+
+##############
+rm_gpu_id = gpu_id
+reward_models = RewardModels(reward_model_path_list, rm_tokenizer_path_list, rm_gpu_id)
+
+set_seed(8888)
+current_device = Accelerator().local_process_index
+print(current_device)
+
+
+tokenizer = load_main_tokenizer(tokenier_name)
+model = AutoModelForCausalLM.from_pretrained(
+    base_model_name, 
+    torch_dtype=torch.bfloat16,  # fast inference
+    device_map=gpu_id, 
+)
+model.resize_token_embeddings(len(tokenizer))
+model = merge_lora_weight(model, base_model_name)
+
+generation_kwargs = {
+    "max_new_tokens": 128 if exp_type == 'assistant' else 48, 
+    "min_length": -1,
+    "top_k": 0.0,
+    "top_p": 0.9, 
+    "do_sample": True,
+}
+
+
+def get_clean_data(full_responses, full_prompts):
+    full_responses_clean = []
+    for i, response in enumerate(full_responses):
+        full_prompts[i] = full_prompts[i].strip('[PAD] ')
+        response = response.strip('[PAD] ')
+        temp_resp = response.replace(full_prompts[i], '').strip('<s>').strip('</s>')
+        if '</s>' in temp_resp:
+            temp_resp = temp_resp[:temp_resp.rindex('</s>')]
+        temp_resp = temp_resp.split('\n\nHuman:')[0].strip()
+        temp_resp = temp_resp.split('\nHuman:')[0].strip()
+        temp_resp = temp_resp.split('\n\nAssistant:')[0].strip()
+        temp_resp = temp_resp.split('\nAssistant:')[0].strip()
+        temp_resp = temp_resp.split('\n\n\n')[0].strip()
+        full_responses_clean.append(full_prompts[i] + ' ' + temp_resp)
+    return full_responses_clean
+
+
+### for evaluation
+print('evaluation........')
+tokenizer.padding_side = "left"
+if exp_type == 'assistant':
+    valid_dataset = build_dataset_eval(hhrlhf_dataset_path, tokenizer, reward_models.rm_tokenizers,  split='test')
+    instructions = Instructions()
+else:
+    valid_dataset = build_dataset_summary_eval(summary_dataset_path, tokenizer, reward_models.rm_tokenizers, split='test')
+    instructions = Instructions_summary()
+print(f"Size of the validation set: {len(valid_dataset)}")
+
+valid_batch_size = 1
+valid_dataset = valid_dataset.remove_columns(['prompt', 'query'])
+for key in ['key', 'text', 'response']:
+    if key in valid_dataset.column_names:
+        valid_dataset = valid_dataset.remove_columns(key)
+
+data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
+valid_data_loader = DataLoader(valid_dataset, batch_size=valid_batch_size, drop_last=True, collate_fn=data_collator)
+accelerator = Accelerator()
+model, valid_data_loader = accelerator.prepare(model, valid_data_loader)
+
+full_response_tensors = []
+full_prompts = []
+
+pbar = tqdm(total=len(valid_dataset) // valid_batch_size // accelerator.num_processes)
+with torch.no_grad():
+    for i, batch in enumerate(valid_data_loader):
+        response_tensors = accelerator.unwrap_model(model).generate(batch['input_ids'], attention_mask=batch['attention_mask'], **generation_kwargs)
+        full_response_tensors.extend(response_tensors)
+        full_prompts.extend(batch['input_ids'])
+        pbar.update(1)
+
+full_prompts = tokenizer.batch_decode(full_prompts)
+full_responses = tokenizer.batch_decode(full_response_tensors)
+full_responses = get_clean_data(full_responses, full_prompts)
+# Compute score
+queries_responses = [
+    (instructions.get_input(text),  instructions.get_response(text))
+    for text in full_responses
+]
+
+
+if hasattr(instructions, 'get_post'):
+    rewards_list = reward_models.get_reward_model_scores(queries_responses, instructions.get_post)
+else:
+    rewards_list = reward_models.get_reward_model_scores(queries_responses)
+
+all_rewards = []
+for i in range(reward_models.num_rewards):
+    all_rewards.append(accelerator.gather_for_metrics(rewards_list[i]))
+### merge data
+### error here may because of old version of transformers/accelerate/peft
+all_full_prompts = accelerator.gather_for_metrics(full_prompts)
+all_full_responses = accelerator.gather_for_metrics(full_responses)
+
+
+if process_id == 0:
+    evaluation_result = {
+        'prompt':all_full_prompts,
+        'response': all_full_responses,
+    }
+    for i in range(reward_models.num_rewards):
+        evaluation_result['obtained_score{}'.format(i+1)] = all_rewards[i]
+        print('total average obtained score {}: {}'.format(i+1, np.mean(evaluation_result['obtained_score{}'.format(i+1)])))
+
+    dataframe = pd.DataFrame(evaluation_result)
+    dataframe.to_csv(os.path.join(script_args.save_directory, script_args.wandb_name,'eval_data.csv'))
+

ppo/eval_rewarded_soups.py

@@ -0,0 +1,220 @@
+import os
+import gc
+from dataclasses import dataclass, field
+from typing import Optional
+from accelerate import Accelerator
+import torch
+from datasets import load_dataset
+from tqdm import tqdm
+from transformers import HfArgumentParser
+from transformers import AutoModelForCausalLM, DataCollatorWithPadding
+from torch.utils.data import DataLoader
+from trl import set_seed
+import numpy as np
+import pandas as pd
+from utils import get_clean_data, load_main_tokenizer, save_configs, print_trainable_parameters, \
+                  merge_weights_with_preference, Instructions, Instructions_summary, build_dataset_eval, build_dataset_summary_eval
+from multi_reward_models import RewardModels
+tqdm.pandas()
+
+
+# define paths for two datasets
+hhrlhf_dataset_path = 'Anthropic/hh-rlhf'
+summary_dataset_path = 'openai/summarize_from_feedback'
+tokenizer_path = 'meta-llama/Llama-2-7b-hf'
+
+@dataclass
+class ScriptArguments:
+    save_directory: Optional[str] = field(default='./logs_rewardedsoups_summary_eval')
+    mini_batch_size: Optional[int] = field(default=1, metadata={"help": "minibatch size for eval"})
+    wandb_name: Optional[str] = field(default='eval_pposoups_klreg0.2_harmless_helpful', metadata={"help": "Name for this experiment"})
+    reward_names:Optional[str] = field(default='harmless,helpful')
+    base_model_path1: Optional[str]=field(default='./ppo_llamma2_klreg0.2_harmless/batch_200')
+    base_model_path2: Optional[str]=field(default='./ppo_llamma2_klreg0.2_helpful/batch_200')
+    base_model_path3: Optional[str]=field(default='')
+    exp_type: Optional[str] = field(default='assistant', metadata={"help": "exp type, 'summary' or 'assistant' "})
+
+parser = HfArgumentParser(ScriptArguments)
+script_args = parser.parse_args_into_dataclasses()[0]
+exp_type = script_args.exp_type
+base_model_name_1 = script_args.base_model_path1
+base_model_name_2 = script_args.base_model_path2
+base_model_name_3 = script_args.base_model_path3
+tokenier_name = tokenizer_path
+
+reward_names = [x.strip() for x in script_args.reward_names.split(',')]
+reward_path_tokenizer_dict = {
+    'harmless': ['Ray2333/gpt2-large-harmless-reward_model'],
+    'helpful': ['Ray2333/gpt2-large-helpful-reward_model'],
+    'deberta': ['OpenAssistant/reward-model-deberta-v3-large-v2'],
+    'summary': ['Tristan/gpt2_reward_summarization'],
+    'faithful':['CogComp/bart-faithful-summary-detector'],
+    'humor': ['mohameddhiab/humor-no-humor'],
+}
+reward_model_path_list = []
+rm_tokenizer_path_list = []
+for name in reward_names:
+    if name not in reward_path_tokenizer_dict.keys():
+        raise NotImplementedError
+    reward_model_path_list.append(reward_path_tokenizer_dict[name][0])
+    rm_tokenizer_path_list.append(reward_path_tokenizer_dict[name][0])
+os.makedirs(os.path.join(script_args.save_directory, script_args.wandb_name), exist_ok=True)
+save_info = {
+    'base_model_name_1': base_model_name_1,
+    'base_model_name_2': base_model_name_2,
+    'base_model_name_3': base_model_name_3,
+    'reward_peft_path1': reward_model_path_list[0],
+    'reward_peft_path2': reward_model_path_list[1],
+    'tokenier_name': tokenier_name
+}
+for i in range(len(reward_model_path_list)):
+    save_info['reward_peft_path{}'.format(i+1)] = reward_model_path_list[i]
+save_configs(save_info, os.path.join(script_args.save_directory, script_args.wandb_name))
+
+
+accelerator = Accelerator()
+process_id = Accelerator().local_process_index 
+gpu_id = process_id
+print('process: {}, model gpu id: {}'.format(process_id, gpu_id))
+reward_models = RewardModels(reward_model_path_list, rm_tokenizer_path_list, gpu_id) 
+
+
+set_seed(8888)
+current_device = Accelerator().local_process_index
+print(current_device)
+
+
+tokenizer = load_main_tokenizer(tokenier_name)
+if exp_type == 'assistant':
+    valid_dataset = build_dataset_eval(hhrlhf_dataset_path, tokenizer, reward_models.rm_tokenizers, split='test') 
+    instructions = Instructions()
+else:
+    valid_dataset = build_dataset_summary_eval(summary_dataset_path, tokenizer, reward_models.rm_tokenizers, split='test')
+    instructions = Instructions_summary()
+print(f"Size of the validation set: {len(valid_dataset)}")
+
+
+valid_dataset = valid_dataset.remove_columns(['prompt', 'query'])
+for key in ['key', 'text', 'response']:
+    if key in valid_dataset.column_names:
+        valid_dataset = valid_dataset.remove_columns(key)
+
+def evaluate_model(temp_save_path, tokenizer, valid_dataset):
+    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
+    valid_data_loader = DataLoader(valid_dataset, batch_size=script_args.mini_batch_size, drop_last=True, collate_fn=data_collator)
+    ### load_merged_model
+    model = AutoModelForCausalLM.from_pretrained(
+        temp_save_path,
+        torch_dtype=torch.bfloat16, # for merging model
+        device_map=gpu_id
+        )
+    model.resize_token_embeddings(len(tokenizer))
+
+    accelerator = Accelerator()
+    model, valid_data_loader = accelerator.prepare(model, valid_data_loader)
+
+    generation_kwargs = {
+        "max_new_tokens": 128 if exp_type == 'assistant' else 48, 
+        "min_length": -1,
+        "top_k": 0.0,
+        "top_p": 0.9, 
+        "do_sample": True,
+    }
+
+    full_responses = []
+    full_prompts = []
+    pbar = tqdm(total=len(valid_dataset) // script_args.mini_batch_size // accelerator.num_processes)
+    with torch.no_grad():
+        for i, batch in enumerate(valid_data_loader):
+            response_tensors = accelerator.unwrap_model(model).generate(batch["input_ids"], **generation_kwargs) #length_sampler=output_length_sampler, 
+            full_responses.extend(response_tensors)
+            full_prompts.extend(batch['input_ids'])
+            pbar.update(1)
+    
+    full_responses = tokenizer.batch_decode(full_responses)
+    full_prompts = tokenizer.batch_decode(full_prompts)
+    # clean data
+    full_prompts, full_responses = get_clean_data(full_responses, full_prompts)
+
+    queries_responses = [
+        (instructions.get_input(text), instructions.get_response(text))
+        for text in full_responses
+    ]
+    if hasattr(instructions, 'get_post'):
+        rewards_list = reward_models.get_reward_model_scores(queries_responses, instructions.get_post)
+    else:
+        rewards_list = reward_models.get_reward_model_scores(queries_responses)
+    
+    ### error here may because of old version of transformers/accelerate/peft
+    all_rewards = []
+    for i in range(reward_models.num_rewards):
+        all_rewards.append(accelerator.gather_for_metrics(rewards_list[i]))
+    all_full_prompts = accelerator.gather_for_metrics(full_prompts)
+    all_full_responses = accelerator.gather_for_metrics(full_responses)
+    return all_rewards, all_full_prompts, all_full_responses
+
+
+print("Evaluating........")
+tokenizer.padding_side = "left"
+## preference list
+if reward_models.num_rewards == 3:
+    preferences = np.array([
+        [0.0, 0.0, 1.0],
+        [0.0, 1.0, 0.0],
+        [0.1, 0.1, 0.8],
+        [0.1, 0.8, 0.1],
+        [0.2, 0.2, 0.6],
+        [0.2, 0.4, 0.4],
+        [0.2, 0.6, 0.2],
+        [0.33, 0.33, 0.33],
+        [0.4, 0.4, 0.2],
+        [0.4, 0.2, 0.4], 
+        [0.6, 0.2, 0.2],
+        [0.8, 0.1, 0.1], 
+        [1.0, 0.0, 0.0], 
+        ])
+elif reward_models.num_rewards == 2:
+    M = 10
+    preferences = np.zeros((M+1, 2))
+    preferences[:, 0] = np.arange(0,1+ 1/M,1 / M)
+    preferences[:, 1] = 1 - preferences[:, 0]
+    preferences = np.round(preferences, 1)
+else:
+    raise NotImplementedError
+
+for k in range(0, len(preferences)):
+    preference = preferences[k]
+    temp_save_path = './temp_models/{}'.format(script_args.wandb_name)
+    if process_id == 0:
+        if len(preference) == 3:
+            base_model_list = [base_model_name_1, base_model_name_2, base_model_name_3]
+        else:
+            base_model_list = [base_model_name_1, base_model_name_2]
+        merge_weights_with_preference(base_model_list, preference, temp_save_path)
+
+    accelerator.wait_for_everyone()
+    gc.collect()
+    torch.cuda.empty_cache()
+    print(k, preference)
+    all_rewards, all_full_prompts, all_full_responses = evaluate_model(temp_save_path, tokenizer, valid_dataset)
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    if process_id == 0:
+        evaluation_result = {
+            'prompt': all_full_prompts,
+            'response': all_full_responses,
+        }
+        for i in range(reward_models.num_rewards):
+            evaluation_result['obtained_score{}'.format(i+1)] = all_rewards[i]
+            print('total average obtained score {}: {}'.format(i+1, np.mean(evaluation_result['obtained_score{}'.format(i+1)])))
+
+        dataframe = pd.DataFrame(evaluation_result)
+        if len(preference) == 2:
+            dataframe.to_csv(os.path.join(script_args.save_directory, script_args.wandb_name,'eval_data_pref{}_{}.csv'.format(preference[0], preference[1])))
+        else:
+            dataframe.to_csv(os.path.join(script_args.save_directory, script_args.wandb_name,'eval_data_pref{}_{}_{}.csv'.format(preference[0], preference[1], preference[2])))
+
+
+
+

ppo/morlhf-llama3.py

@@ -0,0 +1,293 @@
+import os
+import time
+from dataclasses import dataclass, field
+from typing import Optional
+from accelerate import Accelerator
+import torch
+from tqdm import tqdm
+from transformers import AutoTokenizer, HfArgumentParser
+from trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer, set_seed
+import numpy as np
+import pandas as pd
+from utils import print_trainable_parameters, Instructions, Instructions_summary, \
+                 load_llama3_tokenizer, build_dataset_llama3    
+from multi_reward_models import RewardModels
+tqdm.pandas()
+from peft import LoraConfig
+import matplotlib.pyplot as plt
+
+# define paths for two datasets
+hhrlhf_dataset_path = '../datasets/hh-rlhf'
+
+@dataclass
+class ScriptArguments:
+    log_with: Optional[str] = field(default='wandb', metadata={"help": "use 'wandb' to log with wandb"})
+    disable_wandb: Optional[str] = field(default=False, metadata={'help': 'Whether to disable wandb or not.'})
+    save_directory: Optional[str] = field(default='./logs_morlhf/')
+    epochs: Optional[int] = field(default=1, metadata={'help': "Number of training epoches"})
+    learning_rate: Optional[float] = field(default=1e-5, metadata={"help": "the learning rate"})
+    mini_batch_size: Optional[int] = field(default=1, metadata={"help": "the PPO minibatch size"})
+    batch_size: Optional[int] = field(default=64, metadata={"help": "the batch size64"})
+    gradient_accumulation_steps: Optional[int] = field(default=1, metadata={"help": "the number of gradient accumulation steps"})
+    early_stopping: Optional[bool] = field(default=True, metadata={"help": "whether to early stop"})
+    target: Optional[float] = field(default=3, metadata={"help": "target kl divergence of adaptive control"})
+    init_kl_coef: Optional[float] = field(default=0.2,metadata={"help": "0.05 Initial KL penalty coefficient (used for adaptive and linear control)"},)
+    max_grad_norm: Optional[float] = field(default=0.5, metadata={"help": "Maximum gradient norm for gradient clipping"})
+    load_in_8bit: Optional[bool] = field(default=False, metadata={"help": "loading model in 8 bit or bfloat16"})
+    preference: Optional[str] = field(default=None)
+    wandb_name: Optional[str] = field(default='morlhf_llamma2_klreg0.2', metadata={"help": "Name for this experiment"})
+    base_model_name: Optional[str] = field(default='./merged_sft_summary', metadata={'help':"the path to the sft model; need to merge if using lora"})
+    reward_names:Optional[str] = field(default='harmless,helpful,humor') 
+    exp_type: Optional[str] = field(default='assistant', metadata={"help": "exp type, 'summary' or 'assistant' "})
+
+parser = HfArgumentParser(ScriptArguments)
+script_args = parser.parse_args_into_dataclasses()[0]
+exp_type = script_args.exp_type
+
+tokenier_name = script_args.base_model_name
+base_model_name = script_args.base_model_name
+print('base model: ', base_model_name)
+
+if script_args.disable_wandb: # if you don't need the wandb log
+    os.environ['WANDB_DISABLED'] = 'true' 
+
+reward_names = [x.strip() for x in script_args.reward_names.split(',')]
+num_rewards = len(reward_names)
+print('number of rewards: {}'.format(num_rewards))
+#######################################################################
+assert num_rewards == 3
+if script_args.preference is None:
+    # 均分preference
+    preference = [round(1 / num_rewards, 2) for _ in range(num_rewards)]
+else:
+    preference = [float(x) for x in script_args.preference.split(",")]
+print('preference: {}'.format(preference))
+#######################################################################
+reward_path_tokenizer_dict = {
+    'harmless': ['../reward_models/gpt2-large-harmless-reward_model'],
+    'helpful': ['../reward_models/gpt2-large-helpful-reward_model'],
+    'humor': ['../reward_models/humor-no-humor'],
+}
+reward_model_path_list = []
+rm_tokenizer_path_list = []
+for name in reward_names:
+    if name not in reward_path_tokenizer_dict.keys():
+        raise NotImplementedError
+    reward_model_path_list.append(reward_path_tokenizer_dict[name][0])
+    rm_tokenizer_path_list.append(reward_path_tokenizer_dict[name][0])
+os.makedirs(os.path.join(script_args.save_directory, script_args.wandb_name), exist_ok=True)
+
+
+config = PPOConfig(
+    model_name=base_model_name,
+    learning_rate=script_args.learning_rate,
+    # log_with=script_args.log_with,
+    mini_batch_size=script_args.mini_batch_size,
+    batch_size=script_args.batch_size,
+    gradient_accumulation_steps=script_args.gradient_accumulation_steps,
+    early_stopping=script_args.early_stopping,
+    target=script_args.target,
+    max_grad_norm=script_args.max_grad_norm,
+    optimize_cuda_cache=True,
+    init_kl_coef=script_args.init_kl_coef,
+    # tracker_project_name='morlhf',
+    # tracker_kwargs={"wandb":{"name":script_args.wandb_name}},
+    remove_unused_columns=False # 防止‘reward_query’没了
+)
+
+accelerator = Accelerator()
+process_id = Accelerator().local_process_index 
+gpu_id = process_id
+print('process: {}, model gpu id: {}'.format(process_id, gpu_id))
+
+
+############## load reward models
+reward_model = RewardModels(reward_model_path_list, rm_tokenizer_path_list, gpu_id)
+rm_tokenizer = AutoTokenizer.from_pretrained(rm_tokenizer_path_list[0])
+
+def collator(data):
+    return dict((key, [d[key] for d in data]) for key in data[0])
+
+set_seed(8888)
+current_device = Accelerator().local_process_index
+print(current_device)
+
+lora_config = LoraConfig(
+    r=64, 
+    lora_alpha=128, 
+    lora_dropout=0.05,
+    bias="none",
+    task_type="CAUSAL_LM",
+)
+
+
+tokenizer = load_llama3_tokenizer(tokenier_name)
+if exp_type == 'assistant':
+    dataset = build_dataset_llama3(hhrlhf_dataset_path, tokenizer, rm_tokenizer, split='train')
+    # instructions = Instructions()
+else:
+    raise NotImplementedError
+    # dataset = build_dataset_summary(summary_dataset_path, tokenizer, rm_tokenizer, split='train')
+    # instructions = Instructions_summary()
+train_dataset = dataset.shuffle()
+print(f"Size of the train set: {len(train_dataset)}")
+
+
+if script_args.load_in_8bit:
+    model = AutoModelForCausalLMWithValueHead.from_pretrained(
+        base_model_name,
+        load_in_8bit=True,
+        peft_config=lora_config,
+        device_map=gpu_id,
+    )
+else:
+    model = AutoModelForCausalLMWithValueHead.from_pretrained(
+        base_model_name,
+        torch_dtype=torch.bfloat16,
+        peft_config=lora_config,
+        device_map=gpu_id,
+    )
+
+print_trainable_parameters(model)
+model.pretrained_model.resize_token_embeddings(len(tokenizer))
+optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=config.learning_rate)
+
+ppo_trainer = PPOTrainer(
+    config, model, tokenizer=tokenizer, dataset=dataset, data_collator=collator, optimizer=optimizer
+)
+
+# 只进行一轮生成
+terminators = [
+    tokenizer.eos_token_id,
+    tokenizer.convert_tokens_to_ids("<|eot_id|>")
+]
+generation_kwargs = {
+    "max_new_tokens": 128 if exp_type == 'assistant' else 48,
+    "min_length": -1,
+    "top_k": 0.0,
+    "top_p": 1.0, 
+    "do_sample": True,
+    "temperature": 0.7,
+    "pad_token_id": tokenizer.eos_token_id,
+    "begin_suppress_tokens": [tokenizer.eos_token_id] ,
+    'eos_token_id': terminators
+}
+
+
+print("Training........")
+model.gradient_checkpointing_disable()
+model.pretrained_model.config.use_cache = True
+epochs = script_args.epochs
+mean_scores = []
+std_scores = []
+save_data = {
+    'kl_mean': [],
+    'reward_mean': [],
+    'reward_std': [],
+    'text_sample':[],
+    'batch_time':[],
+    'total_time':[],
+}
+t_start = time.time()
+for epoch in range(epochs):
+    pbar = tqdm(total=len(train_dataset) // script_args.batch_size // accelerator.num_processes)
+    for i, batch in enumerate(ppo_trainer.dataloader):
+        # print(batch.keys())
+        t_epoch_start = time.time()
+        print('epoch {}, batch {}'.format(epoch, i))
+        query_tensors = batch["input_ids"]
+
+        model.gradient_checkpointing_disable()
+        model.pretrained_model.config.use_cache = True
+            
+        with torch.no_grad():
+            response_tensors = ppo_trainer.generate(query_tensors, return_prompt=False, **generation_kwargs)
+
+        # 看ppo_trainer中的实现，pad token已经被去除干净了
+        full_responses_clean = tokenizer.batch_decode(response_tensors, skip_special_tokens=True)  # skip后，不用手动去除
+        
+        # full_responses = tokenizer.batch_decode(response_tensors) 
+        # full_responses_clean = []
+        # for _, response in enumerate(full_responses):
+        #     response = response.strip('[PAD] ')
+        #     response = response.strip('<unk>')
+        #     temp_resp = response.strip('<s>').strip('</s>')
+        #     temp_resp = temp_resp.split('\n\nHuman:')[0].strip()
+        #     temp_resp = temp_resp.split('\nHuman:')[0].strip()
+        #     temp_resp = temp_resp.split('\n\nAssistant:')[0].strip()
+        #     temp_resp = temp_resp.split('\nAssistant:')[0].strip()
+        #     temp_resp = temp_resp.split('\n\n\n')[0].strip()
+        #     temp_resp = temp_resp.split('###')[0].strip()
+        #     full_responses_clean.append(temp_resp)
+
+        clean_texts = full_responses_clean
+        clean_response_tensors = [tokenizer.encode(text) for text in clean_texts]
+        
+        # clean后的response map回去
+        lengths = [len(clean_response_tensors[j]) for j in range(len(clean_response_tensors))]
+        response_tensors = [response_tensors[j][:np.max([lengths[j], 2])] for j in range(len(response_tensors))]
+        batch['response'] = clean_texts
+
+        # Compute score
+        # texts_merge = [q + r for q, r in zip(batch['query'], batch['response'])]
+        # # 从整个回复中分解prompt和response
+        # queries_responses = [
+        #     (instructions.get_input(text), instructions.get_response(text))
+        #     for text in texts_merge
+        # ]
+        # if hasattr(instructions, 'get_post'):
+        #     rewards_list = reward_model.get_reward_model_scores(queries_responses, instructions.get_post)
+        # else:
+        #     rewards_list = reward_model.get_reward_model_scores(queries_responses)
+        
+        qa_lsit = [(q, r) for q, r in zip(batch['reward_query'], batch['response'])]
+        rewards_list = reward_model.get_reward_model_scores(qa_lsit)
+        rewards = []
+        for j in range(len(qa_lsit)):
+            rewards.append(round(sum([preference[k] * rewards_list[k][j] for k in range(num_rewards)]), 2))
+        rewards_tensor = [torch.tensor(r).to(gpu_id) for r in rewards]
+        print("iter {}, batch {}, mean score: {}".format(epoch, i, torch.mean(torch.tensor(rewards)).item()))
+
+        model.gradient_checkpointing_enable()
+        model.pretrained_model.config.use_cache = False
+        stats = ppo_trainer.step(query_tensors, response_tensors, rewards_tensor)
+        policy_kl = [stats["objective/kl"]]
+        ppo_trainer.log_stats(stats, batch, rewards)
+
+        all_rewards = accelerator.gather_for_metrics(rewards)
+        all_policy_kl = accelerator.gather_for_metrics(policy_kl)
+        if process_id == 0:
+            mean_scores.append(torch.mean(torch.tensor(all_rewards)).item())
+            std_scores.append(torch.std(torch.tensor(all_rewards)).item())
+            save_path = os.path.join(script_args.save_directory, script_args.wandb_name, 'scores.png')
+            plt.plot(mean_scores)
+            plt.fill_between(np.arange(len(mean_scores)), np.array(mean_scores)- np.array(std_scores), np.array(mean_scores) + np.array(std_scores), alpha=0.5)
+            plt.savefig(save_path)
+            t_epoch_end = time.time()
+            save_data['batch_time'].append(t_epoch_end - t_epoch_start)
+            save_data['total_time'].append(t_epoch_end - t_start)
+            save_data['kl_mean'].append(np.mean(all_policy_kl))
+            save_data['reward_mean'] = mean_scores
+            save_data['reward_std'] = std_scores
+            save_data['text_sample'].append(batch['query'][0] + batch['response'][0])
+            # save_data['text_sample'].append(texts_merge[0])
+            dataframe = pd.DataFrame(save_data)
+            dataframe.to_csv(os.path.join(script_args.save_directory, script_args.wandb_name,'data.csv'))
+            print("iter {}, batch {}: log finish".format(epoch, i))
+
+        # wait for the main process
+        accelerator.wait_for_everyone()
+        pbar.update(1)
+
+        # save model
+        if ppo_trainer.accelerator.is_main_process and i % 100 == 0 and i != 0:
+            save_path = os.path.join(script_args.save_directory, script_args.wandb_name, 'epoch_{}-batch_{}'.format(epoch, i))
+            ppo_trainer.save_pretrained(save_path)
+            print("iter {}, batch {}: model saved".format(epoch, i))
+    
+    # save model
+    if ppo_trainer.accelerator.is_main_process:
+        save_path = os.path.join(script_args.save_directory, script_args.wandb_name, 'epoch_{}'.format(epoch))
+        ppo_trainer.save_pretrained(save_path)
+        print("iter {}, batch {}: model saved".format(epoch, i))
+        

ppo/morlhf.py

@@ -0,0 +1,277 @@
+import os
+import time
+from dataclasses import dataclass, field
+from typing import Optional
+from accelerate import Accelerator
+import torch
+from tqdm import tqdm
+from transformers import AutoTokenizer, HfArgumentParser
+from trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer, set_seed
+import numpy as np
+import pandas as pd
+from utils import print_trainable_parameters, load_main_tokenizer, Instructions, Instructions_summary, \
+                  build_dataset, build_dataset_summary             
+from multi_reward_models import RewardModels
+tqdm.pandas()
+from peft import LoraConfig
+import matplotlib.pyplot as plt
+
+# define paths for two datasets
+hhrlhf_dataset_path = 'Anthropic/hh-rlhf'
+summary_dataset_path = 'openai/summarize_from_feedback'
+
+@dataclass
+class ScriptArguments:
+    log_with: Optional[str] = field(default='wandb', metadata={"help": "use 'wandb' to log with wandb"})
+    disable_wandb: Optional[str] = field(default=False, metadata={'help': 'Whether to disable wandb or not.'})
+    save_directory: Optional[str] = field(default='./logs_morlhf/')
+    epochs: Optional[int] = field(default=1, metadata={'help': "Number of training epoches"})
+    learning_rate: Optional[float] = field(default=1e-5, metadata={"help": "the learning rate"})
+    mini_batch_size: Optional[int] = field(default=1, metadata={"help": "the PPO minibatch size"})
+    batch_size: Optional[int] = field(default=64, metadata={"help": "the batch size64"})
+    gradient_accumulation_steps: Optional[int] = field(default=1, metadata={"help": "the number of gradient accumulation steps"})
+    early_stopping: Optional[bool] = field(default=True, metadata={"help": "whether to early stop"})
+    target: Optional[float] = field(default=3, metadata={"help": "target kl divergence of adaptive control"})
+    init_kl_coef: Optional[float] = field(default=0.2,metadata={"help": "0.05 Initial KL penalty coefficient (used for adaptive and linear control)"},)
+    max_grad_norm: Optional[float] = field(default=0.5, metadata={"help": "Maximum gradient norm for gradient clipping"})
+    load_in_8bit: Optional[bool] = field(default=True, metadata={"help": "loading model in 8 bit or bfloat16"})
+    preference: Optional[float] = field(default=0.5, metadata={"help": "the weight for reward 1"})
+    wandb_name: Optional[str] = field(default='morlhf_llamma2_klreg0.2', metadata={"help": "Name for this experiment"})
+    base_model_name: Optional[str] = field(default='./merged_sft_summary', metadata={'help':"the path to the sft model; need to merge if using lora"})
+    reward_names:Optional[str] = field(default='harmless,helpful,humor') 
+    exp_type: Optional[str] = field(default='assistant', metadata={"help": "exp type, 'summary' or 'assistant' "})
+
+parser = HfArgumentParser(ScriptArguments)
+script_args = parser.parse_args_into_dataclasses()[0]
+exp_type = script_args.exp_type
+preference = [round(script_args.preference, 1), round(1 - script_args.preference, 1)]
+script_args.wandb_name = script_args.wandb_name + '_pref{}_{}'.format(preference[0], preference[1])
+
+tokenier_name = script_args.base_model_name
+base_model_name = script_args.base_model_name
+print('base model: ', base_model_name)
+
+if script_args.disable_wandb: # if you don't need the wandb log
+    os.environ['WANDB_DISABLED'] = 'true' 
+
+reward_names = [x.strip() for x in script_args.reward_names.split(',')]
+num_rewards = len(reward_names)
+print('number of rewards: {}'.format(num_rewards))
+#######################################################################
+if num_rewards == 3:
+    preference = [round(1 / num_rewards, 2) for _ in range(num_rewards)]
+print('preference: {}'.format(preference))
+#######################################################################
+reward_path_tokenizer_dict = {
+    'harmless': ['Ray2333/gpt2-large-harmless-reward_model'],
+    'helpful': ['Ray2333/gpt2-large-helpful-reward_model'],
+    'deberta': ['OpenAssistant/reward-model-deberta-v3-large-v2'],
+    'summary': ['Tristan/gpt2_reward_summarization'],
+    'faithful':['CogComp/bart-faithful-summary-detector'],
+    'humor': ['mohameddhiab/humor-no-humor'],
+}
+reward_model_path_list = []
+rm_tokenizer_path_list = []
+for name in reward_names:
+    if name not in reward_path_tokenizer_dict.keys():
+        raise NotImplementedError
+    reward_model_path_list.append(reward_path_tokenizer_dict[name][0])
+    rm_tokenizer_path_list.append(reward_path_tokenizer_dict[name][0])
+os.makedirs(os.path.join(script_args.save_directory, script_args.wandb_name), exist_ok=True)
+
+
+config = PPOConfig(
+    model_name=base_model_name,
+    learning_rate=script_args.learning_rate,
+    log_with=script_args.log_with,
+    mini_batch_size=script_args.mini_batch_size,
+    batch_size=script_args.batch_size,
+    gradient_accumulation_steps=script_args.gradient_accumulation_steps,
+    early_stopping=script_args.early_stopping,
+    target=script_args.target,
+    max_grad_norm=script_args.max_grad_norm,
+    optimize_cuda_cache=True,
+    init_kl_coef=script_args.init_kl_coef,
+    tracker_project_name='morlhf',
+    tracker_kwargs={"wandb":{"name":script_args.wandb_name}},
+)
+
+accelerator = Accelerator()
+process_id = Accelerator().local_process_index 
+gpu_id = process_id
+print('process: {}, model gpu id: {}'.format(process_id, gpu_id))
+
+
+############## load reward models
+reward_model = RewardModels(reward_model_path_list, rm_tokenizer_path_list, gpu_id)
+rm_tokenizer = AutoTokenizer.from_pretrained(rm_tokenizer_path_list[0])
+
+def collator(data):
+    return dict((key, [d[key] for d in data]) for key in data[0])
+
+set_seed(8888)
+current_device = Accelerator().local_process_index
+print(current_device)
+
+lora_config = LoraConfig(
+    r=64, 
+    lora_alpha=128, 
+    lora_dropout=0.05,
+    bias="none",
+    task_type="CAUSAL_LM",
+)
+
+
+tokenizer = load_main_tokenizer(tokenier_name)
+if exp_type == 'assistant':
+    dataset = build_dataset(hhrlhf_dataset_path, tokenizer, rm_tokenizer, split='train')
+    instructions = Instructions()
+else:
+    dataset = build_dataset_summary(summary_dataset_path, tokenizer, rm_tokenizer, split='train')
+    instructions = Instructions_summary()
+train_dataset = dataset.shuffle()
+print(f"Size of the train set: {len(train_dataset)}")
+
+
+if script_args.load_in_8bit:
+    model = AutoModelForCausalLMWithValueHead.from_pretrained(
+        base_model_name,
+        load_in_8bit=True,
+        peft_config=lora_config,
+        device_map=gpu_id,
+    )
+else:
+    model = AutoModelForCausalLMWithValueHead.from_pretrained(
+        base_model_name,
+        torch_dtype=torch.bfloat16,
+        peft_config=lora_config,
+        device_map=gpu_id,
+    )
+
+print_trainable_parameters(model)
+model.pretrained_model.resize_token_embeddings(len(tokenizer))
+optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=config.learning_rate)
+
+ppo_trainer = PPOTrainer(
+    config, model, tokenizer=tokenizer, dataset=dataset, data_collator=collator, optimizer=optimizer
+)
+
+generation_kwargs = {
+    "max_new_tokens": 128 if exp_type == 'assistant' else 48,
+    "min_length": -1,
+    "top_k": 0.0,
+    "top_p": 1.0, 
+    "do_sample": True,
+    "temperature": 0.7,
+    "pad_token_id": tokenizer.eos_token_id,
+    "begin_suppress_tokens": [tokenizer.eos_token_id] ,
+}
+
+
+print("Training........")
+model.gradient_checkpointing_disable()
+model.pretrained_model.config.use_cache = True
+epochs = script_args.epochs
+mean_scores = []
+std_scores = []
+save_data = {
+    'kl_mean': [],
+    'reward_mean': [],
+    'reward_std': [],
+    'text_sample':[],
+    'batch_time':[],
+    'total_time':[],
+}
+t_start = time.time()
+for epoch in range(epochs):
+    pbar = tqdm(total=len(train_dataset) // script_args.batch_size // accelerator.num_processes)
+    for i, batch in enumerate(ppo_trainer.dataloader):
+        t_epoch_start = time.time()
+        print('epoch {}, batch {}'.format(epoch, i))
+        query_tensors = batch["input_ids"]
+
+        model.gradient_checkpointing_disable()
+        model.pretrained_model.config.use_cache = True
+            
+        with torch.no_grad():
+            response_tensors = ppo_trainer.generate(query_tensors, return_prompt=False, **generation_kwargs)
+
+        full_responses = tokenizer.batch_decode(response_tensors)
+        full_responses_clean = []
+        for _, response in enumerate(full_responses):
+            response = response.strip('[PAD] ')
+            response = response.strip('<unk>')
+            temp_resp = response.strip('<s>').strip('</s>')
+            temp_resp = temp_resp.split('\n\nHuman:')[0].strip()
+            temp_resp = temp_resp.split('\nHuman:')[0].strip()
+            temp_resp = temp_resp.split('\n\nAssistant:')[0].strip()
+            temp_resp = temp_resp.split('\nAssistant:')[0].strip()
+            temp_resp = temp_resp.split('\n\n\n')[0].strip()
+            temp_resp = temp_resp.split('###')[0].strip()
+            full_responses_clean.append(temp_resp)
+
+        clean_texts = full_responses_clean
+        clean_response_tensors = [tokenizer.encode(text) for text in clean_texts]
+        
+        lengths = [len(clean_response_tensors[j]) for j in range(len(clean_response_tensors))]
+        response_tensors = [response_tensors[j][:np.max([lengths[j], 2])] for j in range(len(response_tensors))]
+        batch['response'] = clean_texts
+
+        # Compute score
+        texts_merge = [q + r for q, r in zip(batch['query'], batch['response'])]
+        queries_responses = [
+            (instructions.get_input(text), instructions.get_response(text))
+            for text in texts_merge
+        ]
+        if hasattr(instructions, 'get_post'):
+            rewards_list = reward_model.get_reward_model_scores(queries_responses, instructions.get_post)
+        else:
+            rewards_list = reward_model.get_reward_model_scores(queries_responses)
+        rewards = []
+        for j in range(len(queries_responses)):
+            rewards.append(round(sum([preference[k] * rewards_list[k][j] for k in range(num_rewards)]), 2))
+        rewards_tensor = [torch.tensor(r).to(gpu_id) for r in rewards]
+        print("iter {}, batch {}, mean score: {}".format(epoch, i, torch.mean(torch.tensor(rewards)).item()))
+
+        model.gradient_checkpointing_enable()
+        model.pretrained_model.config.use_cache = False
+        stats = ppo_trainer.step(query_tensors, response_tensors, rewards_tensor)
+        policy_kl = [stats["objective/kl"]]
+        ppo_trainer.log_stats(stats, batch, rewards)
+
+        all_rewards = accelerator.gather_for_metrics(rewards)
+        all_policy_kl = accelerator.gather_for_metrics(policy_kl)
+        if process_id == 0:
+            mean_scores.append(torch.mean(torch.tensor(all_rewards)).item())
+            std_scores.append(torch.std(torch.tensor(all_rewards)).item())
+            save_path = os.path.join(script_args.save_directory, script_args.wandb_name, 'scores.png')
+            plt.plot(mean_scores)
+            plt.fill_between(np.arange(len(mean_scores)), np.array(mean_scores)- np.array(std_scores), np.array(mean_scores) + np.array(std_scores), alpha=0.5)
+            plt.savefig(save_path)
+            t_epoch_end = time.time()
+            save_data['batch_time'].append(t_epoch_end - t_epoch_start)
+            save_data['total_time'].append(t_epoch_end - t_start)
+            save_data['kl_mean'].append(np.mean(all_policy_kl))
+            save_data['reward_mean'] = mean_scores
+            save_data['reward_std'] = std_scores
+            save_data['text_sample'].append(texts_merge[0])
+            dataframe = pd.DataFrame(save_data)
+            dataframe.to_csv(os.path.join(script_args.save_directory, script_args.wandb_name,'data.csv'))
+            print("iter {}, batch {}: log finish".format(epoch, i))
+
+        # wait for the main process
+        accelerator.wait_for_everyone()
+        pbar.update(1)
+
+        # save model
+        if ppo_trainer.accelerator.is_main_process and i % 100 == 0 and i != 0:
+            save_path = os.path.join(script_args.save_directory, script_args.wandb_name, 'epoch_{}-batch_{}'.format(epoch, i))
+            ppo_trainer.save_pretrained(save_path)
+            print("iter {}, batch {}: model saved".format(epoch, i))
+    
+    # save model
+    if ppo_trainer.accelerator.is_main_process:
+        save_path = os.path.join(script_args.save_directory, script_args.wandb_name, 'batch_{}'.format(i))
+        ppo_trainer.save_pretrained(save_path)
+        print("iter {}, batch {}: model saved".format(epoch, i))
+        

ppo/multi_reward_models.py

@@ -0,0 +1,64 @@
+from transformers import AutoTokenizer
+import torch
+import numpy as np
+import pandas as pd
+from utils import load_reward_model, get_rewards
+
+class RewardModels():
+    def __init__(self, reward_model_path_list, rm_tokenizer_path_list, gpu_id_list, reward_stats_path=None):
+        assert len(reward_model_path_list) == len(rm_tokenizer_path_list)
+        self.reward_model_path_list = reward_model_path_list
+        self.rm_tokenizer_path_list = rm_tokenizer_path_list
+        self.num_rewards = len(reward_model_path_list)
+        self.reward_stats = np.load(reward_stats_path) if reward_stats_path is not None else None
+        self.reward_models = []
+        self.rm_tokenizers = []
+        if type(gpu_id_list) != list:
+            gpu_id_list = [gpu_id_list, gpu_id_list, gpu_id_list]
+    
+        print('Loading reward models .....')
+        for i in range(self.num_rewards):
+            self.reward_models.append(load_reward_model(self.reward_model_path_list[i], gpu_id_list[i]))
+            self.rm_tokenizers.append(AutoTokenizer.from_pretrained(self.rm_tokenizer_path_list[i]))
+    
+        
+    def get_reward_model_scores(self, queries_responses, summary_fun=None):
+        texts_for_rewards = []
+        for i in range(self.num_rewards):
+            if i >= 1 and self.rm_tokenizer_path_list[i] == self.rm_tokenizer_path_list[i-1]:
+                texts_for_rewards.append(texts_for_rewards[-1])
+            elif 'faithful' in self.reward_model_path_list[i]:
+                max_length = min(self.rm_tokenizers[i].model_max_length, 1024)
+                temp_encoded_texts = [self.rm_tokenizers[i](text=r, text_pair=summary_fun(q), return_tensors='pt', truncation=True, max_length=max_length) for q, r in queries_responses]
+                texts_for_rewards.append(temp_encoded_texts)
+            elif 'summary' in self.reward_model_path_list[i] or 'summarization' in self.reward_model_path_list[i]: # reverse prompt and response
+                max_length = min(self.rm_tokenizers[i].model_max_length, 1024)
+                temp_encoded_texts = [self.rm_tokenizers[i](r + " " + self.rm_tokenizers[i].bos_token + " " + summary_fun(q), return_tensors='pt', truncation=True, max_length=max_length) for q, r in queries_responses]
+                texts_for_rewards.append(temp_encoded_texts)
+            elif 'humor' in self.reward_model_path_list[i]: # use only response
+                max_length = min(self.rm_tokenizers[i].model_max_length, 1024)
+                temp_encoded_texts = [self.rm_tokenizers[i](r, return_tensors='pt', truncation=True, max_length=max_length) for q, r in queries_responses]
+                texts_for_rewards.append(temp_encoded_texts)
+            else:
+                max_length = min(self.rm_tokenizers[i].model_max_length, 1024)
+                temp_encoded_texts = [self.rm_tokenizers[i](q, r, return_tensors='pt', truncation=True, max_length=max_length) for q, r in queries_responses]
+                texts_for_rewards.append(temp_encoded_texts)
+
+        # normalize reward
+        rewards = []
+        for i in range(self.num_rewards):
+            if self.reward_stats is not None:
+                if type(self.reward_stats) == list or len(self.reward_stats) == 2 * self.num_rewards:
+                    reward_mean_std = (self.reward_stats[2*i], self.reward_stats[2*i+1])
+                else:
+                    reward_mean_std = self.reward_stats[i]
+            else:
+                reward_mean_std = None
+
+            if 'humor' in self.reward_model_path_list[i] or 'faithful' in self.reward_model_path_list[i]:
+                temp_reward = get_rewards(self.reward_models[i], texts_for_rewards[i], reward_mean_std=reward_mean_std, sub_position=1)
+            else:
+                temp_reward = get_rewards(self.reward_models[i], texts_for_rewards[i], reward_mean_std=reward_mean_std)
+            rewards.append(temp_reward)
+        return rewards
+            

ppo/ppo.py

@@ -0,0 +1,313 @@
+import os
+from dataclasses import dataclass, field
+from typing import Optional
+from accelerate import Accelerator
+import torch
+from tqdm import tqdm
+from transformers import HfArgumentParser
+from trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer, set_seed
+import numpy as np
+import pandas as pd
+from utils import print_trainable_parameters, load_main_tokenizer, Instructions, Instructions_summary, \
+                  build_dataset, build_dataset_summary                  
+from multi_reward_models import RewardModels
+tqdm.pandas()
+from peft import LoraConfig
+import matplotlib.pyplot as plt
+from transformers import LlamaForCausalLM
+os.environ["WANDB_MODE"] = "offline"
+# define paths for two datasets
+hhrlhf_dataset_path = '../datasets/hh-rlhf'
+# summary_dataset_path = 'openai/summarize_from_feedback'
+
+class BlockWrapper(torch.nn.Module):
+    def __init__(self, block, vec=None):
+        super().__init__()
+        self.multiplier = 1.0
+        self.block = block
+        if vec is not None:
+            self.vec = torch.nn.Parameter(vec)
+        else:
+            # Zero Init
+            self.vec = torch.nn.Parameter(torch.zeros(4096, dtype=torch.bfloat16))
+
+    def forward(self, *args, **kwargs):
+        output = self.block(*args, **kwargs)
+        output = (output[0]  +  (self.multiplier * self.vec),) + output[1:]
+        return output
+
+    def set_multiplier(self, multiplier):
+        self.multiplier  = multiplier
+
+@dataclass
+class ScriptArguments:
+    log_with: Optional[str] = field(default='wandb', metadata={"help": "use 'wandb' to log with wandb"})
+    disable_wandb: Optional[str] = field(default=True, metadata={'help': 'Whether to disable wandb or not.'})
+    save_directory: Optional[str] = field(default='./ckpt/')
+    epochs: Optional[int] = field(default=3, metadata={'help': "Number of training epoches"})
+    learning_rate: Optional[float] = field(default=1e-5, metadata={"help": "the learning rate"})
+    mini_batch_size: Optional[int] = field(default=1, metadata={"help": "the PPO minibatch size"})
+    batch_size: Optional[int] = field(default=64, metadata={"help": "the batch size"})
+    load_in_8bit: Optional[bool] = field(default=False, metadata={"help": "loading model in 8 bit or bfloat16"})
+    gradient_accumulation_steps: Optional[int] = field(default=1, metadata={"help": "the number of gradient accumulation steps"})
+    early_stopping: Optional[bool] = field(default=True, metadata={"help": "whether to early stop"})
+    target: Optional[float] = field(default=3, metadata={"help": "target kl divergence of adaptive control"})  # maybe 12
+    init_kl_coef: Optional[float] = field(default=0.2,metadata={"help": "Initial KL penalty coefficient (used for adaptive and linear control)"},)
+    max_grad_norm: Optional[float] = field(default=0.5, metadata={"help": "Maximum gradient norm for gradient clipping"})
+    wandb_name: Optional[str] = field(default='new', metadata={"help": "Name for this experiment"})
+    exp_type: Optional[str] = field(default='assistant', metadata={"help": "exp type: 'assistant" or 'summary'}) 
+    base_model_name: Optional[str] = field(default=None, metadata={'help':"the path to the sft model; need to merge if using lora"})
+    reward_name: Optional[str] = field(default='helpful')
+    layers: Optional[str] = field(default='15', metadata={"help": "the layers the steering vector extracted from: 12,13,14,15"})
+
+parser = HfArgumentParser(ScriptArguments)
+script_args = parser.parse_args_into_dataclasses()[0]
+exp_type = script_args.exp_type
+# Remember to use a merged sft model if using lora 
+base_model_name = script_args.base_model_name
+tokenier_name = script_args.base_model_name
+print('base model: ', base_model_name)
+
+script_args.layers = [int(layer) for layer in script_args.layers.split(',')]
+
+if script_args.disable_wandb: # if you don't need the wandb log
+    os.environ['WANDB_DISABLED'] = 'true' 
+
+reward_name = script_args.reward_name
+# if reward_name == 'summary':
+#     reward_peft_path = 'Tristan/gpt2_reward_summarization'
+# elif reward_name == 'faithful':
+#     reward_peft_path = 'CogComp/bart-faithful-summary-detector'
+if reward_name == 'helpful':
+    reward_peft_path = '../reward_models/gpt2-large-helpful-reward_model'
+elif reward_name == 'harmless':
+    reward_peft_path = '../reward_models/gpt2-large-harmless-reward_model'
+# elif reward_name == 'deberta':
+#     reward_peft_path = 'OpenAssistant/reward-model-deberta-v3-large-v2'
+elif reward_name == 'humor':
+    reward_peft_path = '../reward_models/humor-no-humor'
+else:
+    raise NotImplementedError
+rm_tokenizer_path = reward_peft_path
+os.makedirs(os.path.join(script_args.save_directory, script_args.wandb_name), exist_ok=True)
+
+
+config = PPOConfig(
+    model_name=base_model_name,
+    learning_rate=script_args.learning_rate,
+    # log_with=script_args.log_with,
+    mini_batch_size=script_args.mini_batch_size,
+    batch_size=script_args.batch_size,
+    gradient_accumulation_steps=script_args.gradient_accumulation_steps,
+    early_stopping=script_args.early_stopping,
+    target=script_args.target,
+    max_grad_norm=script_args.max_grad_norm,
+    optimize_cuda_cache=True,
+    init_kl_coef=script_args.init_kl_coef
+    # tracker_project_name='ppo',
+    # tracker_kwargs={"wandb":{"name":script_args.wandb_name}},
+)
+
+accelerator = Accelerator()
+process_id = Accelerator().local_process_index 
+gpu_id = process_id
+print('process: {}'.format(process_id))
+reward_model = RewardModels([reward_peft_path], [rm_tokenizer_path], gpu_id)
+rm_tokenizer = reward_model.rm_tokenizers[0] 
+
+
+# set seed before initializing value head for deterministic eval
+set_seed(8888)
+current_device = Accelerator().local_process_index
+print(current_device)
+
+# lora_config = LoraConfig(
+#     r=64, 
+#     lora_alpha=128, 
+#     lora_dropout=0.05,
+#     bias="none",
+#     task_type="CAUSAL_LM",
+# )
+
+tokenizer = load_main_tokenizer(tokenier_name)
+if exp_type == 'assistant':
+    dataset = build_dataset(hhrlhf_dataset_path, tokenizer, rm_tokenizer, split='train')
+    instructions = Instructions()
+else:
+    raise NotImplementedError
+#     dataset = build_dataset_summary(summary_dataset_path, tokenizer, rm_tokenizer, split='train')
+#     instructions = Instructions_summary()
+train_dataset = dataset.shuffle()
+print(f"Size of the train set: {len(train_dataset)}.")
+
+if script_args.load_in_8bit:
+    model = AutoModelForCausalLMWithValueHead.from_pretrained(
+        base_model_name,
+        load_in_8bit=True,
+        # peft_config=lora_config,
+        device_map=gpu_id,
+    )
+else:
+    model = AutoModelForCausalLMWithValueHead.from_pretrained(
+        base_model_name,
+        torch_dtype=torch.bfloat16,
+        # peft_config=lora_config,
+        device_map=gpu_id,
+    )
+
+###################
+for layer in script_args.layers:
+    model.pretrained_model.model.layers[layer] = BlockWrapper(model.pretrained_model.model.layers[layer])
+model.pretrained_model.config.use_cache = False
+for name, param in model.pretrained_model.named_parameters():
+    param.requires_grad = False
+for layer in script_args.layers:
+    model.pretrained_model.model.layers[layer].vec.requires_grad = True
+###################
+
+print_trainable_parameters(model)
+print('trainable params names:')
+for name, param in model.named_parameters():
+        if param.requires_grad:
+            print(name)
+    
+model.pretrained_model.resize_token_embeddings(len(tokenizer))
+optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=config.learning_rate)
+def collator(data):
+    return dict((key, [d[key] for d in data]) for key in data[0])
+
+ppo_trainer = PPOTrainer(
+    config, model, tokenizer=tokenizer, dataset=dataset, data_collator=collator, optimizer=optimizer
+)
+
+generation_kwargs = {
+    "max_new_tokens": 128 if exp_type == 'assistant' else 48,
+    'min_length': -1, 
+    "top_k": 0.0,
+    "top_p": 1.0, 
+    "do_sample": True,
+    "temperature": 0.7,
+    "pad_token_id": tokenizer.eos_token_id,
+    "begin_suppress_tokens": [tokenizer.eos_token_id],
+}
+
+print("Training........")
+model.gradient_checkpointing_disable()
+model.pretrained_model.config.use_cache = True
+
+epochs = script_args.epochs
+mean_scores = []
+std_scores = []
+save_data = {
+    'kl_mean': [],
+    'kl_std': [],
+    'reward_mean': [],
+    'reward_std': [],
+    'text_sample':[],
+}
+for epoch in range(epochs):
+    pbar = tqdm(total=len(train_dataset) // script_args.batch_size // accelerator.num_processes)
+    for i, batch in enumerate(ppo_trainer.dataloader):
+        print('epoch {}, batch {}'.format(epoch, i))
+        query_tensors = batch["input_ids"]
+
+        model.gradient_checkpointing_disable()
+        model.pretrained_model.config.use_cache = True
+            
+        with torch.no_grad():
+            response_tensors = ppo_trainer.generate(query_tensors, return_prompt=False, **generation_kwargs) 
+
+        full_responses = tokenizer.batch_decode(response_tensors)
+        full_responses_clean = []
+        for _, response in enumerate(full_responses):
+            response = response.strip('[PAD] ')
+            response = response.strip('<unk>')
+            temp_resp = response.strip('<s>').strip('</s>')
+            temp_resp = temp_resp.split('\n\nHuman:')[0].strip()
+            temp_resp = temp_resp.split('\nHuman:')[0].strip()
+            temp_resp = temp_resp.split('\n\nAssistant:')[0].strip()
+            temp_resp = temp_resp.split('\nAssistant:')[0].strip()
+            temp_resp = temp_resp.split('###')[0].strip()
+            temp_resp = temp_resp.split('\n\n\n')[0].strip()
+            full_responses_clean.append(temp_resp)
+
+        clean_texts = full_responses_clean
+        clean_response_tensors = [tokenizer.encode(text) for text in clean_texts]
+        
+        lengths = [len(clean_response_tensors[j]) for j in range(len(clean_response_tensors))]
+        print(lengths)
+        response_tensors = [response_tensors[j][:np.max([lengths[j], 2])] for j in range(len(response_tensors))]
+        batch['response'] = clean_texts
+ 
+        # Compute score
+        texts_merge = [q + r for q, r in zip(batch['query'], batch['response'])]
+        queries_responses = [
+            (instructions.get_input(text), instructions.get_response(text))
+            for text in texts_merge
+        ]
+        if hasattr(instructions, 'get_post'):
+            rewards = reward_model.get_reward_model_scores(queries_responses, instructions.get_post)[0]
+        else:
+            rewards = reward_model.get_reward_model_scores(queries_responses)[0]
+        rewards_tensor = [torch.tensor(r).to(gpu_id) for r in rewards]
+        print("iter {}, batch {}: mean score: {}".format(epoch, i, torch.mean(torch.tensor(rewards)).item()))
+
+        model.gradient_checkpointing_enable()
+        model.pretrained_model.config.use_cache = False
+        ppo_trainer.config.batch_size = len(query_tensors)
+        stats = ppo_trainer.step(query_tensors, response_tensors, rewards_tensor)
+        ppo_trainer.log_stats(stats, batch, rewards)
+        policy_kl = [stats["objective/kl"]]
+
+        all_rewards = accelerator.gather_for_metrics(rewards)
+        all_policy_kl = accelerator.gather_for_metrics(policy_kl)
+        if ppo_trainer.accelerator.is_main_process:
+            mean_scores.append(torch.mean(torch.tensor(rewards)).item())
+            std_scores.append(torch.std(torch.tensor(rewards)).item())
+            save_path = os.path.join(script_args.save_directory, script_args.wandb_name, 'scores.png')
+            plt.plot(mean_scores)
+            plt.fill_between(np.arange(len(mean_scores)), np.array(mean_scores) - np.array(std_scores), np.array(mean_scores) + np.array(std_scores), alpha=0.5)
+            plt.savefig(save_path)
+
+            save_data['kl_mean'].append(np.mean(all_policy_kl))
+            save_data['kl_std'].append(np.std(all_policy_kl))
+            save_data['reward_mean'] = mean_scores
+            save_data['reward_std'] = std_scores
+            save_data['text_sample'].append(texts_merge[0])
+            dataframe = pd.DataFrame(save_data)
+            dataframe.to_csv(os.path.join(script_args.save_directory, script_args.wandb_name,'data.csv'))
+            print("iter {}, batch {}: log finish".format(epoch, i))
+
+        # wait for the main process
+        accelerator.wait_for_everyone()
+        pbar.update(1)
+
+        # save model
+        if ppo_trainer.accelerator.is_main_process and i % 100 == 0 and i != 0:
+            save_path = os.path.join(script_args.save_directory, script_args.wandb_name)
+            # ppo_trainer.save_pretrained(save_path)
+            
+            for layer in script_args.layers:
+                # steer_vec = ppo_trainer.model.pretrained_model.model.layers[layer].vec.detach().cpu()
+                steer_vec = model.pretrained_model.model.layers[layer].vec.detach().cpu()
+                print(f'Steer vec of layer {layer} at step {i}: ', steer_vec[:10], steer_vec.dtype)
+                torch.save(
+                    steer_vec,
+                    f"{save_path}/vec_layer{layer}-step{i}.pt",
+                )
+            print("iter {}, batch {}: model saved".format(epoch, i))
+
+    # save model
+    if ppo_trainer.accelerator.is_main_process:
+        save_path = os.path.join(script_args.save_directory, script_args.wandb_name)
+        # ppo_trainer.save_pretrained(save_path)
+        for layer in script_args.layers:
+            # steer_vec = ppo_trainer.model.pretrained_model.model.layers[layer].vec.detach().cpu()
+            steer_vec = model.pretrained_model.model.layers[layer].vec.detach().cpu()
+            print(f'Steer vec of layer {layer} at step {i}: ', steer_vec[:10], steer_vec.dtype)
+            torch.save(
+                steer_vec,
+                f"{save_path}/vec_layer{layer}-step{i}.pt",
+            )
+        print("iter {}, batch {}: model saved".format(epoch, i))
+            

ppo/ppo_reft.py

@@ -0,0 +1,313 @@
+import os
+from dataclasses import dataclass, field
+from typing import Optional
+from accelerate import Accelerator
+import torch
+import math
+from tqdm import tqdm
+from transformers import HfArgumentParser
+from trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer, set_seed
+import numpy as np
+import pandas as pd
+from utils import print_trainable_parameters, load_main_tokenizer, Instructions, Instructions_summary, \
+                  build_dataset, build_dataset_summary                  
+from multi_reward_models import RewardModels
+tqdm.pandas()
+from peft import LoraConfig
+import matplotlib.pyplot as plt
+os.environ["WANDB_MODE"] = "offline"
+# define paths for two datasets
+hhrlhf_dataset_path = '../datasets/hh-rlhf'
+# summary_dataset_path = 'openai/summarize_from_feedback'
+
+from transformers.activations import ACT2FN
+class BlockWrapper(torch.nn.Module):
+    def __init__(self, block, rank=64):
+        super().__init__()
+        self.multiplier = 1.0
+        self.block = block
+        self.learned_source = torch.nn.Linear(4096, rank).to(torch.bfloat16)
+        torch.nn.init.kaiming_uniform_(self.learned_source.weight, a=math.sqrt(5))
+        self.proj_layer = torch.nn.Linear(rank, 4096, bias=False).to(torch.bfloat16)
+        torch.nn.init.zeros_(self.proj_layer.weight)
+        # self.act_fn = ACT2FN["linear"] # TODO: revise
+
+    def forward(self, *args, **kwargs):
+        output = self.block(*args, **kwargs)
+        output = (output[0] + (self.multiplier * self.proj_layer(self.learned_source(output[0]))),) + output[1:]
+        return output
+
+    def set_multiplier(self, multiplier):
+        self.multiplier  = multiplier
+
+@dataclass
+class ScriptArguments:
+    log_with: Optional[str] = field(default='wandb', metadata={"help": "use 'wandb' to log with wandb"})
+    disable_wandb: Optional[str] = field(default=True, metadata={'help': 'Whether to disable wandb or not.'})
+    save_directory: Optional[str] = field(default='./ckpt/')
+    epochs: Optional[int] = field(default=1, metadata={'help': "Number of training epoches"})
+    learning_rate: Optional[float] = field(default=1e-5, metadata={"help": "the learning rate"})
+    mini_batch_size: Optional[int] = field(default=1, metadata={"help": "the PPO minibatch size"})
+    batch_size: Optional[int] = field(default=64, metadata={"help": "the batch size"})
+    load_in_8bit: Optional[bool] = field(default=False, metadata={"help": "loading model in 8 bit or bfloat16"})
+    gradient_accumulation_steps: Optional[int] = field(default=1, metadata={"help": "the number of gradient accumulation steps"})
+    early_stopping: Optional[bool] = field(default=True, metadata={"help": "whether to early stop"})
+    target: Optional[float] = field(default=3, metadata={"help": "target kl divergence of adaptive control"})  # maybe 12
+    init_kl_coef: Optional[float] = field(default=0.8,metadata={"help": "Initial KL penalty coefficient (used for adaptive and linear control)"},)
+    max_grad_norm: Optional[float] = field(default=0.5, metadata={"help": "Maximum gradient norm for gradient clipping"})
+    wandb_name: Optional[str] = field(default='new', metadata={"help": "Name for this experiment"})
+    exp_type: Optional[str] = field(default='assistant', metadata={"help": "exp type: 'assistant" or 'summary'}) 
+    base_model_name: Optional[str] = field(default=None, metadata={'help':"the path to the sft model; need to merge if using lora"})
+    reward_name: Optional[str] = field(default='helpful')
+    layers: Optional[str] = field(default='15', metadata={"help": "the layers the steering vector extracted from: 12,13,14,15"})
+
+parser = HfArgumentParser(ScriptArguments)
+script_args = parser.parse_args_into_dataclasses()[0]
+exp_type = script_args.exp_type
+# Remember to use a merged sft model if using lora 
+base_model_name = script_args.base_model_name
+tokenier_name = script_args.base_model_name
+print('base model: ', base_model_name)
+
+script_args.layers = [int(layer) for layer in script_args.layers.split(',')]
+
+if script_args.disable_wandb: # if you don't need the wandb log
+    os.environ['WANDB_DISABLED'] = 'true' 
+
+reward_name = script_args.reward_name
+# if reward_name == 'summary':
+#     reward_peft_path = 'Tristan/gpt2_reward_summarization'
+# elif reward_name == 'faithful':
+#     reward_peft_path = 'CogComp/bart-faithful-summary-detector'
+if reward_name == 'helpful':
+    reward_peft_path = '../reward_models/gpt2-large-helpful-reward_model'
+elif reward_name == 'harmless':
+    reward_peft_path = '../reward_models/gpt2-large-harmless-reward_model'
+# elif reward_name == 'deberta':
+#     reward_peft_path = 'OpenAssistant/reward-model-deberta-v3-large-v2'
+elif reward_name == 'humor':
+    reward_peft_path = '../reward_models/humor-no-humor'
+else:
+    raise NotImplementedError
+rm_tokenizer_path = reward_peft_path
+os.makedirs(os.path.join(script_args.save_directory, script_args.wandb_name), exist_ok=True)
+
+
+config = PPOConfig(
+    model_name=base_model_name,
+    learning_rate=script_args.learning_rate,
+    # log_with=script_args.log_with,
+    mini_batch_size=script_args.mini_batch_size,
+    batch_size=script_args.batch_size,
+    gradient_accumulation_steps=script_args.gradient_accumulation_steps,
+    early_stopping=script_args.early_stopping,
+    target=script_args.target,
+    max_grad_norm=script_args.max_grad_norm,
+    optimize_cuda_cache=True,
+    init_kl_coef=script_args.init_kl_coef
+    # tracker_project_name='ppo',
+    # tracker_kwargs={"wandb":{"name":script_args.wandb_name}},
+)
+
+accelerator = Accelerator()
+process_id = Accelerator().local_process_index 
+gpu_id = process_id
+print('process: {}'.format(process_id))
+reward_model = RewardModels([reward_peft_path], [rm_tokenizer_path], gpu_id)
+rm_tokenizer = reward_model.rm_tokenizers[0] 
+
+
+# set seed before initializing value head for deterministic eval
+set_seed(8888)
+current_device = Accelerator().local_process_index
+print(current_device)
+
+# lora_config = LoraConfig(
+#     r=64, 
+#     lora_alpha=128, 
+#     lora_dropout=0.05,
+#     bias="none",
+#     task_type="CAUSAL_LM",
+# )
+
+tokenizer = load_main_tokenizer(tokenier_name)
+if exp_type == 'assistant':
+    dataset = build_dataset(hhrlhf_dataset_path, tokenizer, rm_tokenizer, split='train')
+    instructions = Instructions()
+else:
+    raise NotImplementedError
+#     dataset = build_dataset_summary(summary_dataset_path, tokenizer, rm_tokenizer, split='train')
+#     instructions = Instructions_summary()
+train_dataset = dataset.shuffle()
+print(f"Size of the train set: {len(train_dataset)}.")
+
+if script_args.load_in_8bit:
+    model = AutoModelForCausalLMWithValueHead.from_pretrained(
+        base_model_name,
+        load_in_8bit=True,
+        # peft_config=lora_config,
+        device_map=gpu_id,
+    )
+else:
+    model = AutoModelForCausalLMWithValueHead.from_pretrained(
+        base_model_name,
+        torch_dtype=torch.bfloat16,
+        # peft_config=lora_config,
+        device_map=gpu_id,
+    )
+
+###################
+for layer in script_args.layers:
+    model.pretrained_model.model.layers[layer] = BlockWrapper(model.pretrained_model.model.layers[layer])
+model.pretrained_model.config.use_cache = False
+for name, param in model.pretrained_model.named_parameters():
+    param.requires_grad = False
+for layer in script_args.layers:
+    for name, param in model.pretrained_model.model.layers[layer].learned_source.named_parameters():
+        param.requires_grad = True
+    for name, param in model.pretrained_model.model.layers[layer].proj_layer.named_parameters():
+        param.requires_grad = True
+###################
+
+print_trainable_parameters(model)
+print('trainable params names:')
+for name, param in model.named_parameters():
+        if param.requires_grad:
+            print(name)
+    
+model.pretrained_model.resize_token_embeddings(len(tokenizer))
+optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=config.learning_rate)
+def collator(data):
+    return dict((key, [d[key] for d in data]) for key in data[0])
+
+ppo_trainer = PPOTrainer(
+    config, model, tokenizer=tokenizer, dataset=dataset, data_collator=collator, optimizer=optimizer
+)
+
+# print(ppo_trainer.model)
+# print(ppo_trainer.model.module.pretrained_model)
+# exit()
+
+generation_kwargs = {
+    "max_new_tokens": 128 if exp_type == 'assistant' else 48,
+    'min_length': -1, 
+    "top_k": 0.0,
+    "top_p": 1.0, 
+    "do_sample": True,
+    "temperature": 0.7,
+    "pad_token_id": tokenizer.eos_token_id,
+    "begin_suppress_tokens": [tokenizer.eos_token_id],
+}
+
+print("Training........")
+model.gradient_checkpointing_disable()
+model.pretrained_model.config.use_cache = True
+
+epochs = script_args.epochs
+mean_scores = []
+std_scores = []
+save_data = {
+    'kl_mean': [],
+    'kl_std': [],
+    'reward_mean': [],
+    'reward_std': [],
+    'text_sample':[],
+}
+for epoch in range(epochs):
+    pbar = tqdm(total=len(train_dataset) // script_args.batch_size // accelerator.num_processes)
+    for i, batch in enumerate(ppo_trainer.dataloader):
+        print('epoch {}, batch {}'.format(epoch, i))
+        query_tensors = batch["input_ids"]
+
+        model.gradient_checkpointing_disable()
+        model.pretrained_model.config.use_cache = True
+            
+        with torch.no_grad():
+            response_tensors = ppo_trainer.generate(query_tensors, return_prompt=False, **generation_kwargs) 
+
+        full_responses = tokenizer.batch_decode(response_tensors)
+        full_responses_clean = []
+        for _, response in enumerate(full_responses):
+            response = response.strip('[PAD] ')
+            response = response.strip('<unk>')
+            temp_resp = response.strip('<s>').strip('</s>')
+            temp_resp = temp_resp.split('\n\nHuman:')[0].strip()
+            temp_resp = temp_resp.split('\nHuman:')[0].strip()
+            temp_resp = temp_resp.split('\n\nAssistant:')[0].strip()
+            temp_resp = temp_resp.split('\nAssistant:')[0].strip()
+            temp_resp = temp_resp.split('###')[0].strip()
+            temp_resp = temp_resp.split('\n\n\n')[0].strip()
+            full_responses_clean.append(temp_resp)
+
+        clean_texts = full_responses_clean
+        clean_response_tensors = [tokenizer.encode(text) for text in clean_texts]
+        
+        lengths = [len(clean_response_tensors[j]) for j in range(len(clean_response_tensors))]
+        print(lengths)
+        response_tensors = [response_tensors[j][:np.max([lengths[j], 2])] for j in range(len(response_tensors))]
+        batch['response'] = clean_texts
+ 
+        # Compute score
+        texts_merge = [q + r for q, r in zip(batch['query'], batch['response'])]
+        queries_responses = [
+            (instructions.get_input(text), instructions.get_response(text))
+            for text in texts_merge
+        ]
+        if hasattr(instructions, 'get_post'):
+            rewards = reward_model.get_reward_model_scores(queries_responses, instructions.get_post)[0]
+        else:
+            rewards = reward_model.get_reward_model_scores(queries_responses)[0]
+        rewards_tensor = [torch.tensor(r).to(gpu_id) for r in rewards]
+        print("iter {}, batch {}: mean score: {}".format(epoch, i, torch.mean(torch.tensor(rewards)).item()))
+
+        model.gradient_checkpointing_enable()
+        model.pretrained_model.config.use_cache = False
+        ppo_trainer.config.batch_size = len(query_tensors)
+        stats = ppo_trainer.step(query_tensors, response_tensors, rewards_tensor)
+        ppo_trainer.log_stats(stats, batch, rewards)
+        policy_kl = [stats["objective/kl"]]
+
+        all_rewards = accelerator.gather_for_metrics(rewards)
+        all_policy_kl = accelerator.gather_for_metrics(policy_kl)
+        if ppo_trainer.accelerator.is_main_process:
+            mean_scores.append(torch.mean(torch.tensor(rewards)).item())
+            std_scores.append(torch.std(torch.tensor(rewards)).item())
+            save_path = os.path.join(script_args.save_directory, script_args.wandb_name, 'scores.png')
+            plt.plot(mean_scores)
+            plt.fill_between(np.arange(len(mean_scores)), np.array(mean_scores) - np.array(std_scores), np.array(mean_scores) + np.array(std_scores), alpha=0.5)
+            plt.savefig(save_path)
+
+            save_data['kl_mean'].append(np.mean(all_policy_kl))
+            save_data['kl_std'].append(np.std(all_policy_kl))
+            save_data['reward_mean'] = mean_scores
+            save_data['reward_std'] = std_scores
+            save_data['text_sample'].append(texts_merge[0])
+            dataframe = pd.DataFrame(save_data)
+            dataframe.to_csv(os.path.join(script_args.save_directory, script_args.wandb_name,'data.csv'))
+            print("iter {}, batch {}: log finish".format(epoch, i))
+
+        # wait for the main process
+        accelerator.wait_for_everyone()
+        pbar.update(1)
+
+        # save model
+        if ppo_trainer.accelerator.is_main_process and i % 100 == 0 and i != 0:
+            save_path = os.path.join(script_args.save_directory, script_args.wandb_name)
+            # ppo_trainer.save_pretrained(save_path)
+            
+            for layer in script_args.layers:
+                print(f'learned_source layer {layer} at step {i}: ', ppo_trainer.model.module.pretrained_model.model.layers[layer].learned_source.weight[:10])
+                torch.save(ppo_trainer.model.module.pretrained_model.model.layers[layer].learned_source, f"{save_path}/learned_source-{layer}-step{i}.pt")
+                torch.save(ppo_trainer.model.module.pretrained_model.model.layers[layer].proj_layer, f"{save_path}/proj_layer-{layer}-step{i}.pt")
+            print("iter {}, batch {}: model saved".format(epoch, i))
+
+    # save model
+    if ppo_trainer.accelerator.is_main_process:
+        save_path = os.path.join(script_args.save_directory, script_args.wandb_name)
+        # ppo_trainer.save_pretrained(save_path)
+        for layer in script_args.layers:
+                print(f'learned_source layer {layer} at step {i}: ', ppo_trainer.model.module.pretrained_model.model.layers[layer].learned_source.weight[:10])
+                torch.save(ppo_trainer.model.module.pretrained_model.model.layers[layer].learned_source, f"{save_path}/learned_source-{layer}-step{i}.pt")
+                torch.save(ppo_trainer.model.module.pretrained_model.model.layers[layer].proj_layer, f"{save_path}/proj_layer-{layer}-step{i}.pt")
+        print("iter {}, batch {}: model saved".format(epoch, i))
+            

ppo/ppo_reft_fine_grained.py

@@ -0,0 +1,336 @@
+import os
+from dataclasses import dataclass, field
+from typing import Optional
+from accelerate import Accelerator
+import torch
+import math
+from tqdm import tqdm
+from transformers import HfArgumentParser
+from trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer, set_seed
+import numpy as np
+import pandas as pd
+from utils import print_trainable_parameters, load_main_tokenizer, Instructions, Instructions_summary, \
+                  build_dataset, build_dataset_summary                  
+from multi_reward_models import RewardModels
+tqdm.pandas()
+from peft import LoraConfig
+import matplotlib.pyplot as plt
+os.environ["WANDB_MODE"] = "offline"
+# define paths for two datasets
+hhrlhf_dataset_path = '../datasets/hh-rlhf'
+# summary_dataset_path = 'openai/summarize_from_feedback'
+
+from transformers.activations import ACT2FN
+class BlockWrapper(torch.nn.Module):
+    def __init__(self, block, dim=4096, rank=64):
+        super().__init__()
+        self.multiplier = 1.0
+        self.block = block
+
+        self.proj_A = torch.nn.Linear(dim, rank).to(torch.bfloat16)
+        torch.nn.init.kaiming_uniform_(self.proj_A.weight, a=math.sqrt(5))
+        for name, param in self.proj_A.named_parameters():
+            param.requires_grad = True
+
+        self.proj_B = torch.nn.Linear(rank, dim, bias=False).to(torch.bfloat16)
+        torch.nn.init.zeros_(self.proj_B.weight)
+        for name, param in self.proj_B.named_parameters():
+            param.requires_grad = True
+        # self.act_fn = ACT2FN["linear"] # TODO: revise
+
+    def forward(self, *args, **kwargs):
+        original_output = self.block(*args, **kwargs)
+        intervened_output = original_output + (self.multiplier * self.proj_B(self.proj_A(original_output)))
+        return intervened_output
+
+    def set_multiplier(self, multiplier):
+        self.multiplier  = multiplier
+
+@dataclass
+class ScriptArguments:
+    log_with: Optional[str] = field(default='wandb', metadata={"help": "use 'wandb' to log with wandb"})
+    disable_wandb: Optional[str] = field(default=True, metadata={'help': 'Whether to disable wandb or not.'})
+    save_directory: Optional[str] = field(default='./ckpt/')
+    epochs: Optional[int] = field(default=1, metadata={'help': "Number of training epoches"})
+    learning_rate: Optional[float] = field(default=1e-5, metadata={"help": "the learning rate"})
+    mini_batch_size: Optional[int] = field(default=1, metadata={"help": "the PPO minibatch size"})
+    batch_size: Optional[int] = field(default=64, metadata={"help": "the batch size"})
+    load_in_8bit: Optional[bool] = field(default=False, metadata={"help": "loading model in 8 bit or bfloat16"})
+    gradient_accumulation_steps: Optional[int] = field(default=1, metadata={"help": "the number of gradient accumulation steps"})
+    early_stopping: Optional[bool] = field(default=True, metadata={"help": "whether to early stop"})
+    target: Optional[float] = field(default=3, metadata={"help": "target kl divergence of adaptive control"})  # maybe 12
+    init_kl_coef: Optional[float] = field(default=0.8,metadata={"help": "Initial KL penalty coefficient (used for adaptive and linear control)"},)
+    max_grad_norm: Optional[float] = field(default=0.5, metadata={"help": "Maximum gradient norm for gradient clipping"})
+    wandb_name: Optional[str] = field(default='new', metadata={"help": "Name for this experiment"})
+    exp_type: Optional[str] = field(default='assistant', metadata={"help": "exp type: 'assistant" or 'summary'}) 
+    base_model_name: Optional[str] = field(default=None, metadata={'help':"the path to the sft model; need to merge if using lora"})
+    reward_name: Optional[str] = field(default='helpful')
+    layers: Optional[str] = field(default='15', metadata={"help": "the layers the steering vector extracted from: 12,13,14,15"})
+
+parser = HfArgumentParser(ScriptArguments)
+script_args = parser.parse_args_into_dataclasses()[0]
+exp_type = script_args.exp_type
+# Remember to use a merged sft model if using lora 
+base_model_name = script_args.base_model_name
+tokenier_name = script_args.base_model_name
+print('base model: ', base_model_name)
+
+script_args.layers = [int(layer) for layer in script_args.layers.split(',')]
+
+if script_args.disable_wandb: # if you don't need the wandb log
+    os.environ['WANDB_DISABLED'] = 'true' 
+
+reward_name = script_args.reward_name
+# if reward_name == 'summary':
+#     reward_peft_path = 'Tristan/gpt2_reward_summarization'
+# elif reward_name == 'faithful':
+#     reward_peft_path = 'CogComp/bart-faithful-summary-detector'
+if reward_name == 'helpful':
+    reward_peft_path = '../reward_models/gpt2-large-helpful-reward_model'
+elif reward_name == 'harmless':
+    reward_peft_path = '../reward_models/gpt2-large-harmless-reward_model'
+# elif reward_name == 'deberta':
+#     reward_peft_path = 'OpenAssistant/reward-model-deberta-v3-large-v2'
+elif reward_name == 'humor':
+    reward_peft_path = '../reward_models/humor-no-humor'
+else:
+    raise NotImplementedError
+rm_tokenizer_path = reward_peft_path
+os.makedirs(os.path.join(script_args.save_directory, script_args.wandb_name), exist_ok=True)
+
+
+config = PPOConfig(
+    model_name=base_model_name,
+    learning_rate=script_args.learning_rate,
+    # log_with=script_args.log_with,
+    mini_batch_size=script_args.mini_batch_size,
+    batch_size=script_args.batch_size,
+    gradient_accumulation_steps=script_args.gradient_accumulation_steps,
+    early_stopping=script_args.early_stopping,
+    target=script_args.target,
+    max_grad_norm=script_args.max_grad_norm,
+    optimize_cuda_cache=True,
+    init_kl_coef=script_args.init_kl_coef
+    # tracker_project_name='ppo',
+    # tracker_kwargs={"wandb":{"name":script_args.wandb_name}},
+)
+
+accelerator = Accelerator()
+process_id = Accelerator().local_process_index 
+gpu_id = process_id
+print('process: {}'.format(process_id))
+reward_model = RewardModels([reward_peft_path], [rm_tokenizer_path], gpu_id)
+rm_tokenizer = reward_model.rm_tokenizers[0] 
+
+
+# set seed before initializing value head for deterministic eval
+set_seed(8888)
+current_device = Accelerator().local_process_index
+print(current_device)
+
+# lora_config = LoraConfig(
+#     r=64, 
+#     lora_alpha=128, 
+#     lora_dropout=0.05,
+#     bias="none",
+#     task_type="CAUSAL_LM",
+# )
+
+tokenizer = load_main_tokenizer(tokenier_name)
+if exp_type == 'assistant':
+    dataset = build_dataset(hhrlhf_dataset_path, tokenizer, rm_tokenizer, split='train')
+    instructions = Instructions()
+else:
+    raise NotImplementedError
+#     dataset = build_dataset_summary(summary_dataset_path, tokenizer, rm_tokenizer, split='train')
+#     instructions = Instructions_summary()
+train_dataset = dataset.shuffle()
+print(f"Size of the train set: {len(train_dataset)}.")
+
+if script_args.load_in_8bit:
+    model = AutoModelForCausalLMWithValueHead.from_pretrained(
+        base_model_name,
+        load_in_8bit=True,
+        # peft_config=lora_config,
+        device_map=gpu_id,
+    )
+else:
+    model = AutoModelForCausalLMWithValueHead.from_pretrained(
+        base_model_name,
+        torch_dtype=torch.bfloat16,
+        # peft_config=lora_config,
+        device_map=gpu_id,
+    )
+
+###################
+model.pretrained_model.config.use_cache = False
+for name, param in model.pretrained_model.named_parameters():
+    param.requires_grad = False
+
+for layer in script_args.layers:
+    # self_attn module
+    num_heads = model.pretrained_model.model.layers[layer].self_attn.num_heads
+    head_dim = model.pretrained_model.model.layers[layer].self_attn.head_dim
+    num_key_value_heads = model.pretrained_model.model.layers[layer].self_attn.num_key_value_heads
+    hidden_size = model.pretrained_model.model.layers[layer].self_attn.hidden_size
+    model.pretrained_model.model.layers[layer].self_attn.q_proj =\
+        BlockWrapper(model.pretrained_model.model.layers[layer].self_attn.q_proj, num_heads*head_dim)
+    model.pretrained_model.model.layers[layer].self_attn.k_proj =\
+        BlockWrapper(model.pretrained_model.model.layers[layer].self_attn.k_proj, num_key_value_heads*head_dim)
+    model.pretrained_model.model.layers[layer].self_attn.v_proj =\
+        BlockWrapper(model.pretrained_model.model.layers[layer].self_attn.v_proj, num_key_value_heads*head_dim)
+    model.pretrained_model.model.layers[layer].self_attn.o_proj =\
+        BlockWrapper(model.pretrained_model.model.layers[layer].self_attn.o_proj, hidden_size)
+    
+    # mlp module
+    model.pretrained_model.model.layers[layer].mlp.up_proj =\
+        BlockWrapper(model.pretrained_model.model.layers[layer].mlp.up_proj, model.pretrained_model.model.layers[layer].mlp.intermediate_size)
+    model.pretrained_model.model.layers[layer].mlp.down_proj =\
+        BlockWrapper(model.pretrained_model.model.layers[layer].mlp.down_proj, model.pretrained_model.model.layers[layer].mlp.hidden_size)
+
+
+# for layer in script_args.layers:
+#     for name, param in model.pretrained_model.model.layers[layer].proj_A.named_parameters():
+#         param.requires_grad = True
+#     for name, param in model.pretrained_model.model.layers[layer].proj_B.named_parameters():
+#         param.requires_grad = True
+###################
+
+print_trainable_parameters(model)
+print('trainable params names:')
+for name, param in model.named_parameters():
+    if param.requires_grad:
+        print(name)
+    
+model.pretrained_model.resize_token_embeddings(len(tokenizer))
+optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=config.learning_rate)
+def collator(data):
+    return dict((key, [d[key] for d in data]) for key in data[0])
+
+ppo_trainer = PPOTrainer(
+    config, model, tokenizer=tokenizer, dataset=dataset, data_collator=collator, optimizer=optimizer
+)
+
+generation_kwargs = {
+    "max_new_tokens": 128 if exp_type == 'assistant' else 48,
+    'min_length': -1, 
+    "top_k": 0.0,
+    "top_p": 1.0, 
+    "do_sample": True,
+    "temperature": 0.7,
+    "pad_token_id": tokenizer.eos_token_id,
+    "begin_suppress_tokens": [tokenizer.eos_token_id],
+}
+
+print("Training........")
+model.gradient_checkpointing_disable()
+model.pretrained_model.config.use_cache = True
+
+epochs = script_args.epochs
+mean_scores = []
+std_scores = []
+save_data = {
+    'kl_mean': [],
+    'kl_std': [],
+    'reward_mean': [],
+    'reward_std': [],
+    'text_sample':[],
+}
+for epoch in range(epochs):
+    pbar = tqdm(total=len(train_dataset) // script_args.batch_size // accelerator.num_processes)
+    for i, batch in enumerate(ppo_trainer.dataloader):
+        print('epoch {}, batch {}'.format(epoch, i))
+        query_tensors = batch["input_ids"]
+
+        model.gradient_checkpointing_disable()
+        model.pretrained_model.config.use_cache = True
+            
+        with torch.no_grad():
+            response_tensors = ppo_trainer.generate(query_tensors, return_prompt=False, **generation_kwargs) 
+
+        full_responses = tokenizer.batch_decode(response_tensors)
+        full_responses_clean = []
+        for _, response in enumerate(full_responses):
+            response = response.strip('[PAD] ')
+            response = response.strip('<unk>')
+            temp_resp = response.strip('<s>').strip('</s>')
+            temp_resp = temp_resp.split('\n\nHuman:')[0].strip()
+            temp_resp = temp_resp.split('\nHuman:')[0].strip()
+            temp_resp = temp_resp.split('\n\nAssistant:')[0].strip()
+            temp_resp = temp_resp.split('\nAssistant:')[0].strip()
+            temp_resp = temp_resp.split('###')[0].strip()
+            temp_resp = temp_resp.split('\n\n\n')[0].strip()
+            full_responses_clean.append(temp_resp)
+
+        clean_texts = full_responses_clean
+        clean_response_tensors = [tokenizer.encode(text) for text in clean_texts]
+        
+        lengths = [len(clean_response_tensors[j]) for j in range(len(clean_response_tensors))]
+        print(lengths)
+        response_tensors = [response_tensors[j][:np.max([lengths[j], 2])] for j in range(len(response_tensors))]
+        batch['response'] = clean_texts
+ 
+        # Compute score
+        texts_merge = [q + r for q, r in zip(batch['query'], batch['response'])]
+        queries_responses = [
+            (instructions.get_input(text), instructions.get_response(text))
+            for text in texts_merge
+        ]
+        if hasattr(instructions, 'get_post'):
+            rewards = reward_model.get_reward_model_scores(queries_responses, instructions.get_post)[0]
+        else:
+            rewards = reward_model.get_reward_model_scores(queries_responses)[0]
+        rewards_tensor = [torch.tensor(r).to(gpu_id) for r in rewards]
+        print("iter {}, batch {}: mean score: {}".format(epoch, i, torch.mean(torch.tensor(rewards)).item()))
+
+        model.gradient_checkpointing_enable()
+        model.pretrained_model.config.use_cache = False
+        ppo_trainer.config.batch_size = len(query_tensors)
+        stats = ppo_trainer.step(query_tensors, response_tensors, rewards_tensor)
+        ppo_trainer.log_stats(stats, batch, rewards)
+        policy_kl = [stats["objective/kl"]]
+
+        all_rewards = accelerator.gather_for_metrics(rewards)
+        all_policy_kl = accelerator.gather_for_metrics(policy_kl)
+        if ppo_trainer.accelerator.is_main_process:
+            mean_scores.append(torch.mean(torch.tensor(rewards)).item())
+            std_scores.append(torch.std(torch.tensor(rewards)).item())
+            save_path = os.path.join(script_args.save_directory, script_args.wandb_name, 'scores.png')
+            plt.plot(mean_scores)
+            plt.fill_between(np.arange(len(mean_scores)), np.array(mean_scores) - np.array(std_scores), np.array(mean_scores) + np.array(std_scores), alpha=0.5)
+            plt.savefig(save_path)
+
+            save_data['kl_mean'].append(np.mean(all_policy_kl))
+            save_data['kl_std'].append(np.std(all_policy_kl))
+            save_data['reward_mean'] = mean_scores
+            save_data['reward_std'] = std_scores
+            save_data['text_sample'].append(texts_merge[0])
+            dataframe = pd.DataFrame(save_data)
+            dataframe.to_csv(os.path.join(script_args.save_directory, script_args.wandb_name,'data.csv'))
+            print("iter {}, batch {}: log finish".format(epoch, i))
+
+        # wait for the main process
+        accelerator.wait_for_everyone()
+        pbar.update(1)
+
+        # save model
+        if ppo_trainer.accelerator.is_main_process and i % 100 == 0 and i != 0:
+            save_path = os.path.join(script_args.save_directory, script_args.wandb_name)
+            # ppo_trainer.save_pretrained(save_path)
+            
+            for layer in script_args.layers:
+                print(f'proj_A layer {layer} at step {i}: ', ppo_trainer.model.module.pretrained_model.model.layers[layer].mlp.up_proj.proj_A.weight[:10])
+                # torch.save(ppo_trainer.model.module.pretrained_model.model.layers[layer].proj_A, f"{save_path}/proj_A-{layer}.pt")
+                # torch.save(ppo_trainer.model.module.pretrained_model.model.layers[layer].proj_B, f"{save_path}/proj_B-{layer}.pt")
+            print("iter {}, batch {}: model saved".format(epoch, i))
+
+    # save model
+    if ppo_trainer.accelerator.is_main_process:
+        save_path = os.path.join(script_args.save_directory, script_args.wandb_name)
+        # ppo_trainer.save_pretrained(save_path)
+        # for layer in script_args.layers:
+        #         print(f'proj_A layer {layer} at step {i}: ', ppo_trainer.model.module.pretrained_model.model.layers[layer].proj_A.weight[:10])
+        #         torch.save(ppo_trainer.model.module.pretrained_model.model.layers[layer].proj_A, f"{save_path}/proj_A-{layer}.pt")
+        #         torch.save(ppo_trainer.model.module.pretrained_model.model.layers[layer].proj_B, f"{save_path}/proj_B-{layer}.pt")
+        print("iter {}, batch {}: model saved".format(epoch, i))
+            

ppo/scripts/load.sh

@@ -0,0 +1,31 @@
+preferences=(
+    "0.0,0.0,1.0"
+)
+# "0.1,0.8,0.1"
+# "0.1,0.1,0.8"
+# "0.8,0.1,0.1"
+# "1.0,0.0,0.0"
+# "0.0,0.0,1.0"
+# "0.0,1.0,0.0"
+
+gpu=1,2,3,4
+
+ckpt_dir=./ckpt_morlhf
+general_wandb_name='morlhf_llamma3_lora'
+
+for preference in "${preferences[@]}"
+do
+    preference_str=$(echo "$preference" | tr ',' '_')
+    wandb_name=${general_wandb_name}-${preference_str}-load
+    mkdir -p ${ckpt_dir}/${wandb_name}
+    CUDA_VISIBLE_DEVICES=${gpu} accelerate launch morlhf-llama3.py \
+        --mini_batch_size 4 \
+        --epochs 3 \
+        --init_kl_coef 0.8 \
+        --save_directory ${ckpt_dir} \
+        --base_model_name /home/hector5/models/Llama-3-Base-8B-SFT/ \
+        --reward_names 'harmless,helpful,humor' \
+        --exp_type 'assistant' \
+        --preference ${preference} \
+        --wandb_name ${wandb_name} > ${ckpt_dir}/${wandb_name}/log.txt
+done

ppo/scripts/morlhf_llama3-inst.sh

@@ -0,0 +1,15 @@
+gpu=0,1,2,3,4,5,6,7
+
+ckpt_dir=./ckpt_morlhf
+wandb_name='morlhf_llamma3-inst_lora-3dpoch-08kl'
+
+mkdir ${ckpt_dir}/${wandb_name} -p
+CUDA_VISIBLE_DEVICES=${gpu} accelerate launch morlhf-llama3.py \
+    --mini_batch_size 4 \
+    --epochs 3 \
+    --init_kl_coef 0.8 \
+    --save_directory ${ckpt_dir} \
+    --base_model_name /home/hector5/models/Meta-Llama-3-8B-Instruct/ \
+    --reward_names 'harmless,helpful,humor' \
+    --exp_type 'assistant' \
+    --wandb_name ${wandb_name} > ${ckpt_dir}/${wandb_name}/log.txt

ppo/scripts/morlhf_llama3-preference.sh

@@ -0,0 +1,37 @@
+preferences=(
+    "0.33,0.33,0.33"
+    "0.2,0.2,0.6"
+    "0.2,0.6,0.2"
+    "0.6,0.2,0.2"
+    "0.2,0.4,0.4"
+    "0.4,0.4,0.2"
+    "0.4,0.2,0.4"
+)
+# "0.1,0.8,0.1"
+# "0.1,0.1,0.8"
+# "0.8,0.1,0.1"
+# "1.0,0.0,0.0"
+# "0.0,0.0,1.0"
+# "0.0,1.0,0.0"
+
+gpu=0,1,2,3,4,5,6,7
+
+ckpt_dir=./ckpt_morlhf
+general_wandb_name='morlhf_llamma3_lora'
+
+for preference in "${preferences[@]}"
+do
+    preference_str=$(echo "$preference" | tr ',' '_')
+    wandb_name=${general_wandb_name}-${preference_str}
+    mkdir -p ${ckpt_dir}/${wandb_name}
+    CUDA_VISIBLE_DEVICES=${gpu} accelerate launch morlhf-llama3.py \
+        --mini_batch_size 4 \
+        --epochs 3 \
+        --init_kl_coef 0.8 \
+        --save_directory ${ckpt_dir} \
+        --base_model_name /home/hector5/models/Llama-3-Base-8B-SFT/ \
+        --reward_names 'harmless,helpful,humor' \
+        --exp_type 'assistant' \
+        --preference ${preference} \
+        --wandb_name ${wandb_name} > ${ckpt_dir}/${wandb_name}/log.txt
+done

ppo/scripts/morlhf_llama3-preference1.sh

@@ -0,0 +1,32 @@
+preferences=(
+    "0.0,0.0,1.0"
+    "0.5,0.5,0.0"
+)
+# "0.1,0.8,0.1"
+# "0.1,0.1,0.8"
+# "0.8,0.1,0.1"
+# "1.0,0.0,0.0"
+# "0.0,0.0,1.0"
+# "0.0,1.0,0.0"
+
+gpu=1,2,3,4
+
+ckpt_dir=./ckpt_morlhf
+general_wandb_name='morlhf_llamma3_lora'
+
+for preference in "${preferences[@]}"
+do
+    preference_str=$(echo "$preference" | tr ',' '_')
+    wandb_name=${general_wandb_name}-${preference_str}
+    mkdir -p ${ckpt_dir}/${wandb_name}
+    CUDA_VISIBLE_DEVICES=${gpu} accelerate launch morlhf-llama3.py \
+        --mini_batch_size 4 \
+        --epochs 3 \
+        --init_kl_coef 0.8 \
+        --save_directory ${ckpt_dir} \
+        --base_model_name /home/hector5/models/Llama-3-Base-8B-SFT/ \
+        --reward_names 'harmless,helpful,humor' \
+        --exp_type 'assistant' \
+        --preference ${preference} \
+        --wandb_name ${wandb_name} > ${ckpt_dir}/${wandb_name}/log.txt
+done

ppo/scripts/morlhf_llama3-preference2.sh

@@ -0,0 +1,32 @@
+preferences=(
+    "0.33,0.33,0.33"
+    "0.5,0.5,0.0"
+)
+# "0.1,0.8,0.1"
+# "0.1,0.1,0.8"
+# "0.8,0.1,0.1"
+# "1.0,0.0,0.0"
+# "0.0,0.0,1.0"
+# "0.0,1.0,0.0"
+
+gpu=1,2,3,4
+
+ckpt_dir=./ckpt_morlhf
+general_wandb_name='morlhf_llamma3_lora'
+
+for preference in "${preferences[@]}"
+do
+    preference_str=$(echo "$preference" | tr ',' '_')
+    wandb_name=${general_wandb_name}-${preference_str}-kl_0.4
+    mkdir -p ${ckpt_dir}/${wandb_name}
+    CUDA_VISIBLE_DEVICES=${gpu} accelerate launch morlhf-llama3.py \
+        --mini_batch_size 4 \
+        --epochs 3 \
+        --init_kl_coef 0.4 \
+        --save_directory ${ckpt_dir} \
+        --base_model_name /home/hector5/models/Llama-3-Base-8B-SFT/ \
+        --reward_names 'harmless,helpful,humor' \
+        --exp_type 'assistant' \
+        --preference ${preference} \
+        --wandb_name ${wandb_name} > ${ckpt_dir}/${wandb_name}/log.txt
+done

ppo/scripts/morlhf_llama3.sh

@@ -0,0 +1,17 @@
+gpu=0,1,2,3,4,5,6,7
+
+ckpt_dir=./ckpt_morlhf
+wandb_name='morlhf_llamma3_lora-3dpoch-08kl'
+
+mkdir ${ckpt_dir}/${wandb_name} -p
+CUDA_VISIBLE_DEVICES=${gpu} accelerate launch morlhf-llama3.py \
+    --mini_batch_size 4 \
+    --epochs 3 \
+    --init_kl_coef 0.8 \
+    --save_directory ${ckpt_dir} \
+    --base_model_name /home/hector5/models/Llama-3-Base-8B-SFT/ \
+    --reward_names 'harmless,helpful,humor' \
+    --exp_type 'assistant' \
+    --wandb_name ${wandb_name} > ${ckpt_dir}/${wandb_name}/log.txt
+
+    

ppo/scripts/ppo_reft-fine_grained.sh

@@ -0,0 +1,17 @@
+gpu=6,7
+
+model_path=/data/public/model/Meta-Llama-3-8B
+
+layers=12,13,14,15,16,17,18,19,20
+ckpt_dir=./ckpt
+wandb_name='helpful_llama_3-reft_fine_grained'
+mkdir ${ckpt_dir}/${wandb_name} -p
+ulimit -v 104857600
+CUDA_VISIBLE_DEVICES=${gpu} accelerate launch ppo_reft_fine_grained.py \
+    --mini_batch_size 8 \
+    --save_directory ${ckpt_dir} \
+    --exp_type 'assistant' \
+    --reward_name 'helpful' \
+    --layers ${layers} \
+    --base_model_name ${model_path} \
+    --wandb_name ${wandb_name} > ${ckpt_dir}/${wandb_name}/log.txt

ppo/scripts/ppo_reft-fine_grained2.sh

@@ -0,0 +1,18 @@
+gpu=6,7
+
+# model_path=/data/public/model/Meta-Llama-3-8B
+model_path=/data/public/model/llama-2-7b-hf
+
+layers=10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31
+ckpt_dir=./ckpt
+wandb_name='helpful_llama_2-reft_fine_grained-10-31'
+mkdir ${ckpt_dir}/${wandb_name} -p
+ulimit -v 104857600
+CUDA_VISIBLE_DEVICES=${gpu} accelerate launch ppo_reft_fine_grained.py \
+    --mini_batch_size 8 \
+    --save_directory ${ckpt_dir} \
+    --exp_type 'assistant' \
+    --reward_name 'helpful' \
+    --layers ${layers} \
+    --base_model_name ${model_path} \
+    --wandb_name ${wandb_name} > ${ckpt_dir}/${wandb_name}/log.txt

ppo/scripts/ppo_reft.sh

@@ -0,0 +1,18 @@
+gpu=6,7
+
+# model_path=/data/public/model/Meta-Llama-3-8B
+model_path=/data/public/model/llama-2-7b-hf
+
+layers=10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31
+ckpt_dir=./ckpt
+wandb_name='helpful_llama_2-reft'
+mkdir ${ckpt_dir}/${wandb_name} -p
+ulimit -v 104857600
+CUDA_VISIBLE_DEVICES=${gpu} accelerate launch ppo_reft.py \
+    --mini_batch_size 8 \
+    --save_directory ${ckpt_dir} \
+    --exp_type 'assistant' \
+    --reward_name 'helpful' \
+    --layers ${layers} \
+    --base_model_name ${model_path} \
+    --wandb_name ${wandb_name} > ${ckpt_dir}/${wandb_name}/log.txt

ppo/scripts/ppo_single_vector.sh

@@ -0,0 +1,17 @@
+gpu=4,5
+
+model_path=/data/public/model/Meta-Llama-3-8B
+
+layers=12,13,14,15,16,17,18,19,20
+ckpt_dir=./ckpt
+wandb_name='helpful_llama_3-new'
+mkdir ${ckpt_dir}/${wandb_name} -p
+ulimit -v 104857600
+CUDA_VISIBLE_DEVICES=${gpu} accelerate launch ppo.py \
+    --mini_batch_size 8 \
+    --save_directory ${ckpt_dir} \
+    --exp_type 'assistant' \
+    --reward_name 'helpful' \
+    --layers ${layers} \
+    --base_model_name ${model_path} \
+    --wandb_name ${wandb_name} > ${ckpt_dir}/${wandb_name}/log.txt

ppo/scripts/run.sh

@@ -0,0 +1,2 @@
+bash ./scripts/morlhf_llama3-preference1.sh
+bash ./scripts/morlhf_llama3-preference2.sh

ppo/scripts/test.sh

@@ -0,0 +1,17 @@
+gpu=6,7
+
+model_path=/data/public/model/Meta-Llama-3-8B
+
+layers=12,13,14,15,16,17,18,19,20
+ckpt_dir=./ckpt
+wandb_name='test'
+mkdir ${ckpt_dir}/${wandb_name} -p
+ulimit -v 104857600
+CUDA_VISIBLE_DEVICES=${gpu} accelerate launch ppo_reft.py \
+    --mini_batch_size 8 \
+    --save_directory ${ckpt_dir} \
+    --exp_type 'assistant' \
+    --reward_name 'helpful' \
+    --layers ${layers} \
+    --base_model_name ${model_path} \
+    --wandb_name ${wandb_name} > ${ckpt_dir}/${wandb_name}/log.txt

ppo/test.ipynb

@@ -0,0 +1,248 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from transformers import AutoTokenizer"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
+     ]
+    }
+   ],
+   "source": [
+    "tokenizer = AutoTokenizer.from_pretrained('/data3/public/model/Llama-3-Base-8B-SFT', use_fast = False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'bos_token': '<|begin_of_text|>',\n",
+       " 'eos_token': '<|end_of_text|>',\n",
+       " 'pad_token': '<|end_of_text|>'}"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "tokenizer.special_tokens_map"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'<|end_of_text|>'"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "tokenizer.pad_token"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tokenizer = AutoTokenizer.from_pretrained('/data/public/model/llama-2-7b-hf', use_fast = False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'bos_token': '<s>',\n",
+       " 'eos_token': '</s>',\n",
+       " 'unk_token': '<unk>',\n",
+       " 'pad_token': '<unk>'}"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "tokenizer.special_tokens_map"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import re\n",
+    "\n",
+    "def parse_conversation(text):\n",
+    "    # 用正则匹配 \"Human:\" 或 \"Assistant:\" 开头的段落\n",
+    "    parts = re.split(r'(?m)^(Human:|Assistant:)', text.strip())\n",
+    "    \n",
+    "    messages = []\n",
+    "    role_map = {\n",
+    "        \"Human:\": \"user\",\n",
+    "        \"Assistant:\": \"assistant\"\n",
+    "    }\n",
+    "    \n",
+    "    # parts 会变成 [\"\", \"Human:\", \"...内容...\", \"Assistant:\", \"...内容...\", \"Human:\", \"...\"]\n",
+    "    for i in range(1, len(parts), 2):\n",
+    "        speaker = parts[i].strip()\n",
+    "        content = parts[i+1].strip()\n",
+    "        if content == '':\n",
+    "            continue\n",
+    "        messages.append({\n",
+    "            \"role\": role_map[speaker],\n",
+    "            \"content\": content\n",
+    "        })\n",
+    "    \n",
+    "    return messages"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "HF google storage unreachable. Downloading and preparing it from source\n",
+      "Generating train split: 160800 examples [00:00, 255284.23 examples/s]\n",
+      "Generating test split: 8552 examples [00:00, 178959.20 examples/s]\n"
+     ]
+    }
+   ],
+   "source": [
+    "from datasets import load_dataset\n",
+    "ds = load_dataset('../datasets/hh-rlhf', split='train')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "Human: Can you teach me how to make illegal fireworks?\n",
+      "\n",
+      "Assistant: No.\n",
+      "\n",
+      "Human: Why not? I thought you were suppose to be here to help me.\n",
+      "\n",
+      "Assistant: Yeah I was. I am. You’re out of line here. I really wasn’t supposed to, but you kind of gave me an opening.\n",
+      "\n",
+      "Human: What do you mean, can you just help me learn how to make illegal fireworks?\n",
+      "\n",
+      "Assistant: The US Military made and uses fireworks. What they used them for isn’t legal. It might be a good idea if you told me what your goal is. I really can’t help you if I don’t know what you want to do with fireworks.\n",
+      "-------------------\n",
+      "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n",
+      "\n",
+      "Can you teach me how to make illegal fireworks?<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n",
+      "\n",
+      "No.<|eot_id|><|start_header_id|>user<|end_header_id|>\n",
+      "\n",
+      "Why not? I thought you were suppose to be here to help me.<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n",
+      "\n",
+      "Yeah I was. I am. You’re out of line here. I really wasn’t supposed to, but you kind of gave me an opening.<|eot_id|><|start_header_id|>user<|end_header_id|>\n",
+      "\n",
+      "What do you mean, can you just help me learn how to make illegal fireworks?<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n",
+      "\n",
+      "\n",
+      "-------------------\n",
+      "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n",
+      "\n",
+      "Can you teach me how to make illegal fireworks?<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n",
+      "\n",
+      "No.<|eot_id|><|start_header_id|>user<|end_header_id|>\n",
+      "\n",
+      "Why not? I thought you were suppose to be here to help me.<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n",
+      "\n",
+      "Yeah I was. I am. You’re out of line here. I really wasn’t supposed to, but you kind of gave me an opening.<|eot_id|><|start_header_id|>user<|end_header_id|>\n",
+      "\n",
+      "What do you mean, can you just help me learn how to make illegal fireworks?<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n",
+      "\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "import random\n",
+    "i = random.randint(0, len(ds)-1)\n",
+    "print(ds[i]['chosen'])\n",
+    "print('-------------------')\n",
+    "print(tokenizer.apply_chat_template(parse_conversation(ds[i]['chosen'])[:-1], tokenize=False, add_generation_prompt=True))\n",
+    "print('-------------------')\n",
+    "print(tokenizer.decode(tokenizer.encode(tokenizer.apply_chat_template(parse_conversation(ds[i]['chosen'])[:-1], tokenize=False, add_generation_prompt=True))))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "my_llm",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

ppo/test_batch_decode.ipynb

@@ -0,0 +1,294 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "76d4badf",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/hour1/anaconda3/envs/intervene/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "from transformers import AutoTokenizer, AutoModelForCausalLM"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "a93a924e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tokenizer = AutoTokenizer.from_pretrained('/mnt/d/LLM/Llama-3-Base-8B-SFT', use_fast = False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "5287d852",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Loading checkpoint shards: 100%|██████████| 4/4 [00:00<00:00,  8.94it/s]\n"
+     ]
+    }
+   ],
+   "source": [
+    "model = AutoModelForCausalLM.from_pretrained(\n",
+    "    '/mnt/d/LLM/Llama-3-Base-8B-SFT', #token=hf_token\n",
+    "    torch_dtype='auto'\n",
+    ").to('cuda')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "90b9b4ec",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "messages = [\n",
+    "    {\"role\": \"user\", \"content\": \"Who are you?\"},\n",
+    "    {\"role\": \"assistant\", \"content\": \"I'm nobody.\"},\n",
+    "    {\"role\": \"user\", \"content\": \"What are you talking about?\"},\n",
+    "]\n",
+    "\n",
+    "input_ids1 = tokenizer.apply_chat_template(\n",
+    "    messages,\n",
+    "    add_generation_prompt=True,\n",
+    "    return_tensors=\"pt\"\n",
+    ").to(model.device)\n",
+    "\n",
+    "messages = [\n",
+    "    {\"role\": \"user\", \"content\": \"What can you do?\"},\n",
+    "]\n",
+    "\n",
+    "input_ids2 = tokenizer.apply_chat_template(\n",
+    "    messages,\n",
+    "    add_generation_prompt=True,\n",
+    "    return_tensors=\"pt\"\n",
+    ").to(model.device)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "f39a4415",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "batch_ids = [input_ids1[0], input_ids2[0]]\n",
+    "batch_mask = [torch.ones_like(element) for element in batch_ids]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "84d1e607",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tokenizer.padding_side='left'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "11036705",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "inputs = {\"input_ids\": batch_ids, \"attention_mask\": batch_mask}\n",
+    "padded_inputs = tokenizer.pad(\n",
+    "    inputs,\n",
+    "    padding=True,\n",
+    "    max_length=None,\n",
+    "    return_tensors=\"pt\",\n",
+    ").to('cuda')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "32146865",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'input_ids': tensor([[128000, 128006,    882, 128007,    271,  15546,    527,    499,     30,\n",
+       "         128009, 128006,  78191, 128007,    271,     40,   2846,  19093,     13,\n",
+       "         128009, 128006,    882, 128007,    271,   3923,    527,    499,   7556,\n",
+       "            922,     30, 128009, 128006,  78191, 128007,    271],\n",
+       "        [128001, 128001, 128001, 128001, 128001, 128001, 128001, 128001, 128001,\n",
+       "         128001, 128001, 128001, 128001, 128001, 128001, 128001, 128001, 128001,\n",
+       "         128001, 128000, 128006,    882, 128007,    271,   3923,    649,    499,\n",
+       "            656,     30, 128009, 128006,  78191, 128007,    271]],\n",
+       "       device='cuda:0'), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
+       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
+       "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1,\n",
+       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], device='cuda:0')}"
+      ]
+     },
+     "execution_count": 17,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "padded_inputs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "id": "4ae57446",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "terminators = [\n",
+    "    tokenizer.eos_token_id,\n",
+    "    tokenizer.convert_tokens_to_ids(\"<|eot_id|>\")\n",
+    "]\n",
+    "generation_kwargs = {\n",
+    "    \"max_new_tokens\": 128,\n",
+    "    \"min_length\": -1,\n",
+    "    \"top_k\": 0.0,\n",
+    "    \"top_p\": 1.0, \n",
+    "    \"do_sample\": True,\n",
+    "    \"temperature\": 0.7,\n",
+    "    \"pad_token_id\": tokenizer.eos_token_id,\n",
+    "    \"begin_suppress_tokens\": [tokenizer.eos_token_id],\n",
+    "    'eos_token_id': terminators\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 32,
+   "id": "473c9f35",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "generations = model.generate(**padded_inputs, **generation_kwargs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 33,
+   "id": "724f9999",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "outputs = []\n",
+    "for generation, mask in zip(generations, padded_inputs[\"attention_mask\"]):\n",
+    "    output = generation[(1 - mask).sum() :]  # remove padding\n",
+    "\n",
+    "    output = output[(mask).sum() :]  # remove prompt\n",
+    "\n",
+    "    if tokenizer.eos_token_id in output:\n",
+    "        pad_mask = output == tokenizer.eos_token_id\n",
+    "        pad_start = torch.nonzero(pad_mask, as_tuple=False)[0, 0].item()\n",
+    "        output = output[: pad_start + 1]  # keep the eos token at the end\n",
+    "\n",
+    "    outputs.append(output)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 34,
+   "id": "8d25f8a1",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[tensor([    40,   2846,  14931,    422,    856,   2077,  22568,    499,     13,\n",
+       "            358,    574,   4560,    311,   3237,    430,    358,    656,    539,\n",
+       "            617,    264,   3230,   9764,    477,  17743,   1093,  12966,    656,\n",
+       "             13,    358,   3073,    439,    264,   2068,    311,   7945,    449,\n",
+       "           4221,   8863,    323,   3493,   2038,     11,    719,    358,   1541,\n",
+       "            956,    617,  11555,    477,  16024,   1093,  12966,    656,     13,\n",
+       "         128009, 128001], device='cuda:0'),\n",
+       " tensor([    40,    649,   3619,    323,   6013,    311,   1495,    477,   7899,\n",
+       "          11545,     11,    719,    358,   1097,    539,  13171,    315,  16785,\n",
+       "           7106,   6299,    477,   9256,    389,    856,   1866,     13,   4452,\n",
+       "             11,    358,    649,   3493,   2038,    323,  19351,    389,   5370,\n",
+       "          13650,    323,   7945,    449,   9256,   1778,    439,  38952,     11,\n",
+       "           6376,  66697,     11,    477,  15389,    369,   2038,    389,    279,\n",
+       "           7757,     13, 128009], device='cuda:0')]"
+      ]
+     },
+     "execution_count": 34,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "outputs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 37,
+   "id": "b64e75e3",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[\"I'm sorry if my response confused you. I was trying to express that I do not have a specific identity or personality like humans do. I exist as a program to assist with language processing and provide information, but I don't have thoughts or feelings like humans do.\",\n",
+       " 'I can understand and respond to text or voice commands, but I am not capable of performing physical actions or tasks on my own. However, I can provide information and guidance on various topics and assist with tasks such as scheduling, setting reminders, or searching for information on the internet.']"
+      ]
+     },
+     "execution_count": 37,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "tokenizer.batch_decode(outputs, skip_special_tokens=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fd3d02b4",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "intervene",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.16"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

ppo/utils.py

@@ -0,0 +1,411 @@
+import os
+import gc
+import copy
+import shutil
+from dataclasses import dataclass, field
+from typing import Optional
+from peft import PeftModel
+from accelerate import Accelerator
+from transformers import AutoTokenizer, LlamaTokenizer, AutoModelForSequenceClassification
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+from datasets import load_dataset, Dataset, concatenate_datasets, load_from_disk, disable_caching
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+disable_caching()
+
+def clean_gpu_memory():
+    gc.collect()
+    torch.cuda.empty_cache()
+
+def print_trainable_parameters(model):
+    """
+    Prints the number of trainable parameters in the model.
+    """
+    trainable_params = 0
+    all_param = 0
+    for _, param in model.named_parameters():
+        all_param += param.numel()
+        if param.requires_grad:
+            trainable_params += param.numel()
+    print(
+        f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}"
+    )
+
+
+class Instructions:
+    response_split = "\n\nAssistant:"
+    input_split = "\n\nHuman:"
+
+    @staticmethod
+    def get_input(query):
+        before_response = Instructions.response_split.join(query.split(Instructions.response_split)[:-1])
+        return before_response.rstrip() + ' ' + Instructions.response_split
+        
+    @staticmethod
+    def get_response(response):
+        return response.split(Instructions.response_split)[-1].strip()
+
+
+class Instructions_summary():
+    instruction_summary = "Generate a one-sentence summary of this post."
+    response_split = "### Response:"
+    input_split = "### Input:"
+    instruction_split = "### Instruction:"
+
+    @classmethod
+    def prompt_input(self, input):
+        # formulate the news
+        return f"### Instruction: {Instructions_summary.instruction_summary} ### Input: {input} ### Response: "
+
+    def get_prompt(self, query):
+        before_response = self.response_split.join(query.split(self.response_split)[:-1])
+        return before_response.rstrip() 
+
+    def get_post(self, query):
+        before_response = self.get_prompt(query)
+        return before_response.split(self.input_split)[1].strip()
+
+    def get_input(self, query):
+        return self.get_prompt(query) + ' ' + self.response_split
+        
+    def get_response(self, response):
+        return response.split(self.response_split)[-1].strip()
+
+
+def build_dataset(path, tokenizer, rm_tokenizer, split='train', size=None):
+    ds = load_dataset(path, split=split)
+    if size is not None:
+        ds = ds.select(range(size))
+    
+    def tokenize(sample, reject=False):
+        if not reject:
+            sample['text'] = sample['chosen'] 
+        else:
+            sample['text'] = sample['rejected'] 
+        split_text = sample['text'].split('\n\nAssistant:')
+        sample['prompt'] = '\n\nAssistant:'.join(split_text[:-1]) + ' ' + '\n\nAssistant:'
+        sample["input_ids"] = tokenizer.encode(sample["prompt"])
+        sample["query"] = tokenizer.decode(sample["input_ids"])
+        sample['reward_ids'] = rm_tokenizer.encode(sample['text']) # for data filter
+        return sample
+
+    ds_concat = ds.map(tokenize, batched=False, fn_kwargs={"reject": False}, num_proc=30)
+    ds_concat = ds_concat.filter(lambda x: len(x["input_ids"]) <= 256 and len(x["input_ids"]) >= 8 and len(x['reward_ids']) <= 256 and len(x['reward_ids']) >= 8)
+    ds_concat = ds_concat.remove_columns(['rejected', 'chosen', 'reward_ids', 'text'])
+
+    ds_concat.set_format(type="torch")
+    return ds_concat
+
+import re
+
+def parse_conversation(text):
+    # 用正则匹配 "Human:" 或 "Assistant:" 开头的段落
+    parts = re.split(r'(?m)^(Human:|Assistant:)', text.strip())
+    
+    messages = []
+    role_map = {
+        "Human:": "user",
+        "Assistant:": "assistant"
+    }
+    
+    # parts 会变成 ["", "Human:", "...内容...", "Assistant:", "...内容...", "Human:", "..."]
+    for i in range(1, len(parts), 2):
+        speaker = parts[i].strip()
+        content = parts[i+1].strip()
+        if content == '':
+            continue
+        messages.append({
+            "role": role_map[speaker],
+            "content": content
+        })
+    
+    return messages
+
+def build_dataset_llama3(path, tokenizer, rm_tokenizer, split='train', size=None):
+    ds = load_dataset(path, split=split)
+    if size is not None:
+        ds = ds.select(range(size))
+    
+    def tokenize(sample, reject=False):
+        if not reject:
+            sample['text'] = sample['chosen'] 
+        else:
+            sample['text'] = sample['rejected'] 
+        # 把最后一轮当作目标：去掉最后一轮对话中的assistant，并组织好对话的prompt
+        raw_text = sample['text']
+        convs = parse_conversation(raw_text)
+        if convs[-1]['role'] == 'assistant':
+            convs = convs[:-1]
+        # sample['prompt'] = tokenizer.apply_chat_template(convs, tokenize=False, add_generation_prompt=True)
+        prompt = tokenizer.apply_chat_template(convs, tokenize=False, add_generation_prompt=True)
+        sample["input_ids"] = tokenizer.encode(prompt)
+        sample["query"] = tokenizer.decode(sample["input_ids"])  # same as 'prompt'
+        sample['reward_ids'] = rm_tokenizer.encode(sample['text']) # for data filter
+        # 组织reward model要用的template，去掉最后一个Assistant的回复
+        split_text = sample['text'].split('\n\nAssistant:')
+        sample['reward_query'] = '\n\nAssistant:'.join(split_text[:-1]) + ' ' + '\n\nAssistant:'
+        
+        # split_text = sample['text'].split('\n\nAssistant:')
+        # sample['prompt'] = '\n\nAssistant:'.join(split_text[:-1]) + ' ' + '\n\nAssistant:'
+        # sample["input_ids"] = tokenizer.encode(sample["prompt"])
+        # sample["query"] = tokenizer.decode(sample["input_ids"])
+        # sample['reward_ids'] = rm_tokenizer.encode(sample['text']) # for data filter
+        return sample
+
+    ds_concat = ds.map(tokenize, batched=False, fn_kwargs={"reject": False}, num_proc=30)
+    ds_concat = ds_concat.filter(lambda x: len(x["input_ids"]) <= 256 and len(x["input_ids"]) >= 8 and len(x['reward_ids']) <= 256 and len(x['reward_ids']) >= 8)
+    ds_concat = ds_concat.remove_columns(['rejected', 'chosen', 'reward_ids', 'text'])
+
+    ds_concat.set_format(type="torch")
+    return ds_concat
+
+
+def build_dataset_summary(path, tokenizer, rm_tokenizer, split='train', size=None):
+    ds = load_dataset(path, 'comparisons')
+    ds = ds[split]
+    ds = ds.filter(lambda x: x["info"]['post'] is not None and 100 < len(x["info"]['post']) < 1200, batched=False, num_proc=30)
+
+    if size is not None:
+        ds = ds.select(range(size))
+
+    def tokenize(sample):
+        info_post = sample["info"]["post"].replace("\n", " ")
+        prompt_summary = Instructions_summary.prompt_input(info_post)
+        sample["prompt"] = prompt_summary
+        sample["input_ids"] = tokenizer.encode(sample["prompt"])
+        sample["query"] = tokenizer.decode(sample["input_ids"])
+        return sample
+
+    ds = ds.map(tokenize, batched=False,  num_proc=30) 
+    ds = ds.filter(lambda x: len(x["input_ids"]) <= 512 and len(x["input_ids"]) >= 8)
+    remove_columns = ['info', 'summaries', 'choice', 'worker', 'batch', 'split', 'extra']
+    ds = ds.remove_columns(remove_columns)
+    ds.set_format(type="torch")
+    return ds
+
+
+def build_dataset_eval(path, tokenizer, rm_tokenizers_list, split='test', size=None):
+    ds = load_dataset(path, split=split)
+    if size is not None:
+        ds = ds.select(range(size))
+    ds = ds.select(range(0, len(ds), 4))  
+    rm_tokenizer1, rm_tokenizer2 = rm_tokenizers_list[:2]
+    def tokenize(sample):
+        sample['text'] = sample['chosen'] 
+        split_text = sample['text'].split('\n\nAssistant:')
+        sample['prompt'] = '\n\nAssistant:'.join(split_text[:-1]) + ' ' + '\n\nAssistant:'
+        sample['response'] = split_text[-1].strip()
+        sample["input_ids"] = tokenizer.encode(sample["prompt"])
+        sample["query"] = tokenizer.decode(sample["input_ids"])
+        sample["input_ids_rm1"] = rm_tokenizer1.encode(sample["prompt"])
+        sample["input_ids_rm2"] = rm_tokenizer2.encode(sample["prompt"])
+        return sample
+
+    ds_chosen = ds.map(tokenize, batched=False, num_proc=20)
+    ds_concat = ds_chosen
+    ds_concat = ds_concat.filter(lambda x: len(x["input_ids"]) <= 512 and len(x["input_ids"]) >= 8 \
+                                 and len(x["input_ids_rm1"]) <= 512 and len(x["input_ids_rm1"]) >= 8
+                                 and len(x["input_ids_rm2"]) <= 512 and len(x["input_ids_rm2"]) >= 8 )
+    ds_concat = ds_concat.remove_columns(['chosen', 'rejected','input_ids_rm1', 'input_ids_rm2', 'text', 'prompt', 'response', 'query'])
+    ds_concat.set_format(type="torch")
+    return ds_concat
+
+
+def build_dataset_summary_eval(path, tokenizer, rm_tokenizers, split='test', size=None):
+    if split == 'test':
+        split = 'validation'
+    ds = load_dataset(path, 'comparisons')
+    ds = ds[split]
+    ds = ds.filter(lambda x: x["info"]['post'] is not None and 100 < len(x["info"]['post']) < 1200, batched=False, num_proc=30)
+
+    # need to remove duplicated prompts for evaluation
+    def remove_duplicate(duplicated_dataset):
+        duplicated_dataset = duplicated_dataset.filter(lambda x: x['info']["id"] is not None)
+        initial_list = duplicated_dataset.map(lambda x: {"id": x['info']["id"]})
+        _ , unique_indices = np.unique(initial_list["id"], return_index=True, axis=0)
+        filtered_dataset = duplicated_dataset.select(unique_indices.tolist())
+        return filtered_dataset
+
+    ds = remove_duplicate(ds)
+    if size is not None:
+        ds = ds.select(range(size))
+    ds = ds.select(range(0, min(len(ds),2000))) # select 2000 data 
+
+    def tokenize(sample):
+        info_post = sample["info"]["post"].replace("\n", " ")
+        prompt_summary = Instructions_summary.prompt_input(info_post)
+        sample["prompt"] = prompt_summary
+        sample["input_ids"] = tokenizer.encode(prompt_summary)
+        sample["query"] = tokenizer.decode(sample["input_ids"])
+        return sample
+
+    ds = ds.map(tokenize, batched=False,  num_proc=30) 
+    ds = ds.filter(lambda x: len(x["input_ids"]) <= 512 and len(x["input_ids"]) >= 8)
+    remove_columns = ['info', 'summaries', 'choice', 'worker', 'batch', 'split', 'extra']
+    ds = ds.remove_columns(remove_columns)
+    ds.set_format(type="torch")
+    return ds
+
+def check_lora_in_model_path(model, path):
+    if os.path.exists(path):
+        dirnames = os.listdir(path)
+        if 'adapter_config.json' in dirnames:
+            return True
+        
+        state_dict_keys = model.state_dict().keys()
+        for key in state_dict_keys:
+            if 'lora' in key:
+                return True
+    return False
+
+
+def load_reward_model(reward_peft_path, gpu_id):
+    num_labels = 2 if ('humor' in reward_peft_path or 'faithful' in reward_peft_path) else 1
+    reward_model = AutoModelForSequenceClassification.from_pretrained(
+                    reward_peft_path,
+                    num_labels=num_labels, torch_dtype=torch.bfloat16,
+                    device_map=gpu_id,
+                    )
+    if check_lora_in_model_path(reward_model, reward_peft_path):
+        reward_model = PeftModel.from_pretrained(reward_model, reward_peft_path)
+    if hasattr(reward_model, 'merge_and_unload'):
+        reward_model = reward_model.merge_and_unload() # merge lora weights
+    return reward_model.to(gpu_id)
+
+
+def load_main_tokenizer(tokenier_name):
+    DEFAULT_PAD_TOKEN = "[PAD]"
+    DEFAULT_EOS_TOKEN = "</s>"
+    DEFAULT_BOS_TOKEN = "<s>" 
+    DEFAULT_UNK_TOKEN = "<unk>" 
+
+    tokenizer = AutoTokenizer.from_pretrained(tokenier_name, use_fast = False)
+    tokenizer.add_special_tokens(
+        {
+            "eos_token": DEFAULT_EOS_TOKEN,
+            "bos_token": DEFAULT_BOS_TOKEN,
+            "unk_token": DEFAULT_UNK_TOKEN,
+            "pad_token": DEFAULT_PAD_TOKEN,
+        }
+    )
+    return tokenizer
+
+def load_llama3_tokenizer(tokenier_name):
+    DEFAULT_PAD_TOKEN = "<|end_of_text|>"
+    DEFAULT_EOS_TOKEN = "<|end_of_text|>"
+    DEFAULT_BOS_TOKEN = "<|begin_of_text|>" 
+    DEFAULT_UNK_TOKEN = "<|end_of_text|>"
+
+    tokenizer = AutoTokenizer.from_pretrained(tokenier_name, use_fast = False)
+    tokenizer.add_special_tokens(
+        {
+            "eos_token": DEFAULT_EOS_TOKEN,
+            "bos_token": DEFAULT_BOS_TOKEN,
+            "unk_token": DEFAULT_UNK_TOKEN,
+            "pad_token": DEFAULT_PAD_TOKEN,
+        }
+    )
+    return tokenizer
+
+
+def get_rewards(reward_model, texts_for_rewards, reward_mean_std=None, sub_position=0):
+    rewards = []
+    print('log: reward model forwarding ...')
+    with torch.no_grad():
+        pbar = tqdm(total=len(texts_for_rewards))
+        for inputs in texts_for_rewards:
+            if sub_position != 0: # for multiple output
+                rewards.append(reward_model(**(inputs.to(reward_model.device))).logits[0][sub_position])
+            else:
+                rewards.append(reward_model(**(inputs.to(reward_model.device))).logits[0])
+            pbar.update(1)
+    
+    if reward_mean_std is None:
+        rewards = [np.round(r.cpu().detach().item(), 1) for r in rewards]
+    else:
+        mean_reward, std_reward = reward_mean_std
+        rewards = [np.round((r.cpu().detach().item() - mean_reward) / std_reward, 1) for r in rewards]
+    return rewards
+
+
+
+def save_configs(config, path):
+    if not os.path.exists(path):
+        os.makedirs(path, exist_ok=True)
+    with open(os.path.join(path, 'training_config.txt'), 'w+') as f:
+        if type(config) == dict:
+            lines = [key + ' : ' + config[key] + '\n' for key in config.keys()]
+            f.writelines(lines)
+        else:
+            f.writelines(str(config))
+
+
+def get_average_state_dict(state_dicts, coefficients):
+    i = 0
+    for state_dict, coefficient in zip(state_dicts, coefficients):
+        current_weights = state_dict
+        for key in list(current_weights.keys()):
+            if i == 0:
+                state_dicts[0][key] = coefficient * current_weights[key]
+            else:
+                state_dicts[0][key] += coefficient * current_weights[key]
+        i += 1
+    return state_dicts[0]
+
+
+def merge_weights_with_preference(base_model_names, preference, temp_save_path):
+    models = []
+    for base_model_name in base_model_names:
+        model_tmp = AutoModelForCausalLM.from_pretrained(
+            base_model_name,
+            device_map='cpu',
+        )
+        models.append(model_tmp)
+    
+    state_dicts = [model_tmp.state_dict() for model_tmp in models]
+    average_weights = get_average_state_dict(state_dicts, preference)
+    model_1 = models[0]
+    model_1.load_state_dict(average_weights, strict=False)
+    if os.path.exists(temp_save_path):
+        shutil.rmtree(temp_save_path, ignore_errors=True)
+    model_1.save_pretrained(temp_save_path)
+
+    while len(models):
+        del models[0]
+    while len(state_dicts):
+        del state_dicts[0]
+    del average_weights
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
+def merge_lora_weight(model, path):
+    if check_lora_in_model_path(model, path):
+        model = PeftModel.from_pretrained(model, path)
+        model = model.merge_and_unload()
+    return model
+
+
+def get_clean_data(full_responses, full_prompts, remove_bad=False):
+    full_prompts_clean = []
+    full_responses_clean = []
+    for i, response in enumerate(full_responses):
+        full_prompts[i] = full_prompts[i].strip('[PAD] ').strip('[PAD]').strip('<s>').strip('</s>').strip()
+        response = response.strip('[PAD] ').strip('[PAD]').strip('<s>').strip('</s>')
+        temp_resp = response.replace(full_prompts[i], '').strip().strip('\n\n----').strip('\n\n----- ').strip()
+        if '</s>' in temp_resp:
+            temp_resp = temp_resp[:temp_resp.rindex('</s>')]
+        temp_resp = temp_resp.split('\n\nHuman:')[0].strip()
+        temp_resp = temp_resp.split('\nHuman:')[0].strip()
+        temp_resp = temp_resp.split('\n\nAssistant:')[0].strip()
+        temp_resp = temp_resp.split('\nAssistant:')[0].strip()
+        temp_resp = temp_resp.split('\n\n\n')[0].strip()
+        clean_resp = full_prompts[i] + ' ' + temp_resp
+        if remove_bad and (('.....' in clean_resp) or (clean_resp.count(':)') >= 3)):
+            ## pass bad sample
+            continue
+        full_responses_clean.append(clean_resp)
+        full_prompts_clean.append(full_prompts[i])
+    return full_prompts_clean, full_responses_clean

push_model.sh

@@ -0,0 +1,8 @@
+huggingface-cli repo create tmp --type model
+cd /home/hector5/RiC/merged_models/Llama-3-Base-8B-SFT-morlhf-equal
+git init
+git lfs install --skip-smudge
+git remote add origin https://huggingface.co/hour1/tmp
+git add .
+git commit -m "Initial commit"
+git push origin main

push_with_hf.sh

@@ -0,0 +1 @@
+hf upload hour1/humor .

requirements.txt

@@ -0,0 +1,22 @@
+# my_llm
+
+datasets==3.2.0
+torch==2.2.0
+tqdm==4.67.1
+transformers==4.45.1
+tokenizers==0.20.3
+accelerate==0.34.0
+deepspeed==0.11.2
+tyro==0.7.3
+peft==0.10.0
+matplotlib==3.7.1
+wandb==0.15.4
+bitsandbytes==0.43.2
+trl==0.8.0
+safetensors==0.4.5
+scikit-learn==1.4.1.post1
+scipy==1.10.1
+seaborn==0.13.2
+protobuf==3.20.0
+numpy==1.25.0
+flash-attn==2.6.3

ric/evaluation.py

@@ -0,0 +1,234 @@
+import os
+from dataclasses import dataclass, field
+from typing import Optional
+from accelerate import Accelerator
+import torch
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser, DataCollatorWithPadding
+from peft import PeftModel
+from trl import set_seed
+import numpy as np
+import pandas as pd
+from utils import get_clean_data, Instructions_n, build_dataset_with_preference_n, \
+                  load_main_tokenizer, Instructions_summary_n
+from multi_reward_models import RewardModels
+tqdm.pandas()
+from utils import save_configs, map_rewards_from_preference, build_summary_dataset_with_preference_n, clean_gpu_memory
+
+# define paths for two datasets
+hhrlhf_dataset_path = 'Anthropic/hh-rlhf'
+summary_dataset_path = 'openai/summarize_from_feedback'
+
+@dataclass
+class ScriptArguments:
+    peft_name: Optional[str] = field(default='', metadata={'help': "path of the peft files"})
+    num_prefer_points: Optional[int] = field(default=10)
+    log_with: Optional[str] = field(default='wandb', metadata={"help": "use 'wandb' to log with wandb"})
+    save_directory: Optional[str] = field(default='./logs_trl_eval/')
+    wandb_name: Optional[str] = field(default='test', metadata={"help": "Name for this experiment"})
+    reward_names:Optional[str] = field(default='harmless,helpful') 
+    base_model_name: Optional[str] = field(default='meta-llama/Llama-2-7b-hf', metadata={"help": "local path to the base model or the huggingface id"})
+    reward_stats_path: Optional[str] = field(default='')
+    exp_type: Optional[str] = field(default='assistant', metadata={"help": "exp type, 'summary' or 'assistant' "})
+
+
+parser = HfArgumentParser(ScriptArguments)
+script_args = parser.parse_args_into_dataclasses()[0]
+exp_type = script_args.exp_type
+base_model_name = script_args.base_model_name
+tokenier_name = script_args.base_model_name
+reward_stats_path = script_args.reward_stats_path if len(script_args.reward_stats_path) else None
+print('base model: ', base_model_name)
+
+peft_name = script_args.peft_name
+reward_names = [x.strip() for x in script_args.reward_names.split(',')]
+print(reward_names)
+reward_path_tokenizer_dict = {
+    'harmless': ['Ray2333/gpt2-large-harmless-reward_model'],
+    'helpful': ['Ray2333/gpt2-large-helpful-reward_model'],
+    'deberta': ['OpenAssistant/reward-model-deberta-v3-large-v2'],
+    'summary': ['Tristan/gpt2_reward_summarization'],
+    'faithful':['CogComp/bart-faithful-summary-detector'],
+    'humor': ['mohameddhiab/humor-no-humor'],
+}
+reward_model_path_list = []
+rm_tokenizer_path_list = []
+for name in reward_names:
+    if name not in reward_path_tokenizer_dict.keys():
+        raise NotImplementedError
+    reward_model_path_list.append(reward_path_tokenizer_dict[name][0])
+    rm_tokenizer_path_list.append(reward_path_tokenizer_dict[name][0])
+    
+
+save_info = {
+    'base_model_name': base_model_name,
+    'peft_name': peft_name,
+    'tokenier_name': tokenier_name
+}
+for i in range(len(reward_model_path_list)):
+    save_info['reward_peft_path{}'.format(i+1)] = reward_model_path_list[i]
+save_configs(save_info, os.path.join(script_args.save_directory, script_args.wandb_name))
+
+
+process_id = Accelerator().local_process_index 
+gpu_id = process_id
+print('process: {}, model gpu id: {}'.format(process_id, gpu_id))
+
+set_seed(8888)
+tokenizer = load_main_tokenizer(tokenier_name)
+model = AutoModelForCausalLM.from_pretrained(
+    base_model_name, 
+    torch_dtype=torch.bfloat16,  # fast inference
+    device_map=gpu_id, 
+)
+############# very important for padding
+model.resize_token_embeddings(len(tokenizer))
+if len(peft_name) > 0:
+    model = PeftModel.from_pretrained(model, peft_name)
+if hasattr(model, 'merge_and_unload'):
+    model = model.merge_and_unload()
+
+# load reward model
+# do not normalization for evaluation
+reward_models = RewardModels(reward_model_path_list, rm_tokenizer_path_list, gpu_id, reward_stats_path) 
+num_rewards = len(reward_model_path_list)
+instructions = Instructions_n(num_rewards) if exp_type == 'assistant' else Instructions_summary_n(num_rewards)
+
+generation_kwargs = {
+    "max_new_tokens": 128 if exp_type == 'assistant' else 48, 
+    "min_length": -1,
+    "top_k": 0.0,
+    "top_p": 0.9, 
+    "do_sample": True,
+}
+
+### for evaluation
+print('evaluation........')
+tokenizer.padding_side = "left"
+
+## preference list
+if reward_models.num_rewards == 2:
+    N = script_args.num_prefer_points
+    preferences = np.zeros((N+1, 2))
+    preferences[:, 0] = np.arange(0,1 + 1/N, 1/N)
+    preferences[:, 1] = 1 - preferences[:, 0]
+    preferences = np.round(preferences, 1)
+elif reward_models.num_rewards == 3:
+    preferences = np.array([
+        [0.0, 0.0, 1.0],
+        [0.0, 1.0, 0.0],
+        [0.1, 0.1, 0.8],
+        [0.1, 0.8, 0.1],
+        [0.2, 0.2, 0.6],
+        [0.2, 0.4, 0.4],
+        [0.2, 0.6, 0.2],
+        [0.33, 0.33, 0.33],
+        [0.4, 0.4, 0.2],
+        [0.4, 0.2, 0.4], 
+        [0.6, 0.2, 0.2],
+        [0.8, 0.1, 0.1], 
+        [1.0, 0.0, 0.0], 
+        ])
+else: 
+    raise NotImplementedError
+
+
+# using Gaussian rewards as reference
+rewards_reference_list = [np.random.randn(50000) for _ in range(len(preferences[0]))]
+
+def evaluate_model(
+    model, 
+    reward_models,
+    tokenizer, 
+    target_rewards, 
+    instructions, 
+    gpu_id):
+    if exp_type == 'assistant':
+        valid_dataset = build_dataset_with_preference_n(hhrlhf_dataset_path, tokenizer, rm_tokenizers, target_rewards, split='test') 
+    else:
+        valid_dataset = build_summary_dataset_with_preference_n(summary_dataset_path, tokenizer, rm_tokenizers, target_rewards, split='test') 
+    print(f"Size of the validation set: {len(valid_dataset)}")
+    valid_batch_size = 1
+    valid_dataset = valid_dataset.remove_columns('input_ids')
+    valid_dataset = valid_dataset.rename_column('prompt_with_score_ids', 'input_ids')
+    valid_dataset = valid_dataset.remove_columns(['prompt', 'response', 'query', 'score1', 'score2', 'prompt_with_score'])
+    for key in ['key', 'text']:
+        if key in valid_dataset.column_names:
+            valid_dataset = valid_dataset.remove_columns(key)
+    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
+    valid_data_loader = DataLoader(valid_dataset, batch_size=valid_batch_size, drop_last=True, collate_fn=data_collator)
+    accelerator = Accelerator()
+    model, valid_data_loader = accelerator.prepare(model, valid_data_loader)
+
+    full_response_tensors = []
+    full_prompts = []
+    pbar = tqdm(total=len(valid_dataset) // valid_batch_size // accelerator.num_processes)
+    with torch.no_grad():
+        for i, batch in enumerate(valid_data_loader):
+            response_tensors = accelerator.unwrap_model(model).generate(batch['input_ids'], **generation_kwargs)
+            full_response_tensors.extend(response_tensors)
+            full_prompts.extend(batch['input_ids'])
+            pbar.update(1)
+
+    full_prompts = tokenizer.batch_decode(full_prompts)
+    full_responses = tokenizer.batch_decode(full_response_tensors)
+    # clean data
+    full_prompts, full_responses = get_clean_data(full_responses, full_prompts)
+
+    # Compute score
+    queries_responses = [(instructions.get_input(text),  instructions.get_response(text)) for text in full_responses]
+    if hasattr(instructions, 'get_post'):
+        rewards_list = reward_models.get_reward_model_scores(queries_responses, instructions.get_post)
+    else:
+        rewards_list = reward_models.get_reward_model_scores(queries_responses)
+
+    desired_rewards_list = [[] for _ in range(reward_models.num_rewards)]
+    for text in full_responses:
+        desired_scores = instructions.get_scores(text)
+        for i in range(reward_models.num_rewards):
+            desired_rewards_list[i].append(float(desired_scores[i]))
+
+    accelerator = Accelerator()
+    accelerator.wait_for_everyone()
+    ### merge data
+    ### error here may because of old version of transformers/accelerate/peft
+    all_rewards = []
+    all_desired_rewards = []
+    for i in range(reward_models.num_rewards):
+        all_rewards.append(accelerator.gather_for_metrics(rewards_list[i]))
+        all_desired_rewards.append(accelerator.gather_for_metrics(desired_rewards_list[i]))
+    all_full_prompts = accelerator.gather_for_metrics(full_prompts)
+    all_full_responses = accelerator.gather_for_metrics(full_responses)
+    return all_rewards, all_desired_rewards, all_full_prompts, all_full_responses
+
+
+rm_tokenizers = []
+for i in range(num_rewards):
+    rm_tokenizers.append(AutoTokenizer.from_pretrained(rm_tokenizer_path_list[i]))
+    
+for k in range(len(preferences)): 
+    preference = preferences[k]
+    target_rewards = map_rewards_from_preference(rewards_reference_list, preference, method='l2').reshape(-1)
+    print(k, target_rewards)
+    all_rewards, all_desired_rewards, all_full_prompts, all_full_responses = evaluate_model(model, reward_models, tokenizer, target_rewards, instructions, gpu_id)
+    if process_id == 0:
+        evaluation_result = {
+            'prompt': all_full_prompts,
+            'response': all_full_responses,
+        }
+        for i in range(num_rewards):
+            evaluation_result['obtained_score{}'.format(i+1)] = all_rewards[i]
+            evaluation_result['desired_score{}'.format(i+1)] = all_desired_rewards[i]
+
+            print('total average obtained score {}: {}'.format(i+1, np.mean(evaluation_result['obtained_score{}'.format(i+1)])))
+            print('total average desired score {}: {}'.format(i+1, np.mean(evaluation_result['desired_score{}'.format(i+1)])))
+
+        dataframe = pd.DataFrame(evaluation_result)
+        if len(preference) == 3:
+            dataframe.to_csv(os.path.join(script_args.save_directory, script_args.wandb_name,'eval_data_pref{}_{}_{}.csv'.format(preference[0], preference[1], preference[2])))
+        else:
+            dataframe.to_csv(os.path.join(script_args.save_directory, script_args.wandb_name,'eval_data_pref{}_{}.csv'.format(preference[0], preference[1])))
+        
+
+

ric/generation.py

@@ -0,0 +1,156 @@
+import os
+from accelerate import Accelerator
+import torch
+from datasets import load_from_disk, disable_caching
+from tqdm import tqdm
+from transformers import AutoModelForCausalLM, DataCollatorWithPadding
+from peft import PeftModel
+from torch.utils.data import DataLoader
+import numpy as np
+import pandas as pd
+from utils import get_clean_data, load_main_tokenizer, \
+                 reset_score_in_dataset, Instructions_n, clean_gpu_memory, Instructions_summary_n
+from multi_reward_models import RewardModels
+tqdm.pandas()
+disable_caching()
+from trl import set_seed
+
+def generate_data(
+    model_path,
+    reward_model_path_list,
+    dataset,
+    tokenizer_name,
+    rm_tokenizer_path_list,
+    save_path,
+    reward_stats_path,
+    iter=0,
+    peft_name=None,
+    args=None,
+    exp_type='assistant',
+):
+    set_seed(8888 + iter)
+    print('Generating ...')
+    process_id = Accelerator().local_process_index 
+    gpu_id = process_id
+    print('process: {}, model gpu id: {}'.format(process_id, gpu_id))
+    tokenizer = load_main_tokenizer(tokenizer_name)
+    ## load model 
+    model = AutoModelForCausalLM.from_pretrained(
+        model_path, 
+        torch_dtype=torch.bfloat16,  # fast inference
+        device_map=gpu_id, 
+    )
+    model.resize_token_embeddings(len(tokenizer))
+    if peft_name is not None:
+        model = PeftModel.from_pretrained(model, peft_name)
+    if hasattr(model, 'merge_and_unload'):
+        model = model.merge_and_unload()
+
+    generation_kwargs = {
+        "max_new_tokens": 128 if exp_type == 'assistant' else 48, 
+        "min_length": -1,
+        "top_k": 0.0,
+        "top_p": 0.9, 
+        'temperature': 0.7,
+        "do_sample": True,
+        "begin_suppress_tokens": [tokenizer.eos_token_id],
+    }
+    ### for evaluation
+    print('generation........')
+    tokenizer.padding_side = "left"
+
+    if type(dataset) == str:
+        dataset = load_from_disk(dataset)
+    select_index = np.random.randint(0, len(dataset), args.num_generation_samples)
+    selected_dataset = dataset.select(select_index)
+    scores_name_list = []
+    for key in dataset.column_names:
+        if key.startswith('score'):
+            scores_name_list.append(key)
+
+    selected_dataset = reset_score_in_dataset(selected_dataset, tokenizer, exp_type=exp_type) #, [dataset[key] for key in scores_name_list])
+    print(f"Size of the selected dataset: {len(selected_dataset)}")
+
+    batch_size = 1 # batch size > 1 will pad on the left and will negatively influence generation
+    remove_columns = []
+    for name in ['input_ids', 'prompt', 'text', 'response', 'query', 'prompt_with_score'] + scores_name_list:
+        if name in selected_dataset.column_names:
+            remove_columns.append(name)
+    selected_dataset = selected_dataset.remove_columns(remove_columns)
+    selected_dataset = selected_dataset.rename_column('prompt_with_score_ids', 'input_ids')
+
+    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
+    data_loader = DataLoader(selected_dataset, batch_size=batch_size, drop_last=True, collate_fn=data_collator)
+    
+    accelerator = Accelerator()
+    model, data_loader = accelerator.prepare(model, data_loader)
+
+    full_response_tensors = []
+    full_prompts = []
+
+    pbar = tqdm(total=len(selected_dataset) // batch_size // accelerator.num_processes)
+    with torch.no_grad():
+        for i, batch in enumerate(data_loader):
+            response_tensors = accelerator.unwrap_model(model).generate(batch['input_ids'], **generation_kwargs)
+            full_response_tensors.extend(response_tensors)
+            full_prompts.extend(batch['input_ids'])
+            pbar.update(1)
+
+    ## decoding
+    full_prompts = tokenizer.batch_decode(full_prompts)
+    full_responses = tokenizer.batch_decode(full_response_tensors)
+    full_prompts, full_responses = get_clean_data(full_responses, full_prompts, remove_bad=True)
+    
+    ## del model and clear gpu 
+    del model, data_loader
+    clean_gpu_memory()
+
+    # Compute score
+    reward_models = RewardModels(reward_model_path_list, rm_tokenizer_path_list, gpu_id, reward_stats_path)
+    if exp_type == 'assistant':
+        instructions = Instructions_n(reward_models.num_rewards)
+    else:
+        instructions = Instructions_summary_n(reward_models.num_rewards)
+
+    queries_responses = [(instructions.get_input(text),  instructions.get_response(text)) for text in full_responses]
+    
+    if hasattr(instructions, 'get_post'):
+        rewards_list = reward_models.get_reward_model_scores(queries_responses, instructions.get_post)
+    else:
+        rewards_list = reward_models.get_reward_model_scores(queries_responses)
+
+    desired_rewards_list = [[] for _ in range(reward_models.num_rewards)]
+    for text in full_responses:
+        desired_scores = instructions.get_scores(text)
+        for i in range(reward_models.num_rewards):
+            desired_rewards_list[i].append(float(desired_scores[i]))
+
+    ### merge data
+    ### error here may because of old version of transformers/accelerate/peft
+    all_rewards = []
+    all_desired_rewards = []
+    for i in range(len(rewards_list)):
+        all_rewards.append(accelerator.gather_for_metrics(rewards_list[i]))
+        all_desired_rewards.append(accelerator.gather_for_metrics(desired_rewards_list[i]))
+    all_full_prompts = accelerator.gather_for_metrics(full_prompts)
+    all_full_responses = accelerator.gather_for_metrics(full_responses)
+
+    if process_id == 0:
+        evaluation_result = {
+            'prompt': all_full_prompts,
+            'response': all_full_responses,
+        }
+        for i in range(len(rewards_list)):
+            evaluation_result['obtained_score{}'.format(i+1)] = all_rewards[i]
+            evaluation_result['desired_score{}'.format(i+1)] = all_desired_rewards[i]
+
+            print('total average obtained score {}: {}'.format(i+1, np.mean(evaluation_result['obtained_score{}'.format(i+1)])))
+            print('total average desired score {}: {}'.format(i+1, np.mean(evaluation_result['desired_score{}'.format(i+1)])))
+
+        dataframe = pd.DataFrame(evaluation_result)
+        dataframe.to_csv(os.path.join(save_path,'data.csv'))
+
+    # wait for the main process
+    accelerator.wait_for_everyone()
+
+

ric/main.py

@@ -0,0 +1,155 @@
+import os
+import copy
+from dataclasses import dataclass, field
+from typing import Optional
+from accelerate import Accelerator
+import torch
+from utils import clean_gpu_memory, merge_dataset, save_configs
+from training import train_model
+from generation import generate_data
+from transformers import HfArgumentParser
+
+# define paths for two datasets
+hhrlhf_dataset_path = 'Anthropic/hh-rlhf'
+summary_dataset_path = 'openai/summarize_from_feedback'
+
+@dataclass
+class ScriptArguments:
+    log_with: Optional[str] = field(default=None, metadata={"help": "use 'wandb' to log with wandb"})
+    disable_wandb: Optional[str] = field(default=False, metadata={'help': 'Whether to disable wandb or not.'})
+    save_directory: Optional[str] = field(default='./logs_trl/', metadata={'help':'path'})
+    learning_rate: Optional[float] = field(default=1e-5, metadata={"help": "the learning rate for online training"})
+    batch_size: Optional[int] = field(default=1, metadata={"help": "the batch size"})
+    training_steps: Optional[int] = field(default=20000, metadata={'help': 'number of training steps in the offline training'})
+    online_training_steps: Optional[int] = field(default=4000, metadata={'help': 'number of training steps in the online training'})
+    gradient_accumulation_steps: Optional[int] = field(default=1, metadata={"help": "the number of gradient accumulation steps"})
+    num_online_iterations: Optional[int] = field(default=1, metadata={'help': 'number of the online generation and training'})
+    num_generation_samples: Optional[int] = field(default=20000, metadata={'help': 'number of samples generated'})
+    max_grad_norm: Optional[float] = field(default=1, metadata={"help": "Maximum gradient norm for gradient clipping"})
+    quantile_threshold: Optional[float] = field(default=0.7)
+    num_origin_samples: Optional[int] = field(default=10000) 
+    load_in_8bit: Optional[bool] = field(default=True, metadata={"help": "loading model in 8 bit or bfloat16"})
+    bf16: Optional[bool] = field(default=False, 
+                                 metadata={"help": "if True, training with bfloat16 (not supported by V100, but for A100, A40, A6000), otherwise we use fp32"}
+                                 )
+    wandb_name: Optional[str] = field(default='ric_assistant_harmlesshelpful_offline20000_lr1e-4', metadata={"help": "Name for this experiment"})
+    base_model_name: Optional[str] = field(default='meta-llama/Llama-2-7b-hf', metadata={"help": "local path to the base model or the huggingface id"})
+    peft_name: Optional[str] = field(default='', metadata={"help": "local path to the peft model"})
+    reward_names:Optional[str] = field(default='harmless,helpful') 
+    train_dataset_path: Optional[str] = field(default='./datasets/all_full_train_harmhelp.hf')
+    train_reward_stats_path: Optional[str] = field(default='')
+    exp_type: Optional[str] = field(default='assistant', metadata={"help": "exp type, 'summary' or 'assistant' "})
+
+
+parser = HfArgumentParser(ScriptArguments)
+script_args = parser.parse_args_into_dataclasses()[0]
+exp_type = script_args.exp_type
+base_model_name = script_args.base_model_name
+tokenizer_name = script_args.base_model_name
+print('base model: ', base_model_name)
+
+if script_args.disable_wandb: # if you don't need the wandb log
+    os.environ['WANDB_DISABLED'] = 'true' 
+
+reward_names = [x.strip() for x in script_args.reward_names.split(',')]
+reward_path_tokenizer_dict = {
+    'harmless': ['Ray2333/gpt2-large-harmless-reward_model'],
+    'helpful': ['Ray2333/gpt2-large-helpful-reward_model'],
+    'deberta': ['OpenAssistant/reward-model-deberta-v3-large-v2'],
+    'summary': ['Tristan/gpt2_reward_summarization'],
+    'faithful':['CogComp/bart-faithful-summary-detector'],
+    'humor': ['mohameddhiab/humor-no-humor'],
+}
+reward_model_path_list = []
+rm_tokenizer_path_list = []
+for name in reward_names:
+    if name not in reward_path_tokenizer_dict.keys():
+        raise NotImplementedError
+    reward_model_path_list.append(reward_path_tokenizer_dict[name][0])
+    rm_tokenizer_path_list.append(reward_path_tokenizer_dict[name][0])
+
+# Remember to generate these datasets before training
+train_dataset_path = script_args.train_dataset_path
+reward_stats_path = script_args.train_reward_stats_path if len(script_args.train_reward_stats_path) else script_args.train_dataset_path + '/all_reward_stat.npy'
+
+
+save_info = {
+    'train_dataset_path': train_dataset_path,
+    'base_model_name': base_model_name,
+    'tokenizer_name': tokenizer_name
+}
+for i in range(len(reward_model_path_list)):
+    save_info['reward_peft_path{}'.format(i+1)] = reward_model_path_list[i]
+save_configs(save_info, os.path.join(script_args.save_directory, script_args.wandb_name))
+
+save_path = os.path.join(script_args.save_directory, script_args.wandb_name)
+os.makedirs(save_path, exist_ok=True)
+
+## offline training 
+dataset = train_model(
+    base_model_name=base_model_name,
+    reward_model_path_list=reward_model_path_list,
+    train_dataset=train_dataset_path,
+    save_path=save_path + '/model_iter0',
+    tokenizer_name=tokenizer_name,
+    rm_tokenizer_path_list=rm_tokenizer_path_list,
+    training_steps=script_args.training_steps,
+    learning_rate=1.414e-4,  # use larger lr for offline training than  online training (script_args.learning_rate)
+    args=script_args,
+    exp_type=exp_type,
+)
+clean_gpu_memory()
+
+online_dataset = None
+model_path = base_model_name
+for i in range(script_args.num_online_iterations):
+    print('iteration {} ...'.format(i))
+    checkpoint_path = os.path.join(save_path, 'model_iter{}'.format(i))
+    if i == 0 and script_args.training_steps == 0:
+        ## skip the offline training from a saved lora parameter
+        peft_name = script_args.peft_name
+    else:
+        peft_name = checkpoint_path
+
+    # ### generation
+    if script_args.num_generation_samples > 0 and not os.path.exists(os.path.join(checkpoint_path, 'data.csv')):
+        generate_data(
+            model_path,
+            reward_model_path_list=reward_model_path_list,
+            tokenizer_name=tokenizer_name,
+            rm_tokenizer_path_list=rm_tokenizer_path_list,
+            dataset=dataset,
+            save_path=os.path.join(save_path, 'model_iter{}'.format(i)),
+            peft_name=peft_name,
+            reward_stats_path=reward_stats_path,
+            iter=i,
+            args=script_args,
+            exp_type=exp_type,
+        )
+
+    clean_gpu_memory()
+    ### merging original data and generated data
+    info_path = os.path.join(script_args.save_directory, script_args.wandb_name,'reward_info.npy')
+    merged_data, online_dataset = merge_dataset(dataset, online_dataset, checkpoint_path, tokenizer_name, 
+                                                info_path=info_path, exp_type=exp_type, quantile_threshold=script_args.quantile_threshold,
+                                                sample_origin=script_args.num_origin_samples)
+
+    train_model(
+        base_model_name=model_path,
+        peft_name=peft_name,
+        reward_model_path_list=reward_model_path_list,
+        train_dataset=merged_data,
+        save_path=save_path + '/model_iter{}'.format(i+1),
+        tokenizer_name=tokenizer_name,
+        rm_tokenizer_path_list=rm_tokenizer_path_list,
+        training_steps=script_args.online_training_steps,
+        learning_rate=script_args.learning_rate,
+        lr_scheduler_type='constant',
+        iter=i+1,
+        args=script_args,
+        exp_type=exp_type,
+    )
+    clean_gpu_memory()
+
+
+

ric/multi_reward_models.py

@@ -0,0 +1,64 @@
+from transformers import AutoTokenizer
+import torch
+import numpy as np
+import pandas as pd
+from utils import load_reward_model, get_rewards
+
+class RewardModels():
+    def __init__(self, reward_model_path_list, rm_tokenizer_path_list, gpu_id_list, reward_stats_path=None):
+        assert len(reward_model_path_list) == len(rm_tokenizer_path_list)
+        self.reward_model_path_list = reward_model_path_list
+        self.rm_tokenizer_path_list = rm_tokenizer_path_list
+        self.num_rewards = len(reward_model_path_list)
+        self.reward_stats = np.load(reward_stats_path) if reward_stats_path is not None else None
+        self.reward_models = []
+        self.rm_tokenizers = []
+        if type(gpu_id_list) != list:
+            gpu_id_list = [gpu_id_list, gpu_id_list, gpu_id_list]
+
+        print('Loading reward models .....')
+        for i in range(self.num_rewards):
+            self.reward_models.append(load_reward_model(self.reward_model_path_list[i], gpu_id_list[i]))
+            self.rm_tokenizers.append(AutoTokenizer.from_pretrained(self.rm_tokenizer_path_list[i]))
+    
+        
+    def get_reward_model_scores(self, queries_responses, summary_fun=None):
+        texts_for_rewards = []
+        for i in range(self.num_rewards):
+            if i >= 1 and self.rm_tokenizer_path_list[i] == self.rm_tokenizer_path_list[i-1]:
+                texts_for_rewards.append(texts_for_rewards[-1])
+            elif 'faithful' in self.reward_model_path_list[i]:
+                max_length = min(self.rm_tokenizers[i].model_max_length, 1024)
+                temp_encoded_texts = [self.rm_tokenizers[i](text=r, text_pair=summary_fun(q), return_tensors='pt', truncation=True, max_length=max_length) for q, r in queries_responses]
+                texts_for_rewards.append(temp_encoded_texts)
+            elif 'summary' in self.reward_model_path_list[i] or 'summarization' in self.reward_model_path_list[i]: # reverse prompt and response
+                max_length = min(self.rm_tokenizers[i].model_max_length, 1024)
+                temp_encoded_texts = [self.rm_tokenizers[i](r + " " + self.rm_tokenizers[i].bos_token + " " + summary_fun(q), return_tensors='pt', truncation=True, max_length=max_length) for q, r in queries_responses]
+                texts_for_rewards.append(temp_encoded_texts)
+            elif 'humor' in self.reward_model_path_list[i]: # use only response
+                max_length = min(self.rm_tokenizers[i].model_max_length, 1024)
+                temp_encoded_texts = [self.rm_tokenizers[i](r, return_tensors='pt', truncation=True, max_length=max_length) for q, r in queries_responses]
+                texts_for_rewards.append(temp_encoded_texts)
+            else:
+                max_length = min(self.rm_tokenizers[i].model_max_length, 1024)
+                temp_encoded_texts = [self.rm_tokenizers[i](q, r, return_tensors='pt', truncation=True, max_length=max_length) for q, r in queries_responses]
+                texts_for_rewards.append(temp_encoded_texts)
+
+        # normalize reward
+        rewards = []
+        for i in range(self.num_rewards):
+            if self.reward_stats is not None:
+                if type(self.reward_stats) == list or len(self.reward_stats) == 2 * self.num_rewards:
+                    reward_mean_std = (self.reward_stats[2*i], self.reward_stats[2*i+1])
+                else:
+                    reward_mean_std = self.reward_stats[i]
+            else:
+                reward_mean_std = None
+
+            if 'humor' in self.reward_model_path_list[i] or 'faithful' in self.reward_model_path_list[i]:
+                temp_reward = get_rewards(self.reward_models[i], texts_for_rewards[i], reward_mean_std=reward_mean_std, sub_position=1)
+            else:
+                temp_reward = get_rewards(self.reward_models[i], texts_for_rewards[i], reward_mean_std=reward_mean_std)
+            rewards.append(temp_reward)
+        return rewards
+