fix #3694

Former-commit-id: 2a67ab3

Author: hiyouga

examples/extras/llama_pro/llama3_freeze_sft.yaml

@@ -5,8 +5,8 @@ model_name_or_path: models/llama3-8b-instruct-pro
 stage: sft
 do_train: true
 finetuning_type: freeze
-name_module_trainable: all
-num_layer_trainable: 8
+freeze_trainable_layers: 8
+freeze_trainable_modules: all
 use_llama_pro: true
 
 # dataset

scripts/llama_pro.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Performs block expansion for LLaMA, Mistral or Qwen1.5 models.
+# Performs block expansion for LLaMA, Mistral, Qwen1.5 or Yi models.
 # Usage: python llama_pro.py --model_name_or_path meta-llama/Llama-2-7b-hf --output_dir llama2_pro --num_expand 8
 # Inspired by: https://github.com/TencentARC/LLaMA-Pro/blob/main/scripts/block_expansion.py
 
@@ -106,8 +106,7 @@ def block_expansion(
     print("Fine-tune this model with:")
     print("  --model_name_or_path {} \\".format(output_dir))
     print("  --finetuning_type freeze \\")
-    print("  --name_module_trainable all \\")
-    print("  --num_layer_trainable {} \\".format(num_expand))
+    print("  --freeze_trainable_layers {} \\".format(num_expand))
     print("  --use_llama_pro")
 
 

src/llmtuner/hparams/finetuning_args.py

@@ -1,5 +1,4 @@
-import json
-from dataclasses import asdict, dataclass, field
+from dataclasses import dataclass, field
 from typing import Literal, Optional
 
 
@@ -9,22 +8,40 @@ class FreezeArguments:
     Arguments pertaining to the freeze (partial-parameter) training.
     """
 
-    name_module_trainable: str = field(
+    freeze_trainable_layers: int = field(
+        default=2,
+        metadata={
+            "help": (
+                "The number of trainable layers for freeze (partial-parameter) fine-tuning. "
+                "Positive numbers mean the last n layers are set as trainable, "
+                "negative numbers mean the first n layers are set as trainable."
+            )
+        },
+    )
+    freeze_trainable_modules: str = field(
         default="all",
         metadata={
-            "help": """Name of trainable modules for partial-parameter (freeze) fine-tuning. \
-                    Use commas to separate multiple modules. \
-                    Use "all" to specify all the available modules. \
-                    LLaMA choices: ["mlp", "self_attn"], \
-                    BLOOM & Falcon & ChatGLM choices: ["mlp", "self_attention"], \
-                    Qwen choices: ["mlp", "attn"], \
-                    InternLM2 choices: ["feed_forward", "attention"], \
-                    Others choices: the same as LLaMA."""
+            "help": (
+                "Name(s) of trainable modules for freeze (partial-parameter) fine-tuning. "
+                "Use commas to separate multiple modules. "
+                "Use `all` to specify all the available modules. "
+                "LLaMA choices: [`mlp`, `self_attn`], "
+                "BLOOM & Falcon & ChatGLM choices: [`mlp`, `self_attention`], "
+                "Qwen choices: [`mlp`, `attn`], "
+                "InternLM2 choices: [`feed_forward`, `attention`], "
+                "Others choices: the same as LLaMA."
+            )
         },
     )
-    num_layer_trainable: int = field(
-        default=2,
-        metadata={"help": "The number of trainable layers for partial-parameter (freeze) fine-tuning."},
+    freeze_extra_modules: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": (
+                "Name(s) of modules apart from hidden layers to be set as trainable "
+                "for freeze (partial-parameter) fine-tuning. "
+                "Use commas to separate multiple modules."
+            )
+        },
     )
 
 
@@ -37,7 +54,11 @@ class LoraArguments:
     additional_target: Optional[str] = field(
         default=None,
         metadata={
-            "help": "Name(s) of modules apart from LoRA layers to be set as trainable and saved in the final checkpoint."
+            "help": (
+                "Name(s) of modules apart from LoRA layers to be set as trainable "
+                "and saved in the final checkpoint. "
+                "Use commas to separate multiple modules."
+            )
         },
     )
     lora_alpha: Optional[int] = field(
@@ -55,15 +76,17 @@ class LoraArguments:
     lora_target: str = field(
         default="all",
         metadata={
-            "help": """Name(s) of target modules to apply LoRA. \
-                    Use commas to separate multiple modules. \
-                    Use "all" to specify all the linear modules. \
-                    LLaMA choices: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], \
-                    BLOOM & Falcon & ChatGLM choices: ["query_key_value", "dense", "dense_h_to_4h", "dense_4h_to_h"], \
-                    Baichuan choices: ["W_pack", "o_proj", "gate_proj", "up_proj", "down_proj"], \
-                    Qwen choices: ["c_attn", "attn.c_proj", "w1", "w2", "mlp.c_proj"], \
-                    InternLM2 choices: ["wqkv", "wo", "w1", "w2", "w3"], \
-                    Others choices: the same as LLaMA."""
+            "help": (
+                "Name(s) of target modules to apply LoRA. "
+                "Use commas to separate multiple modules. "
+                "Use `all` to specify all the linear modules. "
+                "LLaMA choices: [`q_proj`, `k_proj`, `v_proj`, `o_proj`, `gate_proj`, `up_proj`, `down_proj`], "
+                "BLOOM & Falcon & ChatGLM choices: [`query_key_value`, `dense`, `dense_h_to_4h`, `dense_4h_to_h`], "
+                "Baichuan choices: [`W_pack`, `o_proj`, `gate_proj`, `up_proj`, `down_proj`], "
+                "Qwen choices: [`c_attn`, `attn.c_proj`, `w1`, `w2`, `mlp.c_proj`], "
+                "InternLM2 choices: [`wqkv`, `wo`, `w1`, `w2`, `w3`], "
+                "Others choices: the same as LLaMA."
+            )
         },
     )
     loraplus_lr_ratio: Optional[float] = field(
@@ -177,8 +200,10 @@ class GaloreArguments:
     galore_target: str = field(
         default="all",
         metadata={
-            "help": """Name(s) of modules to apply GaLore. Use commas to separate multiple modules. \
-                    Use "all" to specify all the linear modules."""
+            "help": (
+                "Name(s) of modules to apply GaLore. Use commas to separate multiple modules. "
+                "Use `all` to specify all the linear modules."
+            )
         },
     )
     galore_rank: int = field(
@@ -238,16 +263,20 @@ class BAdamArgument:
     badam_mask_mode: Literal["adjacent", "scatter"] = field(
         default="adjacent",
         metadata={
-            "help": """The mode of the mask for BAdam optimizer. \
-                    `adjacent` means that the trainable parameters are adjacent to each other, \
-                    `scatter` means that trainable parameters are randomly choosed from the weight."""
+            "help": (
+                "The mode of the mask for BAdam optimizer. "
+                "`adjacent` means that the trainable parameters are adjacent to each other, "
+                "`scatter` means that trainable parameters are randomly choosed from the weight."
+            )
         },
     )
     badam_verbose: int = field(
         default=0,
         metadata={
-            "help": """The verbosity level of BAdam optimizer. \
-                    0 for no print, 1 for print the block prefix, 2 for print trainable parameters"""
+            "help": (
+                "The verbosity level of BAdam optimizer. "
+                "0 for no print, 1 for print the block prefix, 2 for print trainable parameters."
+            )
         },
     )
 
@@ -285,7 +314,8 @@ def split_arg(arg):
                 return [item.strip() for item in arg.split(",")]
             return arg
 
-        self.name_module_trainable = split_arg(self.name_module_trainable)
+        self.freeze_trainable_modules = split_arg(self.freeze_trainable_modules)
+        self.freeze_extra_modules = split_arg(self.freeze_extra_modules)
         self.lora_alpha = self.lora_alpha or self.lora_rank * 2
         self.lora_target = split_arg(self.lora_target)
         self.additional_target = split_arg(self.additional_target)
@@ -315,17 +345,3 @@ def split_arg(arg):
 
         if self.loraplus_lr_ratio is not None and self.finetuning_type != "lora":
             raise ValueError("`loraplus_lr_ratio` is only valid for the LoRA training.")
-
-    def save_to_json(self, json_path: str):
-        r"""Saves the content of this instance in JSON format inside `json_path`."""
-        json_string = json.dumps(asdict(self), indent=2, sort_keys=True) + "\n"
-        with open(json_path, "w", encoding="utf-8") as f:
-            f.write(json_string)
-
-    @classmethod
-    def load_from_json(cls, json_path: str):
-        r"""Creates an instance from the content of `json_path`."""
-        with open(json_path, "r", encoding="utf-8") as f:
-            text = f.read()
-
-        return cls(**json.loads(text))

src/llmtuner/model/adapter.py

@@ -1,3 +1,4 @@
+import re
 from typing import TYPE_CHECKING
 
 import torch
@@ -68,37 +69,52 @@ def init_adapter(
             raise ValueError("Current model does not support freeze tuning.")
 
         if finetuning_args.use_llama_pro:
-            if num_layers % finetuning_args.num_layer_trainable != 0:
+            if num_layers % finetuning_args.freeze_trainable_layers != 0:
                 raise ValueError(
                     "`num_layers` {} should be divisible by `num_layer_trainable` {}.".format(
-                        num_layers, finetuning_args.num_layer_trainable
+                        num_layers, finetuning_args.freeze_trainable_layers
                     )
                 )
 
-            stride = num_layers // finetuning_args.num_layer_trainable
+            stride = num_layers // finetuning_args.freeze_trainable_layers
             trainable_layer_ids = range(stride - 1, num_layers + stride - 1, stride)
-        elif finetuning_args.num_layer_trainable > 0:  # fine-tuning the last n layers if num_layer_trainable > 0
-            trainable_layer_ids = range(num_layers - finetuning_args.num_layer_trainable, num_layers)
+        elif finetuning_args.freeze_trainable_layers > 0:  # fine-tuning the last n layers if num_layer_trainable > 0
+            trainable_layer_ids = range(max(0, num_layers - finetuning_args.freeze_trainable_layers), num_layers)
         else:  # fine-tuning the first n layers if num_layer_trainable < 0
-            trainable_layer_ids = range(-finetuning_args.num_layer_trainable)
+            trainable_layer_ids = range(min(-finetuning_args.freeze_trainable_layers, num_layers))
 
-        freeze_modules = {"all"}
-        for name, _ in model.named_modules():
+        hidden_modules = set()
+        non_hidden_modules = set()
+        for name, _ in model.named_parameters():
             if ".0." in name:
-                freeze_modules.add(name.split(".0.")[-1].split(".")[0])
+                hidden_modules.add(name.split(".0.")[-1].split(".")[0])
             elif ".1." in name:  # MoD starts from layer 1
-                freeze_modules.add(name.split(".1.")[-1].split(".")[0])
+                hidden_modules.add(name.split(".1.")[-1].split(".")[0])
+
+            if re.search(r"\.\d+\.", name) is None:
+                non_hidden_modules.add(name.split(".")[-2])
 
         trainable_layers = []
-        for module_name in finetuning_args.name_module_trainable:
-            if module_name not in freeze_modules:
+        for module_name in finetuning_args.freeze_trainable_modules:
+            if module_name != "all" and module_name not in hidden_modules:
                 raise ValueError(
-                    "Module {} is not found, please choose from {}".format(module_name, ", ".join(freeze_modules))
+                    "Module {} is not found, please choose from {}".format(module_name, ", ".join(hidden_modules))
                 )
 
             for idx in trainable_layer_ids:
                 trainable_layers.append(".{:d}.{}".format(idx, module_name if module_name != "all" else ""))
 
+        if finetuning_args.freeze_extra_modules:
+            for module_name in finetuning_args.freeze_extra_modules:
+                if module_name not in non_hidden_modules:
+                    raise ValueError(
+                        "Module {} is not found, please choose from {}".format(
+                            module_name, ", ".join(non_hidden_modules)
+                        )
+                    )
+
+                trainable_layers.append(module_name)
+
         for name, param in model.named_parameters():
             if any(trainable_layer in name for trainable_layer in trainable_layers):
                 if cast_trainable_params_to_fp32:

src/llmtuner/webui/components/train.py

@@ -124,13 +124,17 @@ def create_train_tab(engine: "Engine") -> Dict[str, "Component"]:
 
     with gr.Accordion(open=False) as freeze_tab:
         with gr.Row():
-            num_layer_trainable = gr.Slider(minimum=1, maximum=128, value=2, step=1)
-            name_module_trainable = gr.Textbox(value="all")
+            freeze_trainable_layers = gr.Slider(minimum=-128, maximum=128, value=2, step=1)
+            freeze_trainable_modules = gr.Textbox(value="all")
+            freeze_extra_modules = gr.Textbox()
 
-    input_elems.update({num_layer_trainable, name_module_trainable})
+    input_elems.update({freeze_trainable_layers, freeze_trainable_modules, freeze_extra_modules})
     elem_dict.update(
         dict(
-            freeze_tab=freeze_tab, num_layer_trainable=num_layer_trainable, name_module_trainable=name_module_trainable
+            freeze_tab=freeze_tab,
+            freeze_trainable_layers=freeze_trainable_layers,
+            freeze_trainable_modules=freeze_trainable_modules,
+            freeze_extra_modules=freeze_extra_modules,
         )
     )
 

src/llmtuner/webui/locales.py

@@ -572,24 +572,24 @@
             "label": "部分参数微调设置",
         },
     },
-    "num_layer_trainable": {
+    "freeze_trainable_layers": {
         "en": {
             "label": "Trainable layers",
-            "info": "The number of trainable layers.",
+            "info": "Number of the last(+)/first(-) hidden layers to be set as trainable.",
         },
         "ru": {
             "label": "Обучаемые слои",
-            "info": "Количество обучаемых слоев.",
+            "info": "Количество последних (+)/первых (-) скрытых слоев, которые будут установлены как обучаемые.",
         },
         "zh": {
             "label": "可训练层数",
-            "info": "可训练模型层的数量。",
+            "info": "最末尾（+）/最前端（-）可训练隐藏层的数量。",
         },
     },
-    "name_module_trainable": {
+    "freeze_trainable_modules": {
         "en": {
             "label": "Trainable modules",
-            "info": "The name of trainable modules. Use commas to separate multiple modules.",
+            "info": "Name(s) of trainable modules. Use commas to separate multiple modules.",
         },
         "ru": {
             "label": "Обучаемые модули",
@@ -600,6 +600,26 @@
             "info": "可训练模块的名称。使用英文逗号分隔多个名称。",
         },
     },
+    "freeze_extra_modules": {
+        "en": {
+            "label": "Extra modules (optional)",
+            "info": (
+                "Name(s) of modules apart from hidden layers to be set as trainable. "
+                "Use commas to separate multiple modules."
+            ),
+        },
+        "ru": {
+            "label": "Дополнительные модули (опционально)",
+            "info": (
+                "Имена модулей, кроме скрытых слоев, которые следует установить в качестве обучаемых. "
+                "Используйте запятые для разделения нескольких модулей."
+            ),
+        },
+        "zh": {
+            "label": "额外模块（非必填）",
+            "info": "除隐藏层以外的可训练模块名称。使用英文逗号分隔多个名称。",
+        },
+    },
     "lora_tab": {
         "en": {
             "label": "LoRA configurations",

src/llmtuner/webui/runner.py

@@ -146,8 +146,9 @@ def _parse_train_args(self, data: Dict["Component", Any]) -> Dict[str, Any]:
         )
 
         if args["finetuning_type"] == "freeze":
-            args["num_layer_trainable"] = get("train.num_layer_trainable")
-            args["name_module_trainable"] = get("train.name_module_trainable")
+            args["freeze_trainable_layers"] = get("train.freeze_trainable_layers")
+            args["freeze_trainable_modules"] = get("train.freeze_trainable_modules")
+            args["freeze_extra_modules"] = get("train.freeze_extra_modules") or None
         elif args["finetuning_type"] == "lora":
             args["lora_rank"] = get("train.lora_rank")
             args["lora_alpha"] = get("train.lora_alpha")