fix #2777 #2895

Former-commit-id: 54d5f62d29456a8d9d0c0dd3d0bbfffe48935803

src/llmtuner/train/pt/trainer.py

@@ -0,0 +1,30 @@
+from typing import TYPE_CHECKING
+
+from transformers import Trainer
+
+from ...extras.logging import get_logger
+from ..utils import create_custom_optimzer
+
+
+if TYPE_CHECKING:
+    from ...hparams import FinetuningArguments
+
+
+logger = get_logger(__name__)
+
+
+class CustomTrainer(Trainer):
+    r"""
+    Inherits Trainer for custom optimizer.
+    """
+
+    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)

src/llmtuner/train/pt/workflow.py

@@ -3,12 +3,13 @@
 import math
 from typing import TYPE_CHECKING, List, Optional
 
-from transformers import DataCollatorForLanguageModeling, Trainer
+from transformers import DataCollatorForLanguageModeling
 
 from ...data import get_dataset, split_dataset
 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 .trainer import CustomTrainer
 
 
 if TYPE_CHECKING:
@@ -30,14 +31,13 @@ def run_pt(
     data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
 
     # Initialize our Trainer
-    optimizer = create_custom_optimzer(model, dataset, training_args, finetuning_args)
-    trainer = Trainer(
+    trainer = CustomTrainer(
         model=model,
         args=training_args,
+        finetuning_args=finetuning_args,
         tokenizer=tokenizer,
         data_collator=data_collator,
         callbacks=callbacks,
-        optimizers=(optimizer, None),
         **split_dataset(dataset, data_args, training_args),
     )