support galore

Former-commit-id: b67a4a46a88d83bb2a3459b3317b66cda15e0171

src/llmtuner/train/dpo/workflow.py

@@ -7,9 +7,9 @@ from ...extras.constants import IGNORE_INDEX
 from ...extras.ploting import plot_loss
 from ...hparams import ModelArguments
 from ...model import load_model, load_tokenizer
-from ...train.dpo.collator import DPODataCollatorWithPadding
-from ...train.dpo.trainer import CustomDPOTrainer
-from ...train.utils import create_modelcard_and_push, create_ref_model
+from ..utils import create_custom_optimzer, create_modelcard_and_push, create_ref_model
+from .collator import DPODataCollatorWithPadding
+from .trainer import CustomDPOTrainer
 
 
 if TYPE_CHECKING:
@@ -44,6 +44,7 @@ def run_dpo(
     training_args.remove_unused_columns = False  # important for pairwise dataset
 
     # Initialize our Trainer
+    optimizer = create_custom_optimzer(model, training_args, finetuning_args)
     trainer = CustomDPOTrainer(
         beta=finetuning_args.dpo_beta,
         loss_type=finetuning_args.dpo_loss,
@@ -54,6 +55,7 @@ def run_dpo(
         tokenizer=tokenizer,
         data_collator=data_collator,
         callbacks=callbacks,
+        optimizers=(optimizer, None),
         **split_dataset(dataset, data_args, training_args),
     )