lint

swift/llm/model/model/baai.py

@@ -5,9 +5,9 @@
 
 from modelscope import AutoConfig, AutoModel
 
-from ..model_arch import ModelArch
 from swift.llm import TemplateType
 from ..constant import MLLMModelType
+from ..model_arch import ModelArch
 from ..register import Model, ModelGroup, ModelMeta, get_model_tokenizer_with_flash_attn, register_model
 from ..utils import ModelInfo, git_clone_github, safe_snapshot_download
 

swift/llm/model/model/baichuan.py

@@ -7,10 +7,10 @@
 from torch import Tensor
 from transformers import BitsAndBytesConfig, PretrainedConfig
 
-from ..model_arch import ModelArch
 from swift.llm import TemplateType
 from swift.utils import get_logger
 from ..constant import LLMModelType
+from ..model_arch import ModelArch
 from ..register import Model, ModelGroup, ModelMeta, get_model_tokenizer_with_flash_attn, register_model
 from ..utils import ModelInfo
 

swift/llm/model/model/deepseek.py

@@ -3,9 +3,9 @@
 import sys
 from typing import Any, Dict
 
-from ..model_arch import ModelArch
 from swift.llm import TemplateType
 from ..constant import LLMModelType, MLLMModelType
+from ..model_arch import ModelArch
 from ..patcher import patch_output_clone, patch_output_to_input_device
 from ..register import Model, ModelGroup, ModelMeta, get_model_tokenizer_with_flash_attn, register_model
 from ..utils import ModelInfo, git_clone_github, use_submodel_func

swift/llm/model/model/gemma.py

@@ -1,9 +1,9 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 from typing import Any, Dict
 
-from ..model_arch import ModelArch
 from swift.llm import TemplateType
 from ..constant import LLMModelType, MLLMModelType
+from ..model_arch import ModelArch
 from ..register import Model, ModelGroup, ModelMeta, get_model_tokenizer_with_flash_attn, register_model
 from ..utils import ModelInfo
 

swift/llm/model/model/glm.py

@@ -3,17 +3,16 @@
 
 import torch
 import transformers
-from modelscope import AutoConfig
+from modelscope import AutoConfig, AutoTokenizer
 from packaging import version
-from modelscope import AutoTokenizer
 from transformers import PreTrainedTokenizerBase
 from transformers.dynamic_module_utils import get_class_from_dynamic_module
 from transformers.models.auto.tokenization_auto import get_tokenizer_config
 
-from ..model_arch import ModelArch
 from swift.llm import TemplateType
 from swift.utils import get_dist_setting, get_logger
 from ..constant import LLMModelType, MLLMModelType
+from ..model_arch import ModelArch
 from ..patcher import patch_output_to_input_device
 from ..register import Model, ModelGroup, ModelMeta, get_model_tokenizer_with_flash_attn, register_model
 from ..utils import AttnImpl, ModelInfo

swift/llm/model/model/internlm.py

@@ -6,9 +6,9 @@
 from transformers import PretrainedConfig
 from transformers.dynamic_module_utils import get_class_from_dynamic_module
 
-from ..model_arch import ModelArch
 from swift.llm import TemplateType
 from ..constant import LLMModelType, MLLMModelType
+from ..model_arch import ModelArch
 from ..patcher import patch_output_clone, patch_output_to_input_device
 from ..register import Model, ModelGroup, ModelMeta, get_model_tokenizer_with_flash_attn, register_model
 from ..utils import ModelInfo, use_submodel_func

swift/llm/model/model/llama.py

@@ -5,9 +5,9 @@
 
 from modelscope import AutoConfig
 
-from ..model_arch import ModelArch
 from swift.llm import TemplateType
 from ..constant import LLMModelType, MLLMModelType
+from ..model_arch import ModelArch
 from ..register import (Model, ModelGroup, ModelMeta, get_model_tokenizer_multimodal,
                         get_model_tokenizer_with_flash_attn, register_model)
 from ..utils import ModelInfo, git_clone_github

swift/llm/model/model/llava.py

@@ -7,9 +7,9 @@
 from modelscope import AutoConfig
 from transformers import PretrainedConfig
 
-from ..model_arch import ModelArch
 from swift.llm import TemplateType
 from ..constant import MLLMModelType
+from ..model_arch import ModelArch
 from ..register import Model, ModelGroup, ModelMeta, get_model_tokenizer_with_flash_attn, register_model
 from ..utils import ModelInfo, git_clone_github, safe_snapshot_download
 

swift/llm/model/model/mamba.py

@@ -1,10 +1,10 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 from typing import Any, Dict
 
