Improve the convert and skipped_module parts in lite module

Author: 43758726

lmdeploy/archs.py

@@ -1,5 +1,4 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-import os
 from typing import Literal
 
 from transformers import AutoConfig
@@ -138,9 +137,6 @@ def get_task(backend: str, model_path: str, trust_remote_code: bool = False):
     """Get pipeline type and pipeline class from model config."""
     from lmdeploy.serve.core import AsyncEngine
 
-    if os.path.exists(os.path.join(model_path, 'triton_models', 'weights')):
-        # workspace model
-        return 'llm', AsyncEngine
     _, config = get_model_arch(model_path, trust_remote_code=trust_remote_code)
     if check_vl_llm(backend, config.to_dict()):
         from lmdeploy.serve.core import VLAsyncEngine

lmdeploy/lite/apis/auto_awq.py

@@ -8,8 +8,8 @@
 import torch
 from torch import nn
 
-from lmdeploy.lite.apis.calibrate import LAYER_TYPE_MAP, MOE_MODEL_LIST, calibrate, load_model_and_tokenizer
-from lmdeploy.lite.moe_mlp_modules import CONVERT_MOE_MODELS
+from lmdeploy.lite.apis.calibrate import LAYER_TYPE_MAP, calibrate, load_model_and_tokenizer
+from lmdeploy.lite.model import MODELS
 from lmdeploy.lite.quantization.awq import FC_FCS_MAP, NORM_FCS_MAP, awq_layers, quant_weights, smooth_layers
 from lmdeploy.lite.utils import collect_target_modules
 from lmdeploy.utils import try_import_deeplink
@@ -87,26 +87,31 @@ def auto_awq(model: str,
         model = get_model(model, revision=revision, download_dir=download_dir)
     model_path = model
     if calib_samples == 0:
-        arch, vl_model, model, tokenizer, _, _ = load_model_and_tokenizer(model, dtype=dtype, work_dir=work_dir)
+        arch, vl_model, model, tokenizer, _, _ = load_model_and_tokenizer(model,
+                                                                          dtype=dtype,
+                                                                          work_dir=work_dir,
+                                                                          trust_remote_code=trust_remote_code)
     else:
-        arch, vl_model, model, tokenizer = calibrate(model,
-                                                         calib_dataset,
-                                                         calib_samples,
-                                                         calib_seqlen,
-                                                         work_dir,
-                                                         device,
-                                                         w_bits=w_bits,
-                                                         w_group_size=w_group_size,
-                                                         search_scale=search_scale,
-                                                         dtype=dtype,
-                                                         batch_size=batch_size)
+        arch, vl_model, model, tokenizer, _ = calibrate(model,
+                                                        calib_dataset,
+                                                        calib_samples,
+                                                        calib_seqlen,
+                                                        work_dir,
+                                                        device,
+                                                        w_bits=w_bits,
+                                                        w_group_size=w_group_size,
+                                                        search_scale=search_scale,
+                                                        dtype=dtype,
+                                                        batch_size=batch_size,
+                                                        trust_remote_code=trust_remote_code)
 
     layer_type = LAYER_TYPE_MAP[type(model).__name__]
     layers = collect_target_modules(model, layer_type)
     fcs = {}
+    model_layer = MODELS.get(arch)
     for l_name, layer in layers.items():
-        if arch in MOE_MODEL_LIST:
-            CONVERT_MOE_MODELS.get(arch)(layer)
+        if model_layer:
+            model_layer(layer)
         name2fc = collect_target_modules(layer, nn.Linear, prefix=l_name)
         fcs.update(name2fc)
 
@@ -122,15 +127,14 @@ def auto_awq(model: str,
             act_scales = input_stats['absmax']
             smooth_layers(layers, fc2fcs, norm2fcs, act_scales, w_group_size, device)
 
-    matched_exclude_modules = quant_weights(model, fcs, w_bits, w_sym, w_group_size, device,
-                                            arch=arch)
+    skipped_modules = quant_weights(model, fcs, w_bits, w_sym, arch, w_group_size, device)
     quantization_config = dict(quant_method='awq',
                                version='gemm',
                                bits=w_bits,
                                group_size=w_group_size,
                                zero_point=not w_sym)
-    if matched_exclude_modules:
-        quantization_config['modules_to_not_convert'] = matched_exclude_modules
+    if skipped_modules:
+        quantization_config['modules_to_not_convert'] = skipped_modules
     model.config.update(dict(quantization_config=quantization_config))
 
     if vl_model:

lmdeploy/lite/apis/calibrate.py

@@ -27,6 +27,7 @@
     'LlavaLlamaForCausalLM': 'LlamaDecoderLayer',
     'MGMLlamaForCausalLM': 'LlamaDecoderLayer',  # mini gemini
     'InternLMXComposer2ForCausalLM': 'InternLM2DecoderLayer',
+    'InternS2PreviewForConditionalGeneration': 'InternS2PreviewDecoderLayer',
     'Phi3ForCausalLM': 'Phi3DecoderLayer',
     'ChatGLMForConditionalGeneration': 'GLMBlock',
     'MixtralForCausalLM': 'MixtralDecoderLayer',
@@ -51,6 +52,7 @@
     'LlavaLlamaForCausalLM': 'LlamaRMSNorm',
     'MGMLlamaForCausalLM': 'LlamaRMSNorm',  # mini gemini
     'InternLMXComposer2ForCausalLM': 'InternLM2RMSNorm',
+    'InternS2PreviewForConditionalGeneration': 'InternS2PreviewRMSNorm',
     'Phi3ForCausalLM': 'Phi3RMSNorm',
     'ChatGLMForConditionalGeneration': 'RMSNorm',
     'MixtralForCausalLM': 'MixtralRMSNorm',
@@ -75,6 +77,7 @@
     'LlavaLlamaForCausalLM': 'lm_head',
     'MGMLlamaForCausalLM': 'lm_head',  # mini gemini
     'InternLMXComposer2ForCausalLM': 'output',
+    'InternS2PreviewForConditionalGeneration': 'lm_head',
     'Phi3ForCausalLM': 'lm_head',
     'ChatGLMForConditionalGeneration': 'output_layer',
     'MixtralForCausalLM': 'lm_head',
@@ -83,12 +86,6 @@
     'MistralForCausalLM': 'lm_head',
 }
 
-MOE_MODEL_LIST = [
-    'Qwen3MoeForCausalLM',
-    'Qwen3_5MoeForConditionalGeneration',
-    'MixtralForCausalLM'
-]
-
 
 def check_vl_llm(backend: str, config: dict) -> bool:
     """Check if the model is a vl model from model config."""
@@ -110,7 +107,8 @@ def check_vl_llm(backend: str, config: dict) -> bool:
         'Qwen3_5MoeForConditionalGeneration', 'MllamaForConditionalGeneration', 'MolmoForCausalLM',
         'Gemma3ForConditionalGeneration', 'Llama4ForConditionalGeneration', 'InternVLForConditionalGeneration',
         'InternS1ForConditionalGeneration', 'InternS1ProForConditionalGeneration',
-        'InternS1_1_ForConditionalGeneration', 'Glm4vForConditionalGeneration'
+        'InternS1_1_ForConditionalGeneration', 'Glm4vForConditionalGeneration',
+        'InternS2PreviewForConditionalGeneration'
     ])
     if arch == 'QWenLMHeadModel' and 'visual' in config:
         return True
