support DPO training (2305.18290)

Former-commit-id: 6d98de148e4af63a7028dfaeb6cf86eb56a4488f

src/llmtuner/tuner/ppo/trainer.py

@@ -10,7 +10,7 @@ from trl import PPOTrainer
 from trl.core import LengthSampler
 
 from llmtuner.extras.logging import get_logger
-from llmtuner.extras.misc import AverageMeter, count_parameters, get_logits_processor
+from llmtuner.extras.misc import AverageMeter, count_parameters, get_logits_processor, get_stopping_criteria
 from llmtuner.tuner.core.trainer import PeftTrainer
 from llmtuner.tuner.ppo.utils import cast_layernorm_dtype, replace_model
 
@@ -18,7 +18,7 @@ if TYPE_CHECKING:
     from transformers import Seq2SeqTrainingArguments
     from trl import AutoModelForCausalLMWithValueHead
     from llmtuner.extras.callbacks import LogCallback
-    from llmtuner.hparams import FinetuningArguments
+    from llmtuner.hparams import FinetuningArguments, GeneratingArguments
 
 
 logger = get_logger(__name__)
@@ -33,16 +33,17 @@ class PPOPeftTrainer(PPOTrainer, PeftTrainer):
         self,
         training_args: "Seq2SeqTrainingArguments",
         finetuning_args: "FinetuningArguments",
+        generating_args: "GeneratingArguments",
         callbacks: List["LogCallback"],
         **kwargs
     ):
         PPOTrainer.__init__(self, **kwargs)
         self.args = training_args
         self.finetuning_args = finetuning_args
+        self.generating_args = generating_args
         self.log_callback = callbacks[0]
         self.state = TrainerState()
         self.control = TrainerControl()
-        self._remove_log()
 
     def ppo_train(self, max_target_length: int) -> None:
         r"""
@@ -72,14 +73,10 @@ class PPOPeftTrainer(PPOTrainer, PeftTrainer):
             logger.info(f"  Number of trainable parameters = {count_parameters(self.model)[0]}")
 
         # Keyword arguments for `model.generate`
-        gen_kwargs = {
-            "top_k": 0.0,
-            "top_p": 1.0,
-            "do_sample": True,
-            "pad_token_id": self.tokenizer.pad_token_id,
-            "eos_token_id": self.tokenizer.eos_token_id,
-            "logits_processor": get_logits_processor()
-        }
+        gen_kwargs = self.generating_args.to_dict()
+        gen_kwargs["logits_processor"] = get_logits_processor()
+        gen_kwargs["stopping_criteria"] = get_stopping_criteria(self.tokenizer.additional_special_tokens_ids)
+
         length_sampler = LengthSampler(max_target_length // 2, max_target_length)
         unwrapped_model: "AutoModelForCausalLMWithValueHead" = self.accelerator.unwrap_model(self.model)
 

src/llmtuner/tuner/ppo/workflow.py

@@ -1,11 +1,9 @@
-# Inspired by:
-# https://github.com/lvwerra/trl/blob/main/examples/research_projects/stack_llama/scripts/rl_training.py
+# Inspired by: https://github.com/lvwerra/trl/blob/main/examples/research_projects/stack_llama/scripts/rl_training.py
 
 import math
-from typing import TYPE_CHECKING
 from trl import PPOConfig
 from torch.optim import AdamW
-from typing import Optional, List
+from typing import TYPE_CHECKING, Optional, List
 from transformers import DataCollatorForSeq2Seq
 from transformers.optimization import get_scheduler
 
@@ -16,7 +14,7 @@ from llmtuner.tuner.ppo.trainer import PPOPeftTrainer
 
 if TYPE_CHECKING:
     from transformers import Seq2SeqTrainingArguments, TrainerCallback
-    from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments
+    from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments
 
 
 def run_ppo(
@@ -24,6 +22,7 @@ def run_ppo(
     data_args: "DataArguments",
     training_args: "Seq2SeqTrainingArguments",
     finetuning_args: "FinetuningArguments",
+    generating_args: "GeneratingArguments",
     callbacks: Optional[List["TrainerCallback"]] = None
 ):
     dataset = get_dataset(model_args, data_args)
@@ -42,8 +41,9 @@ def run_ppo(
     )
 
     optimizer = AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=training_args.learning_rate)
-    total_train_batch_size = \
+    total_train_batch_size = (
         training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps * training_args.world_size
+    )
     num_training_steps = training_args.num_train_epochs * math.ceil(len(dataset) / total_train_batch_size)
     lr_scheduler = get_scheduler(
         training_args.lr_scheduler_type,
@@ -56,6 +56,7 @@ def run_ppo(
     ppo_trainer = PPOPeftTrainer(
         training_args=training_args,
         finetuning_args=finetuning_args,
+        generating_args=generating_args,
         callbacks=callbacks,
         config=ppo_config,
         model=model,
@@ -67,8 +68,10 @@ def run_ppo(
         lr_scheduler=lr_scheduler
     )
 
-    ppo_trainer.ppo_train(max_target_length=data_args.max_target_length)
-    ppo_trainer.save_model()
-    ppo_trainer.save_state() # must be after save_model
-    if ppo_trainer.is_world_process_zero() and model_args.plot_loss:
-        plot_loss(training_args.output_dir, keys=["loss", "reward"])
+    # Training
+    if training_args.do_train:
+        ppo_trainer.ppo_train(max_target_length=data_args.max_target_length)
+        ppo_trainer.save_model()
+        ppo_trainer.save_state() # must be called after save_model to have a folder
+        if ppo_trainer.is_world_process_zero() and model_args.plot_loss:
+            plot_loss(training_args.output_dir, keys=["loss", "reward"])