-from ..model_arch import ModelArch
 from swift.llm import TemplateType
 from swift.utils import get_logger
 from ..constant import LLMModelType
+from ..model_arch import ModelArch
 from ..register import Model, ModelGroup, ModelMeta, get_model_tokenizer_from_local, register_model
 from ..utils import ModelInfo
 

swift/llm/model/model/microsoft.py

@@ -2,10 +2,10 @@
 from types import MethodType
 from typing import Any, Dict
 
-from ..model_arch import ModelArch
 from swift.llm import TemplateType
 from swift.utils import get_env_args
 from ..constant import LLMModelType, MLLMModelType
+from ..model_arch import ModelArch
 from ..patcher import patch_output_clone
 from ..register import Model, ModelGroup, ModelMeta, get_model_tokenizer_with_flash_attn, register_model
 from ..utils import ModelInfo, ignore_check_imports, use_submodel_func

swift/llm/model/model/minicpm.py

@@ -6,9 +6,9 @@
 from transformers import PretrainedConfig
 from transformers.dynamic_module_utils import get_class_from_dynamic_module
 
-from ..model_arch import ModelArch
 from swift.llm import TemplateType
 from ..constant import LLMModelType, MLLMModelType
+from ..model_arch import ModelArch
 from ..patcher import patch_fixed_device, patch_output_clone
 from ..register import Model, ModelGroup, ModelMeta, get_model_tokenizer_with_flash_attn, register_model
 from ..utils import ModelInfo, ignore_check_imports, use_submodel_func

swift/llm/model/model/mistral.py

@@ -1,8 +1,8 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 
-from ..model_arch import ModelArch
 from swift.llm import TemplateType
 from ..constant import LLMModelType
+from ..model_arch import ModelArch
 from ..register import Model, ModelGroup, ModelMeta, get_model_tokenizer_with_flash_attn, register_model
 
 register_model(

swift/llm/model/model/models.py

@@ -10,10 +10,10 @@
 from transformers import AutoTokenizer
 from transformers.dynamic_module_utils import get_class_from_dynamic_module
 
-from ..model_arch import ModelArch
 from swift.llm import TemplateType
 from swift.utils import get_logger
 from ..constant import LLMModelType, MLLMModelType
+from ..model_arch import ModelArch
 from ..patcher import patch_output_clone
 from ..register import Model, ModelGroup, ModelMeta, get_model_tokenizer_with_flash_attn, register_model
 from ..utils import ModelInfo, git_clone_github, use_submodel_func
@@ -219,7 +219,8 @@ def get_model_tokenizer_yuan(model_dir: str,
                 Model('OrionStarAI/Orion-14B-Base', 'OrionStarAI/Orion-14B-Base'),
                 Model('OrionStarAI/Orion-14B-Chat', 'OrionStarAI/Orion-14B-Chat'),
             ],
-                       ignore_file_pattern=[r'.+\.gguf$'], requires=['transformers==4.34.1']),
+                       ignore_file_pattern=[r'.+\.gguf$'],
+                       requires=['transformers==4.34.1']),
         ],
         TemplateType.orion,
         get_model_tokenizer_with_flash_attn,
@@ -253,15 +254,15 @@ def get_model_tokenizer_idefics(model_dir: str, *args, **kwargs):
         architectures=['Idefics3ForConditionalGeneration'],
     ))
 
-
 register_model(
     ModelMeta(
         LLMModelType.wizardlm2_moe,
         [
             ModelGroup([
                 Model('AI-ModelScope/WizardLM-2-8x22B', 'alpindale/WizardLM-2-8x22B'),
             ],
-                       requires=['transformers>=4.36'], tags=['skip_test']),
+                       requires=['transformers>=4.36'],
+                       tags=['skip_test']),
         ],
         TemplateType.wizardlm2,
         get_model_tokenizer_with_flash_attn,
@@ -276,7 +277,8 @@ def get_model_tokenizer_idefics(model_dir: str, *args, **kwargs):
             ModelGroup([
                 Model('AI-ModelScope/WizardLM-2-7B-AWQ', 'MaziyarPanahi/WizardLM-2-7B-AWQ'),
             ],
-                       requires=['transformers>=4.34'], tags=['skip_test']),
+                       requires=['transformers>=4.34'],
+                       tags=['skip_test']),
         ],
         TemplateType.wizardlm2_awq,
         get_model_tokenizer_with_flash_attn,