@@ -125,11 +123,7 @@ def check_vl_llm(backend: str, config: dict) -> bool:
 
 def get_task(backend: str, model_path: str):
     """Get pipeline type and pipeline class from model config."""
-    import os
 
-    if os.path.exists(os.path.join(model_path, 'triton_models', 'weights')):
-        # workspace model
-        return 'llm'
     _, config = get_model_arch(model_path)
     if check_vl_llm(backend, config.to_dict()):
         return 'vlm'
@@ -260,22 +254,23 @@ def update_moe_mapping(model, model_type):
 
 def load_model_and_tokenizer(model: str,
                          dtype: Literal['float16', 'bfloat16', 'auto'] = 'auto',
-                         work_dir: str = './work_dir'):
+                         work_dir: str = './work_dir',
+                         trust_remote_code: bool = False):
     """Load model and tokenizer."""
     model_type = get_task(backend='turbomind', model_path=model)
     make_compatible_internvl_config(model)
 
     # Load tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True)
+    tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=trust_remote_code)
 
     # get model arch and config
-    arch, original_config = get_model_arch(model)
+    arch, original_config = get_model_arch(model, trust_remote_code=trust_remote_code)
 
     if model_type == 'llm':
-        model = load_hf_from_pretrained(model, dtype=dtype, trust_remote_code=True)
+        model = load_hf_from_pretrained(model, dtype=dtype, trust_remote_code=trust_remote_code)
         vl_model = None
     elif model_type == 'vlm':
-        vl_model = load_vl_model(model, backend=None, with_llm=True).vl_model
+        vl_model = load_vl_model(model, backend=None, with_llm=True, trust_remote_code=trust_remote_code).vl_model
         model = vl_model
         if hasattr(vl_model, 'language_model'):  # deepseek-vl, ...
             model = vl_model.language_model
@@ -355,7 +350,7 @@ def calibrate(model: str,
         '`neuralmagic_calibration`, `open-platypus`, `openwebtext`.'
 
     arch, vl_model, model, tokenizer, model_type, work_dir = load_model_and_tokenizer(
-        model, dtype=dtype, work_dir=work_dir)
+        model, dtype=dtype, work_dir=work_dir, trust_remote_code=trust_remote_code)
 
     if model_type in ['MixtralForCausalLM']:
         update_moe_mapping(model, model_type)
