[train] support omni seq_cls (#5329)

Author: Jintao-Huang

docs/source/Customization/插件化.md

@@ -120,7 +120,7 @@ class IA3(Tuner):
 
     @staticmethod
     def prepare_model(args: 'TrainArguments', model: torch.nn.Module) -> torch.nn.Module:
-        model_arch: ModelKeys = MODEL_ARCH_MAPPING[model.model_meta.model_arch]
+        model_arch: ModelKeys = model.model_meta.model_arch
         ia3_config = IA3Config(
             target_modules=find_all_linears(model), feedforward_modules='.*' + model_arch.mlp.split('{}.')[1] + '.*')
         return get_peft_model(model, ia3_config)

docs/source_en/Customization/Pluginization.md

@@ -136,7 +136,7 @@ class IA3(Tuner):
 
     @staticmethod
     def prepare_model(args: 'TrainArguments', model: torch.nn.Module) -> torch.nn.Module:
-        model_arch: ModelKeys = MODEL_ARCH_MAPPING[model.model_meta.model_arch]
+        model_arch: ModelKeys = model.model_meta.model_arch
         ia3_config = IA3Config(
             target_modules=find_all_linears(model), feedforward_modules='.*' + model_arch.mlp.split('{}.')[1] + '.*')
         return get_peft_model(model, ia3_config)

examples/notebook/qwen2vl-ocr/ocr-sft.ipynb

@@ -36,7 +36,7 @@
     "os.environ['CUDA_VISIBLE_DEVICES'] = '0'\n",
     "\n",
     "from swift.llm import (\n",
-    "    get_model_tokenizer, load_dataset, get_template, EncodePreprocessor, get_model_arch,\n",
+    "    get_model_tokenizer, load_dataset, get_template, EncodePreprocessor,\n",
     "    get_multimodal_target_regex, LazyLLMDataset\n",
     ")\n",
     "from swift.utils import get_logger, get_model_parameter_info, plot_images, seed_everything\n",

examples/train/multimodal/lora_llm_full_vit/custom_plugin.py

@@ -4,7 +4,7 @@
 import safetensors.torch
 import torch
 
-from swift.llm import deep_getattr, get_model_arch, get_multimodal_target_regex
+from swift.llm import deep_getattr, get_multimodal_target_regex
 from swift.plugin import Tuner, extra_tuners
 from swift.tuners import LoraConfig, Swift
 from swift.utils import get_logger
@@ -46,14 +46,14 @@ def save_pretrained(
                     state_dict[n] = p.detach().cpu()
         model.save_pretrained(save_directory, state_dict=state_dict, safe_serialization=safe_serialization, **kwargs)
         # vit
-        model_arch = get_model_arch(model.model_meta.model_arch)
+        model_arch = model.model_meta.model_arch
         state_dict = {k: v for k, v in state_dict.items() if is_vit_param(model_arch, k)}
         safetensors.torch.save_file(
             state_dict, os.path.join(save_directory, 'vit.safetensors'), metadata={'format': 'pt'})
 
     @staticmethod
     def prepare_model(args: 'TrainArguments', model: torch.nn.Module) -> torch.nn.Module:
-        model_arch = get_model_arch(model.model_meta.model_arch)
+        model_arch = model.model_meta.model_arch
         target_regex = get_multimodal_target_regex(model)
         logger.info(f'target_regex: {target_regex}')
         lora_config = LoraConfig(

swift/llm/argument/base_args/quant_args.py

@@ -89,10 +89,9 @@ def get_quantization_config(self):
         return quantization_config
 
     def get_modules_to_not_convert(self):
-        from swift.llm import get_model_arch
         if not hasattr(self, 'model_meta') or not hasattr(self, 'model_info'):
             return None
-        model_arch = get_model_arch(self.model_meta.model_arch)
+        model_arch = self.model_meta.model_arch
         res = []
         if self.model_info.is_moe_model:
             res += ['mlp.gate', 'mlp.shared_expert_gate']

swift/llm/argument/tuner_args.py

@@ -4,7 +4,6 @@
 
 from transformers.utils import strtobool
 
-from swift.llm import get_model_arch
 from swift.utils import get_logger
 
 logger = get_logger()
@@ -204,7 +203,7 @@ def __post_init__(self):
             self.target_modules = self.target_regex
 
     def _init_multimodal_full(self):
-        model_arch = get_model_arch(self.model_meta.model_arch)
+        model_arch = self.model_meta.model_arch
         if not self.model_meta.is_multimodal or not model_arch or self.train_type != 'full':
             return
         if self.freeze_llm:

swift/llm/export/quant.py

@@ -7,8 +7,8 @@
 import torch.nn as nn
 from tqdm import tqdm
 
-from swift.llm import (ExportArguments, HfConfigFactory, MaxLengthError, ProcessorMixin, deep_getattr, get_model_arch,
-                       load_dataset, prepare_model_template, save_checkpoint, to_device)
+from swift.llm import (ExportArguments, HfConfigFactory, MaxLengthError, ProcessorMixin, deep_getattr, load_dataset,
+                       prepare_model_template, save_checkpoint, to_device)
 from swift.utils import get_logger, get_model_parameter_info
 
 logger = get_logger()
@@ -160,7 +160,7 @@ def awq_model_quantize(self) -> None:
                 self.tokenizer, quant_config=quant_config, n_parallel_calib_samples=args.quant_batch_size)
         quantizer.get_calib_dataset = _origin_get_calib_dataset  # recover
         if self.model.quant_config.modules_to_not_convert:
-            model_arch = get_model_arch(args.model_meta.model_arch)
+            model_arch = args.model_meta.model_arch
             lm_head_key = getattr(model_arch, 'lm_head', None) or 'lm_head'
             if lm_head_key not in self.model.quant_config.modules_to_not_convert:
                 self.model.quant_config.modules_to_not_convert.append(lm_head_key)
@@ -180,7 +180,7 @@ def _patch_gptq(self):
 
     @staticmethod
     def get_block_name_to_quantize(model: nn.Module) -> Optional[str]:
-        model_arch = get_model_arch(model.model_meta.model_arch)
+        model_arch = model.model_meta.model_arch
         prefix = ''
         if hasattr(model_arch, 'language_model'):
             assert len(model_arch.language_model) == 1, f'mllm_arch.language_model: {model_arch.language_model}'

swift/llm/model/patcher.py

@@ -150,12 +150,30 @@ def _check_imports(filename) -> List[str]:
         td.check_imports = _old_check_imports
 
 
+def get_lm_head_model(model, model_meta, lm_heads):
+    llm_prefix_list = getattr(model_meta.model_arch, 'language_model', None)
+    prefix_list = []
+    if llm_prefix_list:
+        prefix_list = llm_prefix_list[0].split('.')
+
+    origin_model = model
+    current_model = model
+    for prefix in [None] + prefix_list:
+        if prefix:
+            current_model = getattr(current_model, prefix)
+        for lm_head in lm_heads:
+            if hasattr(current_model, lm_head):
+                return current_model
+
+    raise ValueError(f'Cannot find the lm_head. model: {origin_model}')
+
+
 def _patch_sequence_classification(model, model_meta):
     hidden_size = HfConfigFactory.get_config_attr(model.config, 'hidden_size')
     initializer_range = HfConfigFactory.get_config_attr(model.config, 'initializer_range')
 
     lm_heads = ['lm_head', 'output', 'embed_out', 'output_layer']
-    llm_model = get_llm_model(model, model_meta=model_meta)
+    llm_model = get_lm_head_model(model, model_meta, lm_heads)
     llm_model.num_labels = model.config.num_labels
     llm_model.score = nn.Linear(hidden_size, llm_model.num_labels, bias=False, dtype=llm_model.dtype)
     if llm_model.score.weight.device == torch.device('meta'):

swift/llm/model/register.py

@@ -117,6 +117,7 @@ def register_model(model_meta: ModelMeta, *, exist_ok: bool = False) -> None:
     model_type: The unique ID for the model type. Models with the same model_type share
         the same architectures, template, get_function, etc.
     """
+    from .model_arch import get_model_arch
     model_type = model_meta.model_type
     if not exist_ok and model_type in MODEL_MAPPING:
         raise ValueError(f'The `{model_type}` has already been registered in the MODEL_MAPPING.')
@@ -125,6 +126,8 @@ def register_model(model_meta: ModelMeta, *, exist_ok: bool = False) -> None:
         model_meta.is_multimodal = True
     if model_type in RMModelType.__dict__:
         model_meta.is_reward = True
+    if model_meta.model_arch:
+        model_meta.model_arch = get_model_arch(model_meta.model_arch)
     MODEL_MAPPING[model_type] = model_meta
 
 

swift/llm/model/utils.py

@@ -351,9 +351,8 @@ def git_clone_github(github_url: str,
 
 
 def get_llm_model(model: torch.nn.Module, model_meta=None):
-    from swift import SwiftModel
+    from swift.tuners import SwiftModel
     from peft import PeftModel
-    from swift.llm import get_model_arch
     from accelerate.utils import extract_model_from_parallel
     model = extract_model_from_parallel(model)
 
@@ -362,7 +361,7 @@ def get_llm_model(model: torch.nn.Module, model_meta=None):
     if model_meta is None:
         model_meta = model.model_meta
 
-    llm_prefix = getattr(get_model_arch(model_meta.model_arch), 'language_model', None)
+    llm_prefix = getattr(model_meta.model_arch, 'language_model', None)
     if llm_prefix:
         llm_model = deep_getattr(model, llm_prefix[0])
     else: