support DPO training (2305.18290)

Former-commit-id: 6d98de148e4af63a7028dfaeb6cf86eb56a4488f

src/llmtuner/tuner/core/adapter.py

@@ -39,7 +39,7 @@ def init_adapter(
     if finetuning_args.finetuning_type == "none" and is_trainable:
         raise ValueError("You cannot use finetuning_type=none while training.")
 
-    if finetuning_args.finetuning_type == "full":
+    if finetuning_args.finetuning_type == "full" and is_trainable:
         logger.info("Fine-tuning method: Full")
         model = model.float()
 

src/llmtuner/tuner/core/loader.py

@@ -34,7 +34,7 @@ check_min_version("4.29.1")
 require_version("datasets>=2.12.0", "To fix: pip install datasets>=2.12.0")
 require_version("accelerate>=0.21.0", "To fix: pip install accelerate>=0.21.0")
 require_version("peft>=0.4.0", "To fix: pip install peft>=0.4.0")
-require_version("trl>=0.4.7", "To fix: pip install trl>=0.4.7")
+require_version("trl>=0.5.0", "To fix: pip install trl>=0.5.0")
 
 
 def load_model_and_tokenizer(
@@ -52,9 +52,6 @@ def load_model_and_tokenizer(
         logger.warning("Checkpoint is not found at evaluation, load the original model.")
         finetuning_args = FinetuningArguments(finetuning_type="none")
 
-    assert stage in ["pt", "sft"] or finetuning_args.finetuning_type == "lora", \
-        "RM and PPO training can only be performed with the LoRA method."
-
     config_kwargs = {
         "trust_remote_code": True,
         "cache_dir": model_args.cache_dir,
@@ -132,8 +129,6 @@ def load_model_and_tokenizer(
                 })
 
         if stage == "ppo": # load reward model
-            assert is_trainable, "PPO stage cannot be performed at evaluation."
-            assert model_args.reward_model is not None, "Reward model is necessary for PPO training."
             logger.info("Load reward model from {}".format(model_args.reward_model))
             model.pretrained_model.load_adapter(model_args.reward_model, "reward", is_trainable=False)
             assert load_valuehead_params(model, model_args.reward_model), "Reward model is not correctly loaded."

src/llmtuner/tuner/core/parser.py

@@ -19,7 +19,7 @@ from llmtuner.hparams import (
 logger = get_logger(__name__)
 
 
-def _parse_args(parser: HfArgumentParser, args: Optional[Dict[str, Any]] = None):
+def _parse_args(parser: HfArgumentParser, args: Optional[Dict[str, Any]] = None) -> Tuple[Any]:
     if args is not None:
         return parser.parse_dict(args)
     elif len(sys.argv) == 2 and sys.argv[1].endswith(".yaml"):
@@ -32,26 +32,53 @@ def _parse_args(parser: HfArgumentParser, args: Optional[Dict[str, Any]] = None)
 
 def parse_train_args(
     args: Optional[Dict[str, Any]] = None
-) -> Tuple[ModelArguments, DataArguments, Seq2SeqTrainingArguments, FinetuningArguments, GeneralArguments]:
+) -> Tuple[
+    ModelArguments,
+    DataArguments,
+    Seq2SeqTrainingArguments,
+    FinetuningArguments,
+    GeneratingArguments,
+    GeneralArguments
+]:
     parser = HfArgumentParser((
-        ModelArguments, DataArguments, Seq2SeqTrainingArguments, FinetuningArguments, GeneralArguments
+        ModelArguments,
+        DataArguments,
+        Seq2SeqTrainingArguments,
+        FinetuningArguments,
+        GeneratingArguments,
+        GeneralArguments
     ))
     return _parse_args(parser, args)
 
 
 def parse_infer_args(
     args: Optional[Dict[str, Any]] = None
-) -> Tuple[ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments]:
+) -> Tuple[
+    ModelArguments,
+    DataArguments,
+    FinetuningArguments,
+    GeneratingArguments
+]:
     parser = HfArgumentParser((
-        ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments
+        ModelArguments,
+        DataArguments,
+        FinetuningArguments,
+        GeneratingArguments
     ))
     return _parse_args(parser, args)
 
 
 def get_train_args(
     args: Optional[Dict[str, Any]] = None
-) -> Tuple[ModelArguments, DataArguments, Seq2SeqTrainingArguments, FinetuningArguments, GeneralArguments]:
-    model_args, data_args, training_args, finetuning_args, general_args = parse_train_args(args)
+) -> Tuple[
+    ModelArguments,
+    DataArguments,
+    Seq2SeqTrainingArguments,
+    FinetuningArguments,
+    GeneratingArguments,
+    GeneralArguments
+]:
+    model_args, data_args, training_args, finetuning_args, generating_args, general_args = parse_train_args(args)
 
     # Setup logging
     if training_args.should_log:
@@ -68,7 +95,7 @@ def get_train_args(
     data_args.init_for_training()
 
     if general_args.stage != "sft" and training_args.predict_with_generate:
-        raise ValueError("`predict_with_generate` cannot be set as True at PT, RM and PPO stages.")
+        raise ValueError("`predict_with_generate` cannot be set as True except SFT.")
 
     if training_args.do_train and training_args.predict_with_generate:
         raise ValueError("`predict_with_generate` cannot be set as True while training.")
@@ -76,6 +103,15 @@ def get_train_args(
     if general_args.stage == "sft" and training_args.do_predict and not training_args.predict_with_generate:
         raise ValueError("Please enable `predict_with_generate` to save model predictions.")
 
+    if general_args.stage in ["rm", "ppo"] and finetuning_args.finetuning_type != "lora":
+        raise ValueError("RM and PPO training can only be performed with the LoRA method.")
+
+    if general_args.stage in ["ppo", "dpo"] and not training_args.do_train:
+        raise ValueError("PPO and DPO stage can only be performed at training.")
+
+    if general_args.stage == "ppo" and model_args.reward_model is None:
+        raise ValueError("Reward model is necessary for PPO training.")
+
     if training_args.max_steps == -1 and data_args.streaming:
         raise ValueError("Please specify `max_steps` in streaming mode.")
 
@@ -133,12 +169,17 @@ def get_train_args(
     # Set seed before initializing model.
     transformers.set_seed(training_args.seed)
 
-    return model_args, data_args, training_args, finetuning_args, general_args
+    return model_args, data_args, training_args, finetuning_args, generating_args, general_args
 
 
 def get_infer_args(
     args: Optional[Dict[str, Any]] = None
-) -> Tuple[ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments]:
+) -> Tuple[
+    ModelArguments,
+    DataArguments,
+    FinetuningArguments,
+    GeneratingArguments
+]:
     model_args, data_args, finetuning_args, generating_args = parse_infer_args(args)
 
     if model_args.quantization_bit is not None and finetuning_args.finetuning_type != "lora":

src/llmtuner/tuner/core/trainer.py

@@ -13,26 +13,25 @@ from llmtuner.extras.logging import get_logger
 from llmtuner.extras.save_and_load import get_state_dict, load_trainable_params
 
 if TYPE_CHECKING:
+    from transformers import PreTrainedTokenizer, Seq2SeqTrainingArguments, TrainerState
     from llmtuner.hparams import FinetuningArguments
 
 
 logger = get_logger(__name__)
 
 
-class PeftTrainer(Seq2SeqTrainer):
+class PeftModelMixin:
     r"""
-    Inherits Seq2SeqTrainer to support parameter-efficient checkpoints.
+    Patches the save and load methods in Hugging Face Trainer for PeftModel and ModelWithValueHead.
     """
 
-    def __init__(self, finetuning_args: "FinetuningArguments", **kwargs):
-        super().__init__(**kwargs)
-        self.finetuning_args = finetuning_args
-        self._remove_log()
-
-    def _remove_log(self):
-        if self.is_world_process_zero() and os.path.exists(os.path.join(self.args.output_dir, "trainer_log.jsonl")):
-            logger.warning("Previous log file in this folder will be deleted.")
-            os.remove(os.path.join(self.args.output_dir, "trainer_log.jsonl"))
+    def __init__(self) -> None: # for type checking
+        self.model: PreTrainedModel = None
+        self.tokenizer: "PreTrainedTokenizer" = None
+        self.args: "Seq2SeqTrainingArguments" = None
+        self.finetuning_args: "FinetuningArguments" = None
+        self.state: "TrainerState" = None
+        raise AssertionError("Mixin should not be initialized.")
 
     def _save(self, output_dir: Optional[str] = None, state_dict: Optional[Dict[str, torch.Tensor]] = None) -> None:
         r"""
@@ -96,3 +95,13 @@ class PeftTrainer(Seq2SeqTrainer):
             model.load_adapter(self.state.best_model_checkpoint, model.active_adapter)
         else: # freeze/full-tuning
             load_trainable_params(model, self.state.best_model_checkpoint)
+
+
+class PeftTrainer(PeftModelMixin, Seq2SeqTrainer):
+    r"""
+    Inherits Seq2SeqTrainer to support parameter-efficient checkpoints.
+    """
+
+    def __init__(self, finetuning_args: "FinetuningArguments", **kwargs):
+        Seq2SeqTrainer.__init__(self, **kwargs)
+        self.finetuning_args = finetuning_args