[DPO] use ref model logprobs if it exists in the data

tests/test_dpo_trainer.py

@@ -28,17 +28,24 @@
 class DPOTrainerTester(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
-        cls.model_id = "trl-internal-testing/dummy-GPT2-correct-vocab"
-        cls.model = AutoModelForCausalLM.from_pretrained(cls.model_id)
-        cls.ref_model = AutoModelForCausalLM.from_pretrained(cls.model_id)
-        cls.tokenizer = AutoTokenizer.from_pretrained(cls.model_id)
-        cls.tokenizer.pad_token = cls.tokenizer.eos_token
-
-        # get t5 as seq2seq example:
-        model_id = "trl-internal-testing/tiny-T5ForConditionalGeneration-correct-vocab"
-        cls.t5_model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
-        cls.t5_ref_model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
-        cls.t5_tokenizer = AutoTokenizer.from_pretrained(model_id)
+        cls.models = {
+            "gpt2": {
+                "model_id": "trl-internal-testing/dummy-GPT2-correct-vocab",
+                "model_type": AutoModelForCausalLM,
+                "tokenizer_type": AutoTokenizer,
+            },
+            "t5": {
+                "model_id": "trl-internal-testing/tiny-T5ForConditionalGeneration-correct-vocab",
+                "model_type": AutoModelForSeq2SeqLM,
+                "tokenizer_type": AutoTokenizer,
+            },
+            # add more models here if needed
+        }
+        for key, model_info in cls.models.items():
+            model_info["model"] = model_info["model_type"].from_pretrained(model_info["model_id"])
+            model_info["ref_model"] = model_info["model_type"].from_pretrained(model_info["model_id"])
+            model_info["tokenizer"] = model_info["tokenizer_type"].from_pretrained(model_info["model_id"])
+        cls.models["gpt2"]["tokenizer"].pad_token = cls.models["gpt2"]["tokenizer"].eos_token
 
     def _init_dummy_dataset(self):
         # fmt: off
@@ -74,8 +81,12 @@ def _init_dummy_dataset(self):
         # fmt: on
         return Dataset.from_dict(dummy_dataset_dict)
 
-    @parameterized.expand([["gpt2", "sigmoid"], ["t5", "hinge"]])
-    def test_dpo_trainer(self, name, loss_type):
+    def _get_models_by_name(self, name):
+        return self.models[name]["model"], self.models[name]["ref_model"], self.models[name]["tokenizer"]
+
+    @parameterized.expand([["gpt2", "sigmoid", True], ["t5", "hinge", False]])
+    def test_dpo_trainer(self, name, loss_type, precompute_ref_log_probs):
+        model, ref_model, tokenizer = self._get_models_by_name(name)
         with tempfile.TemporaryDirectory() as tmp_dir:
             training_args = TrainingArguments(
                 output_dir=tmp_dir,
@@ -89,15 +100,6 @@ def test_dpo_trainer(self, name, loss_type):
 
             dummy_dataset = self._init_dummy_dataset()
 
-            if name == "gpt2":
-                model = self.model
-                ref_model = self.ref_model
-                tokenizer = self.tokenizer
-            elif name == "t5":
-                model = self.t5_model
-                ref_model = self.t5_ref_model
-                tokenizer = self.t5_tokenizer
-
             trainer = DPOTrainer(
                 model=model,
                 ref_model=ref_model,
@@ -107,6 +109,7 @@ def test_dpo_trainer(self, name, loss_type):
                 tokenizer=tokenizer,
                 train_dataset=dummy_dataset,
                 eval_dataset=dummy_dataset,
+                precompute_ref_log_probs=precompute_ref_log_probs,
             )
 
             previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}
@@ -122,7 +125,9 @@ def test_dpo_trainer(self, name, loss_type):
                 if param.sum() != 0:
                     self.assertFalse(torch.equal(param, new_param))
 
-    def test_dpo_trainer_without_providing_ref_model(self):
+    @parameterized.expand([["gpt2", False], ["t5", True]])
+    def test_dpo_trainer_without_providing_ref_model(self, name, precompute_ref_log_probs):
+        model, _, tokenizer = self._get_models_by_name(name)
         with tempfile.TemporaryDirectory() as tmp_dir:
             training_args = TrainingArguments(
                 output_dir=tmp_dir,
@@ -137,13 +142,14 @@ def test_dpo_trainer_without_providing_ref_model(self):
             dummy_dataset = self._init_dummy_dataset()
 
             trainer = DPOTrainer(
-                model=self.model,
+                model=model,
                 ref_model=None,
                 beta=0.1,
                 args=training_args,
-                tokenizer=self.tokenizer,
+                tokenizer=tokenizer,
                 train_dataset=dummy_dataset,
                 eval_dataset=dummy_dataset,
+                precompute_ref_log_probs=precompute_ref_log_probs,
             )
 
             previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}
@@ -161,7 +167,9 @@ def test_dpo_trainer_without_providing_ref_model(self):
 
     @require_peft
     @mark.peft_test
-    def test_dpo_trainer_without_providing_ref_model_with_lora(self):
+    @parameterized.expand([["gpt2", True], ["t5", False]])
+    def test_dpo_trainer_without_providing_ref_model_with_lora(self, name, precompute_ref_log_probs):
+        model, _, tokenizer = self._get_models_by_name(name)
         from peft import LoraConfig
 
         lora_config = LoraConfig(
@@ -186,14 +194,15 @@ def test_dpo_trainer_without_providing_ref_model_with_lora(self):
             dummy_dataset = self._init_dummy_dataset()
 
             trainer = DPOTrainer(
-                model=self.model,
+                model=model,
                 ref_model=None,
                 beta=0.1,
                 args=training_args,
-                tokenizer=self.tokenizer,
+                tokenizer=tokenizer,
                 train_dataset=dummy_dataset,
                 eval_dataset=dummy_dataset,
                 peft_config=lora_config,
+                precompute_ref_log_probs=precompute_ref_log_probs,
             )
 
             previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}
@@ -210,8 +219,10 @@ def test_dpo_trainer_without_providing_ref_model_with_lora(self):
                     if param.sum() != 0:
                         self.assertFalse(torch.equal(param, new_param))
 
+    @parameterized.expand([["gpt2", False], ["t5", True]])
     @require_no_wandb
-    def test_dpo_trainer_generate_during_eval_no_wandb(self):
+    def test_dpo_trainer_generate_during_eval_no_wandb(self, name, precompute_ref_log_probs):
+        model, _, tokenizer = self._get_models_by_name(name)
         with tempfile.TemporaryDirectory() as tmp_dir:
             training_args = TrainingArguments(
                 output_dir=tmp_dir,
@@ -231,12 +242,13 @@ def test_dpo_trainer_generate_during_eval_no_wandb(self):
                 " Please install `wandb` to resolve.",
             ):
                 DPOTrainer(
-                    model=self.model,
+                    model=model,
                     ref_model=None,
                     beta=0.1,
                     args=training_args,
-                    tokenizer=self.tokenizer,
+                    tokenizer=tokenizer,
                     train_dataset=dummy_dataset,
                     eval_dataset=dummy_dataset,
                     generate_during_eval=True,
+                    precompute_ref_log_probs=precompute_ref_log_probs,
                 )