refactor dataset_attr, add eos in pt, fix #757

Former-commit-id: 0feec9a830b917b36686b61938a66e842eccf930

src/llmtuner/tuner/dpo/trainer.py

@@ -1,9 +1,9 @@
 import torch
 from collections import defaultdict
-from peft import PeftModel
 from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
 from transformers import BatchEncoding, Trainer
 from trl import DPOTrainer
+from trl.trainer.utils import disable_dropout_in_model
 
 from llmtuner.extras.constants import IGNORE_INDEX
 from llmtuner.tuner.core.trainer import PeftModelMixin
@@ -18,9 +18,16 @@ class DPOPeftTrainer(PeftModelMixin, DPOTrainer):
     def __init__(
         self,
         finetuning_args: "FinetuningArguments",
+        model: Union["PreTrainedModel", torch.nn.Module],
         ref_model: Optional[Union["PreTrainedModel", torch.nn.Module]] = None,
+        disable_dropout: Optional[bool] = True,
         **kwargs
     ):
+        if disable_dropout:
+            disable_dropout_in_model(model)
+            if ref_model is not None:
+                disable_dropout_in_model(ref_model)
+
         self.finetuning_args = finetuning_args
         self.ref_model = ref_model
         self.use_dpo_data_collator = True # hack to avoid warning
@@ -29,12 +36,16 @@ class DPOPeftTrainer(PeftModelMixin, DPOTrainer):
         self.beta = finetuning_args.dpo_beta
         self._stored_metrics = defaultdict(lambda: defaultdict(list))
 
-        Trainer.__init__(self, **kwargs)
+        Trainer.__init__(self, model=model, **kwargs)
         if not hasattr(self, "accelerator"):
             raise AttributeError("Please update `transformers`.")
 
         if ref_model is not None:
-            self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True)
+            if self.is_deepspeed_enabled:
+                self.ref_model = self.accelerator._prepare_deepspeed(self.ref_model)
+                self.ref_model.eval()
+            else:
+                self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True)
 
     def concatenated_forward(
         self,
@@ -42,27 +53,12 @@ class DPOPeftTrainer(PeftModelMixin, DPOTrainer):
         batch: Optional[Dict[str, torch.Tensor]] = None
     ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
         batch_copied = BatchEncoding({k: v.detach().clone() for k, v in batch.items()}) # avoid error
-        unwrapped_model: "PreTrainedModel" = self.accelerator.unwrap_model(self.model)
 
-        if not torch.is_grad_enabled():
-            unwrapped_model.gradient_checkpointing_disable()
-
-        if model is None and isinstance(unwrapped_model, PeftModel): # peft model has no ref_model
-            with unwrapped_model.disable_adapter():
-                all_logits = self.model(
-                    input_ids=batch_copied["input_ids"],
-                    attention_mask=batch_copied["attention_mask"],
-                    return_dict=True
-                ).logits.to(torch.float32)
-        else:
-            all_logits = model(
-                input_ids=batch_copied["input_ids"],
-                attention_mask=batch_copied["attention_mask"],
-                return_dict=True
-            ).logits.to(torch.float32)
-
-        if not torch.is_grad_enabled():
-            unwrapped_model.gradient_checkpointing_enable()
+        all_logits = model(
+            input_ids=batch_copied["input_ids"],
+            attention_mask=batch_copied["attention_mask"],
+            return_dict=True
+        ).logits.to(torch.float32)
 
         all_logps = self._get_batch_logps(
             all_logits,