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/ppo/workflow.py
+++ b/src/llmtuner/train/ppo/workflow.py
@@ -64,7 +64,7 @@ def run_ppo(
    )

    # Create optimizer and scheduler
-    optimizer = create_custom_optimzer(model, training_args, finetuning_args)
+    optimizer = create_custom_optimzer(model, dataset, training_args, finetuning_args)
    if optimizer is None:
        optimizer = AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=training_args.learning_rate)