fix #2777 #2895

Former-commit-id: 54d5f62d29456a8d9d0c0dd3d0bbfffe48935803

src/llmtuner/train/sft/trainer.py

@@ -4,28 +4,41 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 import torch
-import torch.nn as nn
 from transformers import Seq2SeqTrainer
 
 from ...extras.constants import IGNORE_INDEX
 from ...extras.logging import get_logger
+from ..utils import create_custom_optimzer
 
 
 if TYPE_CHECKING:
     from transformers.trainer import PredictionOutput
 
+    from ...hparams import FinetuningArguments
+
 
 logger = get_logger(__name__)
 
 
 class CustomSeq2SeqTrainer(Seq2SeqTrainer):
     r"""
-    Inherits PeftTrainer to compute generative metrics such as BLEU and ROUGE.
+    Inherits Seq2SeqTrainer to compute generative metrics such as BLEU and ROUGE.
     """
 
+    def __init__(self, finetuning_args: "FinetuningArguments", **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.finetuning_args = finetuning_args
+
+    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 prediction_step(
         self,
-        model: nn.Module,
+        model: "torch.nn.Module",
         inputs: Dict[str, Union[torch.Tensor, Any]],
         prediction_loss_only: bool,
         ignore_keys: Optional[List[str]] = None,

src/llmtuner/train/sft/workflow.py

@@ -9,10 +9,9 @@ from ...extras.constants import IGNORE_INDEX
 from ...extras.misc import get_logits_processor
 from ...extras.ploting import plot_loss
 from ...model import load_model, load_tokenizer
-from ...train.sft.metric import ComputeMetrics
-from ...train.sft.trainer import CustomSeq2SeqTrainer
-from ...train.utils import create_modelcard_and_push
-from ..utils import create_custom_optimzer
+from ..utils import create_modelcard_and_push
+from .metric import ComputeMetrics
+from .trainer import CustomSeq2SeqTrainer
 
 
 if TYPE_CHECKING:
@@ -50,14 +49,13 @@ def run_sft(
     training_args.generation_num_beams = data_args.eval_num_beams or training_args.generation_num_beams
 
     # Initialize our Trainer
-    optimizer = create_custom_optimzer(model, dataset, training_args, finetuning_args)
     trainer = CustomSeq2SeqTrainer(
         model=model,
         args=training_args,
+        finetuning_args=finetuning_args,
         tokenizer=tokenizer,
         data_collator=data_collator,
         callbacks=callbacks,
-        optimizers=(optimizer, None),
         compute_metrics=ComputeMetrics(tokenizer) if training_args.predict_with_generate else None,
         **split_dataset(dataset, data_args, training_args),
     )