support galore

Former-commit-id: b67a4a46a88d83bb2a3459b3317b66cda15e0171

src/llmtuner/train/sft/workflow.py

@@ -12,6 +12,7 @@ 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
 
 
 if TYPE_CHECKING:
@@ -49,12 +50,14 @@ 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, training_args, finetuning_args)
     trainer = CustomSeq2SeqTrainer(
         model=model,
         args=training_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),
     )