support layerwise galore

Former-commit-id: d43a4da0947897d0be3f62fad3107754d4c89f2b
2024-03-10 00:24:11 +08:00
parent c635bbe465
commit 7ff8a064f3
14 changed files with 109 additions and 51 deletions
--- a/src/llmtuner/train/rm/workflow.py
+++ b/src/llmtuner/train/rm/workflow.py
@@ -35,7 +35,7 @@ def run_rm(
    training_args.remove_unused_columns = False  # important for pairwise dataset

    # Initialize our Trainer
-    optimizer = create_custom_optimzer(model, training_args, finetuning_args)
+    optimizer = create_custom_optimzer(model, dataset, training_args, finetuning_args)
    trainer = PairwiseTrainer(
        model=model,
        args=training_args,