fix #2777 #2895

Former-commit-id: 54d5f62d29456a8d9d0c0dd3d0bbfffe48935803

src/llmtuner/train/rm/trainer.py

@@ -6,25 +6,36 @@ import torch
 from transformers import Trainer
 
 from ...extras.logging import get_logger
+from ..utils import create_custom_optimzer
 
 
 if TYPE_CHECKING:
     from transformers.modeling_utils import PreTrainedModel
     from transformers.trainer import PredictionOutput
 
+    from ...hparams import FinetuningArguments
+
 
 logger = get_logger(__name__)
 
 
 class PairwiseTrainer(Trainer):
     r"""
-    Inherits PeftTrainer to compute pairwise loss.
+    Inherits Trainer to compute pairwise loss.
     """
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init__(self, finetuning_args: "FinetuningArguments", **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.finetuning_args = finetuning_args
         self.can_return_loss = True  # override property to return eval_loss
 
+    def create_optimizer_and_scheduler(self, num_training_steps: int) -> None:
+        self.optimizer = create_custom_optimzer(self.model, self.args, self.finetuning_args, num_training_steps)
+        if self.optimizer is None:
+            self.create_optimizer()
+
+        self.create_scheduler(num_training_steps=num_training_steps, optimizer=self.optimizer)
+
     def compute_loss(
         self, model: "PreTrainedModel", inputs: Dict[str, torch.Tensor], return_outputs: bool = False
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[torch.Tensor]]]:

src/llmtuner/train/rm/workflow.py

@@ -7,7 +7,7 @@ from ...extras.callbacks import FixValueHeadModelCallback
 from ...extras.misc import fix_valuehead_checkpoint
 from ...extras.ploting import plot_loss
 from ...model import load_model, load_tokenizer
-from ..utils import create_custom_optimzer, create_modelcard_and_push
+from ..utils import create_modelcard_and_push
 from .collator import PairwiseDataCollatorWithPadding
 from .metric import compute_accuracy
 from .trainer import PairwiseTrainer
@@ -35,14 +35,13 @@ def run_rm(
     training_args.remove_unused_columns = False  # important for pairwise dataset
 
     # Initialize our Trainer
-    optimizer = create_custom_optimzer(model, dataset, training_args, finetuning_args)
     trainer = PairwiseTrainer(
         model=model,
         args=training_args,
+        finetuning_args=finetuning_args,
         tokenizer=tokenizer,
         data_collator=data_collator,
         callbacks=callbacks + [FixValueHeadModelCallback()],
-        optimizers=(optimizer, None),
         compute_metrics=compute_accuracy,
         **split_dataset(dataset, data_args, training_args),
     )