@@ -401,7 +396,7 @@ def calibrate(model: str,
 
     calib_ctx.export(work_dir)
 
-    return arch, vl_model, model, tokenizer
+    return arch, vl_model, model, tokenizer, work_dir
 
 
 if __name__ == '__main__':

lmdeploy/lite/apis/smooth_quant.py

@@ -8,6 +8,7 @@
 from torch import nn
 
 from lmdeploy.lite.apis.calibrate import LAYER_TYPE_MAP, NORM_TYPE_MAP, calibrate
+from lmdeploy.lite.model import MODELS
 from lmdeploy.lite.quantization.awq import FC_FCS_MAP, NORM_FCS_MAP, awq_layers, skipped_module, smooth_layers
 from lmdeploy.lite.utils import collect_target_modules
 from lmdeploy.pytorch.models import QLinear, QRMSNorm
@@ -26,7 +27,8 @@ def smooth_quant(model: str,
                  device: str = 'cuda',
                  quant_dtype: Literal['int8', 'fp8', 'float8_e4m3fn', 'float8_e5m2'] = 'int8',
                  revision: str = None,
-                 download_dir: str = None):
+                 download_dir: str = None,
+                 trust_remote_code: bool = False):
     try_import_deeplink(device)
     if quant_dtype == 'fp8':
         quant_dtype = 'float8_e4m3fn'
@@ -44,17 +46,18 @@ def smooth_quant(model: str,
         from lmdeploy.utils import get_model
         model = get_model(model, revision=revision, download_dir=download_dir)
     model_path = model
-    vl_model, model, tokenizer, work_dir = calibrate(model,
-                                                     calib_dataset,
-                                                     calib_samples,
-                                                     calib_seqlen,
-                                                     work_dir,
-                                                     device,
-                                                     w_bits=w_bits,
-                                                     w_group_size=-1,
-                                                     search_scale=search_scale,
-                                                     dtype=dtype,
-                                                     batch_size=batch_size)
+    arch, vl_model, model, tokenizer, work_dir = calibrate(model,
+                                                           calib_dataset,
+                                                           calib_samples,
+                                                           calib_seqlen,
+                                                           work_dir,
+                                                           device,
+                                                           w_bits=w_bits,
+                                                           w_group_size=-1,
+                                                           search_scale=search_scale,
+                                                           dtype=dtype,
+                                                           batch_size=batch_size,
+                                                           trust_remote_code=trust_remote_code)
 
     # calibrate function exports the calibration statistics
     # (inputs, outputs, keys and values) to `work_dir`.
@@ -95,32 +98,41 @@ def smooth_quant(model: str,
 
     rmsnorms = collect_target_modules(model, norm_type)
 
-    for name, linear in fcs.items():
-        if skipped_module(name, model.config.architectures[0])[0]:
-            continue
-        linear.to(device)
-        q_linear = QLinear.from_float(linear, quant_dtype=quant_dtype)
-        parent_name, _, child_name = name.rpartition('.')
-        parent = model.get_submodule(parent_name)
-        setattr(parent, child_name, q_linear)
-        linear.to('cpu')
-        q_linear.to('cpu')
-        torch.cuda.empty_cache()
-
-    for name, norm in rmsnorms.items():
-        if skipped_module(name, model.config.architectures[0])[0]:
-            continue
-        norm.to(device)
-        q_norm = QRMSNorm.from_float(norm, quant_dtype=quant_dtype)
-        parent_name, _, child_name = name.rpartition('.')
-        parent = model.get_submodule(parent_name)
-        setattr(parent, child_name, q_norm)
-        norm.to('cpu')
-        q_norm.to('cpu')
-        torch.cuda.empty_cache()
+    patterns = []
+    skipped_modules = []
+    arch = model.config.architectures[0]
+    model_layer = MODELS.get(arch)
+    if model_layer:
+        patterns = model_layer.skipped_modules()
+        skipped_modules.extend(patterns)
+
+    for modules, q_cls in ((fcs, QLinear), (rmsnorms, QRMSNorm)):
+        for name, module in modules.items():
+            skipped, skipped_pattern = skipped_module(name, patterns)
+            if skipped:
+                if skipped_pattern and skipped_pattern not in skipped_modules:
+                    skipped_modules.append(skipped_pattern)
+                continue
+
+            module.to(device)
+            q_module = q_cls.from_float(module, quant_dtype=quant_dtype)
+
+            parent_name, _, child_name = name.rpartition('.')
+            parent = model.get_submodule(parent_name)
+            setattr(parent, child_name, q_module)
+
+            module.to('cpu')
+            q_module.to('cpu')
+            torch.cuda.empty_cache()
+
 
     quant_dtype_s = str(quant_dtype).split('.')[1]
-    model.config.update(dict(quantization_config=dict(quant_method='smooth_quant', quant_dtype=f'{quant_dtype_s}')))
+    quantization_config = dict(quant_method='smooth_quant',
+                               quant_dtype=f'{quant_dtype_s}')
+    if skipped_modules:
+        quantization_config['modules_to_not_convert'] = skipped_modules
+    model.config.update(dict(quantization_config=quantization_config))
+    # model.config.update(dict(quantization_config=dict(quant_method='smooth_quant', quant_dtype=f'{quant_dtype_s}')))
 
     if vl_model:
         from .auto_awq import save_vl_model

lmdeploy/lite/model/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+
+from . import (
+    mixtral,  # noqa: F401
+    qwen,  # noqa: F401
+)
+from .base import MODELS
+
+__all__ = ['MODELS']