support mm llamapro (#2738)

Author: tastelikefeet

swift/tuners/llamapro.py

@@ -1,12 +1,12 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 from copy import deepcopy
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, fields
 from typing import Optional
 
 import torch
 from torch import nn
 
-from swift.llm import MODEL_ARCH_MAPPING, ModelKeys
+from swift.llm import MODEL_ARCH_MAPPING, HfConfigFactory, ModelKeys
 from swift.utils.logger import get_logger
 from .utils import ActivationMixin, SwiftAdapter, SwiftConfig, SwiftOutput
 
@@ -46,11 +46,9 @@ class LLaMAPro(SwiftAdapter):
     @staticmethod
     def prepare_model(model: nn.Module, config: LLaMAProConfig, adapter_name: str) -> SwiftOutput:
         """Prepare a model with `LLaMAProConfig`"""
-        num_hidden_layers = None
-        if hasattr(model.config, 'num_hidden_layers'):
-            num_hidden_layers = model.config.num_hidden_layers
-        elif hasattr(model.config, 'num_layers'):
-            num_hidden_layers = model.config.num_layers
+        num_hidden_layers = HfConfigFactory.get_config_attr(model.config, 'num_hidden_layers')
+        if num_hidden_layers is None:
+            num_hidden_layers = HfConfigFactory.get_config_attr(model.config, 'num_layers')
 
         assert num_hidden_layers is not None, 'Cannot find num of layers config'
         assert num_hidden_layers % config.num_new_blocks == 0, f'Model layers {num_hidden_layers} ' \
@@ -60,8 +58,26 @@ def prepare_model(model: nn.Module, config: LLaMAProConfig, adapter_name: str) -
 
         num_stride = num_hidden_layers // config.num_groups
 
-        # We only support decoder only model for now.
-        module_list = LLaMAPro._find_module_list(config, model)
+        try:
+            module_list = LLaMAPro._find_module_list(config, model)
+        except AssertionError as e:
+            model_type = LLaMAPro.search_correct_model_type(model)
+            if model_type is None:
+                language_model_name = SwiftAdapter.get_model_key_mapping(config.model_type, config).language_model
+                if language_model_name:
+                    if isinstance(language_model_name, str):
+                        language_model_name = [language_model_name]
+                    language_model = model.get_submodule(language_model_name[0])
+                    model_type = LLaMAPro.search_correct_model_type(language_model)
+                    if model_type:
+                        model = language_model
+
+            if model_type:
+                config.model_type = model_type
+                module_list = LLaMAPro._find_module_list(config, model)
+            else:
+                raise e
+
         new_module_list = nn.ModuleList()
         new_module_idx = []
         for idx, module in enumerate(module_list):
@@ -107,7 +123,10 @@ def _update_module_attr(config: LLaMAProConfig, module_list):
         if model_type in ('llama', 'mistral', 'qwen2', 'yi', 'gemma', 'deepseek', 'openbuddy', 'xverse', 'orion',
                           'bluelm', 'ziya', 'skywork', 'deepseek-v2', 'minicpm', 'phi3', 'internlm2'):
             for idx, module in enumerate(module_list):
-                getattr(module, attention).layer_idx = idx
+                try:
+                    getattr(module, attention).layer_idx = idx
+                except AttributeError:
+                    getattr(module, 'cross_attn').layer_idx = idx
         elif model_type in ('chatglm', 'glm4'):
             for idx, module in enumerate(module_list):
                 getattr(module, attention).layer_number = idx
@@ -135,6 +154,34 @@ def get_model_key_mapping(cls, model_type, config) -> ModelKeys:
             'LLaMAPro only support models with o_proj and down_proj components.'
         return model_key_mapping
 
+    @classmethod
+    def search_correct_model_type(cls, module: nn.Module):
+        for arch_name, arch_type in MODEL_ARCH_MAPPING.items():
+            arch_type: ModelKeys
+            if getattr(arch_type, 'module_list') is None:
+                # Need to be a LLM arch
+                continue
+
+            matched = True
+            for f in fields(arch_type):
+                arch_str = getattr(arch_type, f.name)
+                if f.name == 'arch_name' or arch_str is None:
+                    continue
+
+                arch_str = arch_str.replace('{}', '0')
+                try:
+                    sub_module = module.get_submodule(arch_str)
+                    if sub_module is None:
+                        matched = False
+                except AttributeError:
+                    matched = False
+
+                if not matched:
+                    break
+
+            if matched:
+                return arch_name
+
     @staticmethod
     def _update_module_weight(config: LLaMAProConfig, module_list, new_module_idx):
         model_key_mapping = LLaMAPro.get_model_key_mapping(config.model_type, config)