Upload folder using huggingface_hub

.gitignore

@@ -0,0 +1,4 @@
+wandb
+src/__pycache__
+scripts/run_orpo.sh
+src/accelerate/fsdp.yaml

README.md

@@ -0,0 +1,53 @@
+# **ORPO**
+
+This is the official repository for <a class="link" href="https://arxiv.org/abs/2403.07691">**Reference-free Monolithic Preference Optimization with Odds Ratio**</a>. The detailed results in the paper can be found in:
+- [Open LLM Leaderboard](https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard?query=kaist-ai%2Fmistral-orpo-beta)
+- [AlpacaEval](#alpacaeval)
+- [MT-Bench](#mt-bench)
+- [IFEval](#ifeval)
+
+&nbsp;
+
+### **`Model Checkpoints`**
+
+Our models trained with ORPO can be found in:
+
+- [X] **Mistral-ORPO-⍺**: 🤗 <a class="link" href="https://huggingface.co/kaist-ai/mistral-orpo-alpha">kaist-ai/mistral-orpo-alpha</a>
+- [X] **Mistral-ORPO-β**: 🤗 <a class="link" href="https://huggingface.co/kaist-ai/mistral-orpo-beta">kaist-ai/mistral-orpo-beta</a>
+
+And the corresponding logs for the average log probabilities of chosen/rejected responses during training are reported in:
+
+- [X] **Mistral-ORPO-⍺**: <a class="link" href="https://wandb.ai/jiwooya1000/PREF/reports/Mistral-ORPO-7B-Training-Log--Vmlldzo3MTE1NzE0?accessToken=rms6o4mg5vo3feu1bvbpk632m4cspe19l0u1p4he3othx5bgean82chn9neiile6">Wandb Report for Mistral-ORPO-⍺</a>
+- [X] **Mistral-ORPO-β**: <a class="link" href="https://wandb.ai/jiwooya1000/PREF/reports/Mistral-ORPO-7B-Training-Log--Vmlldzo3MTE3MzMy?accessToken=dij4qbp6dcrofsanzbgobjsne9el8a2zkly2u5z82rxisd4wiwv1rhp0s2dub11e">Wandb Report for Mistral-ORPO-β</a>
+
+&nbsp;
+
+### **`AlpacaEval`**
+
+<figure>
+  <img class="png" src="/assets/img/alpaca_blog.png" alt="Description of the image">
+  <figcaption><b>Figure 1.</b> AlpacaEval 2.0 score for the models trained with different alignment methods.</figcaption>
+</figure>
+
+&nbsp;
+
+### **`MT-Bench`**
+
+<figure>
+  <img class="png" src="/assets/img/mtbench_hf.png" alt="Description of the image">
+  <figcaption><b>Figure 2.</b> MT-Bench result by category.</figcaption>
+</figure>
+
+&nbsp;
+
+### **`IFEval`**
+
+IFEval scores are measured with <a class="link" href="https://github.com/EleutherAI/lm-evaluation-harness">EleutherAI/lm-evaluation-harness</a> by applying the chat template. The scores for Llama-2-Chat (70B), Zephyr-β (7B), and Mixtral-8X7B-Instruct-v0.1 are originally reported in <a class="link" href="https://twitter.com/wiskojo/status/1739767758462877823">this tweet</a>.
+
+| **Model Type**     | **Prompt-Strict** | **Prompt-Loose** | **Inst-Strict** | **Inst-Loose** |
+|--------------------|:-----------------:|:----------------:|:---------------:|----------------|
+| **Llama-2-Chat (70B)** |       0.4436      |      0.5342      |      0.5468     |     0.6319     |
+| **Zephyr-β (7B)** |       0.4233      |      0.4547      |      0.5492     |     0.5767     |
+| **Mixtral-8X7B-Instruct-v0.1** |       0.5213      |      **0.5712**      |      0.6343     |     **0.6823**     |
+| **Mistral-ORPO-⍺ (7B)** |       0.5009      |      0.5083      |      0.5995     |     0.6163     |
+| **Mistral-ORPO-β (7B)** |       **0.5287**      |      0.5564      |      **0.6355**     |     0.6619     |

main.py

@@ -0,0 +1,200 @@
+
+import os
+import time
+import wandb
+import torch
+import argparse
+from datasets import load_dataset
+from typing import List, Dict, Union
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    TrainingArguments,
+    DataCollatorForLanguageModeling
+)
+
+from src.args import default_args
+from src.orpo_trainer import ORPOTrainer
+from src.utils import preprocess_logits_for_metrics, dataset_split_selector
+
+class ORPO(object):
+    def __init__(self, args) -> None:
+        self.start = time.gmtime()
+        self.args = args
+
+        # Load Tokenizer
+        print(">>> 1. Loading Tokenizer")
+        self.tokenizer = AutoTokenizer.from_pretrained(self.args.model_name, cache_dir=self.args.cache_dir)
+        if self.tokenizer.chat_template is None:
+            self.tokenizer.chat_template = "{% for message in messages %}\n{% if message['role'] == 'user' %}\n{{ '<|user|>\n' + message['content'] + eos_token }}\n{% elif message['role'] == 'system' %}\n{{ '<|system|>\n' + message['content'] + eos_token }}\n{% elif message['role'] == 'assistant' %}\n{{ '<|assistant|>\n'  + message['content'] + eos_token }}\n{% endif %}\n{% if loop.last and add_generation_prompt %}\n{{ '<|assistant|>' }}\n{% endif %}\n{% endfor %}"
+            print("     1-1. Chat Template Applied (<|user|> <|assistant|>)")
+        else:
+            pass
+        self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+
+        # Load Model
+        print(">>> 2. Loading Model")
+        if self.args.flash_attention_2:
+            self.model = AutoModelForCausalLM.from_pretrained(self.args.model_name, 
+                                                              cache_dir=self.args.cache_dir,
+                                                              torch_dtype=torch.bfloat16,
+                                                              attn_implementation="flash_attention_2")
+        else:
+            self.model = AutoModelForCausalLM.from_pretrained(self.args.model_name, 
+                                                              cache_dir=self.args.cache_dir,
+                                                              torch_dtype=torch.bfloat16)
+                                                          
+        # Load Dataset
+        print(">>> 3. Loading Dataset")
+        self.data = load_dataset(self.args.data_name, cache_dir=self.args.cache_dir)
+
+        # Preprocess Dataset
+        print(">>> 4. Filtering and Preprocessing Dataset")
+        data_split = dataset_split_selector(self.data)
+
+        if len(data_split) == 1:
+            self.is_test = False
+            train_split = data_split[0]
+        else:
+            self.is_test = True
+            train_split = data_split[0]
+            test_split = data_split[0]
+
+            test = self.data[test_split].filter(self.filter_dataset)
+            self.test = test.map(self.preprocess_dataset, batched=True, num_proc=self.args.num_proc, remove_columns=self.data[test_split].column_names)       
+
+        train = self.data[train_split].filter(self.filter_dataset)
+        print(f"\n\n>>> {len(train)} / {len(self.data[train_split])} rows left after filtering by prompt length.")
+        self.train = train.map(self.preprocess_dataset, batched=True, num_proc=self.args.num_proc, remove_columns=self.data[train_split].column_names)                       
+                
+        # Set WANDB & Logging Configurations
+        self.run_name = f"{self.args.model_name.split('/')[-1]}-{self.args.data_name.split('/')[-1]}-ORPO-{self.start.tm_mday}-{self.start.tm_hour}-{self.start.tm_min}"
+        self.save_dir = os.path.join('./checkpoints/', f"{self.args.data_name.split('/')[-1]}/{self.run_name}")
+        self.log_dir = os.path.join('./checkpoints/', f"{self.args.data_name.split('/')[-1]}/{self.run_name}/logs")
+        
+        os.makedirs(self.save_dir, exist_ok=True)
+        os.makedirs(self.log_dir, exist_ok=True)
+
+    def preprocess_dataset(self, examples: Union[List, Dict]):
+        if 'instruction' in examples.keys():
+            prompt_key = 'instruction'
+            prompt = [self.tokenizer.apply_chat_template([{'role': 'user', 'content': item}], tokenize=False, add_generation_prompt=True) for item in examples[prompt_key]]
+            chosen = [self.tokenizer.apply_chat_template([{'role': 'user', 'content': item_prompt}, {'role': 'assistant', 'content': item_chosen}], tokenize=False) for item_prompt, item_chosen in zip(examples[prompt_key], examples['chosen'])]
+            rejected = [self.tokenizer.apply_chat_template([{'role': 'user', 'content': item_prompt}, {'role': 'assistant', 'content': item_rejected}], tokenize=False) for item_prompt, item_rejected in zip(examples[prompt_key], examples['rejected'])]
+        else:
+            prompt = [self.tokenizer.apply_chat_template([item[0]], tokenize=False, add_generation_prompt=True) for item in examples['chosen']]
+            chosen = [self.tokenizer.apply_chat_template(item, tokenize=False) for item in examples['chosen']]
+            rejected = [self.tokenizer.apply_chat_template(item, tokenize=False) for item in examples['rejected']]
+    
+        model_inputs = self.tokenizer(prompt,
+                                      max_length=self.args.response_max_length,
+                                      padding='max_length',
+                                      truncation=True,
+                                      return_tensors='pt')
+        pos_labels = self.tokenizer(chosen,
+                                    max_length=self.args.response_max_length,
+                                    padding='max_length',
+                                    truncation=True,
+                                    return_tensors='pt')
+        neg_labels = self.tokenizer(rejected,
+                                    max_length=self.args.response_max_length,
+                                    padding='max_length',
+                                    truncation=True,
+                                    return_tensors='pt') 
+                
+        model_inputs['positive_input_ids'] = pos_labels['input_ids']
+        model_inputs['positive_attention_mask'] = pos_labels['attention_mask']
+        
+        model_inputs['negative_input_ids'] = neg_labels['input_ids']
+        model_inputs['negative_attention_mask'] = neg_labels['attention_mask']
+        
+        return model_inputs
+
+    def filter_dataset(self, examples: Union[List, Dict]):
+        if 'instruction' in examples.keys():
+            query = examples['instruction']
+            prompt_length = self.tokenizer.apply_chat_template([{'content': query, 'role': 'user'}], tokenize=True, add_generation_prompt=True, return_tensors='pt').size(-1)
+        else:
+            prompt_length = self.tokenizer.apply_chat_template([examples['chosen'][0]], tokenize=True, add_generation_prompt=True, return_tensors='pt').size(-1)  
+           
+        if prompt_length < self.args.prompt_max_length:    
+            return True
+        else:
+            return False
+
+    def prepare_trainer(self):
+        wandb.init(name=self.run_name)
+        arguments = TrainingArguments(
+            torch_compile=self.args.torch_compile,
+            output_dir=self.save_dir,  # The output directory
+            logging_dir=self.log_dir,
+            logging_steps=50,
+            learning_rate=self.args.lr,
+            overwrite_output_dir=True,  # overwrite the content of the output directory
+            num_train_epochs=self.args.num_train_epochs,  # number of training epochs
+            per_device_train_batch_size=self.args.per_device_train_batch_size,  # batch size for training
+            per_device_eval_batch_size=self.args.per_device_eval_batch_size,  # batch size for evaluation
+            evaluation_strategy=self.args.evaluation_strategy,  # batch size for evaluation
+            save_strategy=self.args.evaluation_strategy,
+            optim=self.args.optim,
+            warmup_steps=self.args.warmup_steps,
+            gradient_accumulation_steps=self.args.gradient_accumulation_steps,
+            gradient_checkpointing=True, #if ('llama' in self.args.model_name.lower()) or ('mistral' in self.args.model_name.lower()) else False,
+            gradient_checkpointing_kwargs={'use_reentrant':True},
+            load_best_model_at_end=True,
+            do_train=True,
+            do_eval= self.is_test,
+            lr_scheduler_type=self.args.lr_scheduler_type,
+            remove_unused_columns=False,
+            report_to='wandb',
+            run_name=self.run_name,
+            bf16=True
+        )
+        
+        data_collator = DataCollatorForLanguageModeling(tokenizer=self.tokenizer, mlm=False)
+        
+        self.trainer = ORPOTrainer(
+            model=self.model,
+            alpha=self.args.alpha,
+            pad=self.tokenizer.pad_token_id,
+            args=arguments,
+            train_dataset=self.train,
+            eval_dataset=self.test if self.is_test else None,
+            data_collator=data_collator,
+            preprocess_logits_for_metrics=preprocess_logits_for_metrics
+        )
+        
+    def run(self):
+        print(">>> 5. Preparing ORPOTrainer")
+        self.prepare_trainer()
+        self.trainer.train()
+
+        # Saving code for FSDP
+        if self.trainer.is_fsdp_enabled:
+            self.trainer.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")
+        self.trainer.save_model()
+        
+        
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser("ORPO")
+    args = default_args(parser)
+
+    # Set WANDB configurations
+    if args.wandb_entity is not None and args.wandb_project_name is not None:
+        os.environ["WANDB_ENTITY"] = args.wandb_entity
+        os.environ["WANDB_PROJECT"] = args.wandb_project_name
+    else:
+        pass
+    os.environ["TOKENIZERS_PARALLELISM"] = 'false'
+
+    print("================================================================================================\n")
+    print(f">>> Fine-tuning {args.model_name} with ORPO on {args.data_name}\n")
+    print("================================================================================================")
+    print("\n\n>>> Summary:")
+    print(f"    - Lambda              : {args.alpha}")
+    print(f"    - Training Epochs     : {args.num_train_epochs}")
+    print(f"    - Prompt Max Length   : {args.prompt_max_length}")
+    print(f"    - Response Max Length : {args.response_max_length}")
+
+    item = ORPO(args=args)
+    item.run()

src/accelerate/ds2.yaml

@@ -0,0 +1,21 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+deepspeed_config:
+  gradient_accumulation_steps: 1
+  offload_optimizer_device: none
+  offload_param_device: none
+  zero3_init_flag: false
+  zero_stage: 2
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 2
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

src/args.py

@@ -0,0 +1,34 @@
+def default_args(parser):
+    parser.add_argument("--cache_dir", default=None, type=str)
+    parser.add_argument("--save_dir", default='./saved', type=str)
+    parser.add_argument("--data_name", default='HuggingfaceH4/UltraFeedback', type=str)
+    parser.add_argument("--model_name", default="gpt2", type=str)
+
+    # Training Arguments
+    parser.add_argument("--torch_compile", default=True, type=bool)
+    parser.add_argument("--flash_attention_2", action='store_true')
+    parser.add_argument("--lr_scheduler_type", default="cosine", type=str)
+    parser.add_argument("--optim", default="paged_adamw_32bit", type=str)
+    parser.add_argument("--overwrite_output_dir", default=True, type=bool)
+    parser.add_argument("--lr", default=2e-5, type=float)
+    parser.add_argument("--num_proc", default=1, type=int)
+    parser.add_argument("--num_train_epochs", default=10, type=int)
+    parser.add_argument("--per_device_train_batch_size", default=2, type=int)
+    parser.add_argument("--per_device_eval_batch_size", default=2, type=int)
+    parser.add_argument("--warmup_steps", default=5000, type=int)    
+    parser.add_argument("--evaluation_strategy", default='epoch', type=str)
+    parser.add_argument("--do_eval", action='store_true')
+    parser.add_argument("--gradient_accumulation_steps", default=1, type=int)
+    parser.add_argument("--save_strategy", default='epoch', type=str)
+    parser.add_argument("--prompt_max_length", default=256, type=int)
+    parser.add_argument("--response_max_length", default=1024, type=int)
+    parser.add_argument("--alpha", default=1.0, type=float)
+
+    # Wandb Configurations
+    parser.add_argument("--wandb_entity", default=None, type=str)
+    parser.add_argument("--wandb_project_name", default=None, type=str)
+
+    
+    args = parser.parse_args()
+
+    return args

src/orpo_trainer.py

@@ -0,0 +1,83 @@
+
+import torch
+import wandb
+from transformers import Trainer
+
+
+class ORPOTrainer(Trainer):
+    def __init__(self, alpha, pad, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.pad = pad
+        self.alpha = alpha
+        self.loss_fct = torch.nn.CrossEntropyLoss(reduction='none')
+        print("Pad Token ID: ", self.pad)
+        
+    def compute_custom_loss(self, logits, labels):
+        
+        logits = logits.contiguous()
+        
+        if labels is not None:
+            # move labels to correct device to enable model parallelism
+            labels = labels.to(logits.device)
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            
+            # Flatten the tokens
+            loss = self.loss_fct(shift_logits.transpose(2, 1), shift_labels).mean(dim=-1)
+            
+        return loss
+    
+    def compute_logps(self, prompt_attention_mask, chosen_inputs, chosen_attention_mask, logits):
+        mask = chosen_attention_mask[:, :-1] - prompt_attention_mask[:, 1:]
+        per_token_logps = torch.gather(logits[:, :-1, :].log_softmax(-1), dim=2, 
+                                       index=(mask * chosen_inputs[:, 1:]).unsqueeze(2)).squeeze(2)
+        return torch.mul(per_token_logps, mask.to(dtype=torch.bfloat16)).sum(dim=1).to(dtype=torch.float64) / mask.sum(dim=1).to(dtype=torch.float64)
+        
+    def compute_loss(self, model, inputs, return_outputs=False):
+        if self.label_smoother is not None and "labels" in inputs:
+            labels = inputs.pop("labels")
+        else:
+            labels = None
+        
+        # Generate the hidden states for 'chosen' and 'reject'
+        neg_labels = inputs['negative_input_ids'].clone()
+        pos_labels = inputs['positive_input_ids'].clone()
+
+        neg_labels[neg_labels == self.pad] = -100
+        pos_labels[pos_labels == self.pad] = -100
+
+        outputs_neg = model(**{'input_ids': inputs['negative_input_ids'],
+                               'attention_mask': inputs['negative_attention_mask'],
+                               'labels': neg_labels,}, output_hidden_states=True)      
+        outputs_pos = model(**{'input_ids': inputs['positive_input_ids'],
+                               'attention_mask': inputs['positive_attention_mask'],
+                               'labels': pos_labels,}, output_hidden_states=True)
+            
+        # Calculate NLL loss
+        pos_loss = self.compute_custom_loss(logits=outputs_pos.logits, labels=inputs['positive_input_ids']) 
+
+        # Calculate Log Probability
+        pos_prob = self.compute_logps(prompt_attention_mask=inputs['attention_mask'], 
+                                      chosen_inputs=inputs['positive_input_ids'], 
+                                      chosen_attention_mask=inputs['positive_attention_mask'], 
+                                      logits=outputs_pos.logits)
+        neg_prob = self.compute_logps(prompt_attention_mask=inputs['attention_mask'], 
+                                      chosen_inputs=inputs['negative_input_ids'], 
+                                      chosen_attention_mask=inputs['negative_attention_mask'], 
+                                      logits=outputs_neg.logits)
+
+        # Calculate log odds
+        log_odds = (pos_prob - neg_prob) - (torch.log(1 - torch.exp(pos_prob)) - torch.log(1 - torch.exp(neg_prob)))
+        sig_ratio = torch.nn.functional.sigmoid(log_odds)
+        ratio = torch.log(sig_ratio)
+        
+        # Calculate the Final Loss
+        loss = torch.mean(pos_loss - self.alpha * ratio).to(dtype=torch.bfloat16)
+        
+        wandb.log({'Positive Geometric Mean': torch.mean(pos_prob).item(),
+                   'Negative Geometric Mean': torch.mean(neg_prob).item(),
+                   'Log Odds Ratio': torch.mean(ratio).item(),
+                   'Log Odds': torch.mean(log_odds).item()})
+        
+        return (loss, outputs_pos) if return_outputs else loss

src/utils.py

@@ -0,0 +1,20 @@
+
+from typing import List
+
+def preprocess_logits_for_metrics(logits, labels):
+    if isinstance(logits, tuple):
+        logits = logits[0]
+    return logits.argmax(dim=-1)
+
+def dataset_split_selector(data) -> List:
+    """
+    This is a function for automating the process of selecting data split.
+    Will be further updated.
+    """
+    if len(data.keys()) == 1:
+        return ['train']
+    else:
+        if 'train_prefs' in data.keys():
+            return ['train_prefs', 'test_prefs']
+        else:
+            return ['train', 'test']