@@ -304,7 +306,8 @@ def get_model_tokenizer_idefics(model_dir: str, *args, **kwargs):
         [
             ModelGroup([
                 Model('OpenBuddy/openbuddy-deepseek-67b-v15.2', 'OpenBuddy/openbuddy-deepseek-67b-v15.2'),
-            ], tags=['skip_test']),
+            ],
+                       tags=['skip_test']),
         ],
         TemplateType.openbuddy,
         get_model_tokenizer_with_flash_attn,
@@ -412,7 +415,8 @@ def get_model_tokenizer_idefics(model_dir: str, *args, **kwargs):
         [
             ModelGroup([
                 Model('OpenBuddy/openbuddy-llama-65b-v8-bf16', 'OpenBuddy/openbuddy-llama-65b-v8-bf16'),
-            ], tags=['skip_test']),
+            ],
+                       tags=['skip_test']),
         ],
         TemplateType.openbuddy,
         get_model_tokenizer_with_flash_attn,
@@ -522,7 +526,8 @@ def update(self, key_states: torch.Tensor, value_states: torch.Tensor, layer_idx
                 Model('AI-ModelScope/Llama-3.1-Nemotron-70B-Instruct-HF', 'nvidia/Llama-3.1-Nemotron-70B-Instruct-HF'),
             ],
                        requires=['transformers>=4.43'],
-                       ignore_file_pattern=[r'.+\.pth$'], tags=['skip_test']),
+                       ignore_file_pattern=[r'.+\.pth$'],
+                       tags=['skip_test']),
         ],
         TemplateType.llama3,
         get_model_tokenizer_with_flash_attn,
@@ -537,7 +542,8 @@ def update(self, key_states: torch.Tensor, value_states: torch.Tensor, layer_idx
             ModelGroup([
                 Model('LLM-Research/Reflection-Llama-3.1-70B', 'mattshumer/Reflection-Llama-3.1-70B'),
             ],
-                       requires=['transformers>=4.43'], tags=['skip_test']),
+                       requires=['transformers>=4.43'],
+                       tags=['skip_test']),
         ],
         TemplateType.reflection,
         get_model_tokenizer_with_flash_attn,
@@ -639,7 +645,8 @@ def get_model_tokenizer_got_ocr2(*args, **kwargs):
                 Model('xverse/XVERSE-65B-2', 'xverse/XVERSE-65B-2'),
                 Model('xverse/XVERSE-65B-Chat', 'xverse/XVERSE-65B-Chat'),
                 Model('xverse/XVERSE-13B-256K', 'xverse/XVERSE-13B-256K', ms_revision='v1.0.0'),
-            ], requires=['transformers==4.38.2']),
+            ],
+                       requires=['transformers==4.38.2']),
         ],
         TemplateType.xverse,
         get_model_tokenizer_with_flash_attn,
@@ -653,7 +660,9 @@ def get_model_tokenizer_got_ocr2(*args, **kwargs):
         [
             ModelGroup([
                 Model('xverse/XVERSE-MoE-A4.2B', 'xverse/XVERSE-MoE-A4.2B'),
-            ], tags=['moe', 'skip_test'], requires=['transformers==4.38.2']),
+            ],
+                       tags=['moe', 'skip_test'],
+                       requires=['transformers==4.38.2']),
         ],
         TemplateType.xverse,
         get_model_tokenizer_with_flash_attn,

swift/llm/model/model/qwen.py

@@ -466,12 +466,14 @@ def _get_cast_dtype(self) -> torch.dtype:
                 Model('Qwen/Qwen1.5-MoE-A2.7B', 'Qwen/Qwen1.5-MoE-A2.7B'),
                 Model('Qwen/Qwen1.5-MoE-A2.7B-Chat', 'Qwen/Qwen1.5-MoE-A2.7B-Chat'),
                 Model('Qwen/Qwen1.5-MoE-A2.7B-Chat-GPTQ-Int4', 'Qwen/Qwen1.5-MoE-A2.7B-Chat-GPTQ-Int4'),
-            ], tags=['skip_test']),
+            ],
+                       tags=['skip_test']),
             ModelGroup([
                 Model('Qwen/Qwen2-57B-A14B', 'Qwen/Qwen2-57B-A14B'),
                 Model('Qwen/Qwen2-57B-A14B-Instruct', 'Qwen/Qwen2-57B-A14B-Instruct'),
                 Model('Qwen/Qwen2-57B-A14B-Instruct-GPTQ-Int4', 'Qwen/Qwen2-57B-A14B-Instruct-GPTQ-Int4'),
-            ], tags=['skip_test'])
+            ],
+                       tags=['skip_test'])
         ],
         TemplateType.qwen,
         get_model_tokenizer_with_flash_attn,

