fix #2777 #2895

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),
     )