a monkey patch for lora_target

hiyouga · Jul 17, 2023 · 262252d · 262252d
1 parent f8193e8
commit 262252d
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/src/llmtuner/extras/constants.py b/src/llmtuner/extras/constants.py
@@ -29,3 +29,12 @@
     "InternLM-7B-Base": "internlm/internlm-7b",
     "InternLM-7B-Chat": "internlm/internlm-chat-7b"
 }
+
+DEFAULT_MODULE = { # will be deprecated
+    "LLaMA": "q_proj,v_proj",
+    "BLOOM": "query_key_value",
+    "BLOOMZ": "query_key_value",
+    "Falcon": "query_key_value",
+    "Baichuan": "W_pack",
+    "InternLM": "q_proj,v_proj"
+}
diff --git a/src/llmtuner/webui/runner.py b/src/llmtuner/webui/runner.py
@@ -6,6 +6,7 @@
 from typing import Optional, Tuple
 
 from llmtuner.extras.callbacks import LogCallback
+from llmtuner.extras.constants import DEFAULT_MODULE # will be deprecated
 from llmtuner.extras.logging import LoggerHandler
 from llmtuner.extras.misc import torch_gc
 from llmtuner.tuner import get_train_args, run_sft
@@ -79,6 +80,7 @@ def run_train(
             model_name_or_path=model_name_or_path,
             do_train=True,
             finetuning_type=finetuning_type,
+            lora_target=DEFAULT_MODULE.get(model_name.split("-")[0], None) or "q_proj,v_proj",
             prompt_template=template,
             dataset=",".join(dataset),
             dataset_dir=dataset_dir,