swift/llm/model/model/telechat.py

@@ -1,8 +1,8 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 
-from ..model_arch import ModelArch
 from swift.llm import TemplateType
 from ..constant import LLMModelType
+from ..model_arch import ModelArch
 from ..register import Model, ModelGroup, ModelMeta, get_model_tokenizer_with_flash_attn, register_model
 
 register_model(

swift/llm/model/model/yi.py

@@ -7,11 +7,11 @@
 from transformers import PretrainedConfig
 
 from swift.llm import TemplateType
-from ..model_arch import ModelArch
 from swift.utils import get_logger
 from ..constant import LLMModelType, MLLMModelType
+from ..model_arch import ModelArch
 from ..register import Model, ModelGroup, ModelMeta, get_model_tokenizer_with_flash_attn, register_model
-from ..utils import git_clone_github, ModelInfo
+from ..utils import ModelInfo, git_clone_github
 
 logger = get_logger()
 

swift/llm/template/template/__init__.py

@@ -1,19 +1,3 @@
-from . import llama
-from . import qwen
-from . import deepseek
-from . import emu3
-from . import gemma
-from . import glm
-from . import got_ocr
-from . import idefics3
-from . import internlm
-from . import internvl
-from . import llava
 # from . import llm
-from . import microsoft
-from . import minicpm
-from . import mplug
-from . import openbuddy
-from . import pixtral
-from . import rlhf
-from . import yi
+from . import (deepseek, emu3, gemma, glm, got_ocr, idefics3, internlm, internvl, llama, llava, microsoft, minicpm,
+               mplug, openbuddy, pixtral, qwen, rlhf, yi)

swift/llm/tuner.py

@@ -1,9 +1,9 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
-import json
 import os
 import types
 from typing import List
 
+import json
 import numpy as np
 import torch
 import transformers
@@ -15,8 +15,8 @@
 from swift.plugin.callback import extra_callbacks
 from swift.plugin.optimizer import optimizers_map
 from swift.plugin.tuner import Tuner, extra_tuners
-from swift.tuners import (AdaLoraConfig, AdapterConfig, BOFTConfig, LLaMAProConfig, LongLoRAModelType,
-                          LoraConfig, LoRAConfig, ReftConfig, Swift, VeraConfig)
+from swift.tuners import (AdaLoraConfig, AdapterConfig, BOFTConfig, LLaMAProConfig, LongLoRAModelType, LoraConfig,
+                          LoRAConfig, ReftConfig, Swift, VeraConfig)
 from swift.utils import activate_model_parameters, freeze_model_parameters, get_logger, use_torchacc
 from swift.utils.torch_utils import find_all_linears, find_embedding
 

swift/tuners/utils.py

@@ -20,7 +20,7 @@
 from peft.utils import ModulesToSaveWrapper as _ModulesToSaveWrapper
 from peft.utils import _get_submodules
 
-from swift.llm import ModelKeys, MODEL_ARCH_MAPPING
+from swift.llm import MODEL_ARCH_MAPPING, ModelKeys
 from swift.utils.constants import BIN_EXTENSIONS
 from swift.utils.logger import get_logger
 

tests/llm/load_model.py

@@ -1,12 +1,24 @@
-from swift.llm import get_model_tokenizer
 import argparse
+from dataclasses import fields
 
+import torch
 
-def get_model_and_tokenizer(ms_model_id):
+from swift.llm import MODEL_ARCH_MAPPING, ModelKeys, get_model_tokenizer
+
+
+def get_model_and_tokenizer(ms_model_id, model_arch=None):
     try:
         import transformers
         print(f'Test model: {ms_model_id} with transformers version: {transformers.__version__}')
         model_ins, tokenizer = get_model_tokenizer(ms_model_id)
+        model_ins: torch.nn.Module
+        if model_arch:
+            model_arch: ModelKeys = MODEL_ARCH_MAPPING[model_arch]
+            for f in fields(model_arch):
+                value = getattr(model_arch, f.name)
+                if value is not None:
+                    value = value.replace('{}', '[0]')
+                    model_ins.get_submodule(value)
     except Exception:
         import traceback
         print(traceback.format_exc())
@@ -19,7 +31,11 @@ def get_model_and_tokenizer(ms_model_id):
     type=str,
     required=True,
 )
+parser.add_argument(
+    '--model_arch',
+    type=str,
+    required=True,
+)
 args = parser.parse_args()
 
-get_model_and_tokenizer(args.ms_model_id)
-
+get_model_and_tokenizer(args.ms_model_id, args.model_arch)