fix vlm bugs

llmc/__main__.py

@@ -34,6 +34,7 @@ def main(config):
     eval_model(model, None, eval_list, eval_pos='pretrain')
 
     for modality in config.quant.get('quant_objects', ['language']):
+        model.get_key_info(modality)
         if not config.get('calib', False):
             blockwise_opt = ALGO_REGISTRY[config.quant.method](
                 model,

llmc/models/base_model.py

@@ -37,11 +37,14 @@ def __init__(self, config, device_map=None, use_cache=False):
         self.build_tokenizer()
         self.build_model()
         self.model.eval()
-        self.find_blocks()
+        self.kvcache_buffer = []
+        self.get_key_info(modality='language')
+
+    def get_key_info(self, modality='language'):
+        self.find_blocks(modality=modality)
         self.find_embed_layers()
         self.find_block_name()
-        self.add_layernorms_class()
-        self.kvcache_buffer = []
+        self.add_layernorms_class(modality=modality)
 
     def reset_kv(self):
         for kvcache in self.kvcache_buffer:
@@ -190,10 +193,10 @@ def build_model(self):
         )
         logger.info(f'self.model : {self.model}')
 
-    def add_layernorms_class(self):
+    def add_layernorms_class(self, modality='language'):
         ln_class_list = []
-        single_block = self.get_blocks()[0]
-        ln_dict = self.get_layernorms_in_block(single_block)
+        single_block = self.blocks[0]
+        ln_dict = self.get_layernorms_in_block(single_block, modality=modality)
         for ln_name in ln_dict:
             ln_class = ln_dict[ln_name].__class__
             if ln_class not in ln_class_list:
@@ -207,7 +210,6 @@ def add_layernorms_class(self):
     def collect_first_block_input(self, calib_data, padding_mask=None, modality='language'):
         first_block_input = defaultdict(list)
 
-        self.find_blocks(modality)
         Catcher = self.get_catcher(first_block_input)
 
         if not self.use_cpu_to_save_cuda_mem_for_catcher:

llmc/models/bloom.py

@@ -37,7 +37,7 @@ def skip_layer_name(self):
     def has_bias(self):
         return True
 
-    def get_layernorms_in_block(self, block):
+    def get_layernorms_in_block(self, block, modality='language'):
         return {
             'input_layernorm': block.input_layernorm,
             'post_attention_layernorm': block.post_attention_layernorm,

llmc/models/chatglm.py

@@ -43,7 +43,7 @@ def skip_layer_name(self):
     def has_bias(self):
         return False
 
-    def get_layernorms_in_block(self, block):
+    def get_layernorms_in_block(self, block, modality='language'):
         return {
             'input_layernorm': block.input_layernorm,
             'post_attention_layernorm': block.post_attention_layernorm,

llmc/models/deepseekv2.py

@@ -34,7 +34,7 @@ def skip_layer_name(self):
     def has_bias(self):
         return False
 
-    def get_layernorms_in_block(self, block):
+    def get_layernorms_in_block(self, block, modality='language'):
         return {
             'input_layernorm': block.input_layernorm,
             'post_attention_layernorm': block.post_attention_layernorm,

llmc/models/falcon.py

@@ -30,7 +30,7 @@ def get_layers_except_blocks(self):
     def has_bias(self):
         return False
 
-    def get_layernorms_in_block(self, block):
+    def get_layernorms_in_block(self, block, modality='language'):
         if block.config.architectures[0] == 'RWForCausalLM':
             new_decoder_architecture = False
         elif block.config.architectures[0] == 'FalconForCausalLM':

llmc/models/gemma2.py

@@ -54,7 +54,7 @@ def skip_layer_name(self):
     def has_bias(self):
         return False
 
-    def get_layernorms_in_block(self, block):
+    def get_layernorms_in_block(self, block, modality='language'):
         return {
             'input_layernorm': block.input_layernorm,
             'pre_feedforward_layernorm': block.pre_feedforward_layernorm,

llmc/models/internlm2.py

@@ -43,7 +43,7 @@ def skip_layer_name(self):
     def has_bias(self):
         return False
 
-    def get_layernorms_in_block(self, block):
+    def get_layernorms_in_block(self, block, modality='language'):
         return {
             'attention_norm': block.attention_norm,
             'ffn_norm': block.ffn_norm,

llmc/models/internvl2.py

@@ -184,6 +184,7 @@ def __init__(self, config, device_map=None, use_cache=False):
 
 class InternVL2SharedBehavior():
     def build_model(self):
+        self.eval_name = 'InternVL2'
         self.vlm_model_config = AutoConfig.from_pretrained(
             self.model_path, trust_remote_code=True
         )
@@ -311,6 +312,18 @@ def find_blocks(self, modality='language'):
         elif modality == 'vision':
             self.blocks = self.vision_model.encoder.layers
 
+    def get_layernorms_in_block(self, block, modality='language'):
+        if modality == 'language':
+            return {
+                'attention_norm': block.attention_norm,
+                'ffn_norm': block.ffn_norm,
+            }
+        elif modality == 'vision':
+            return {
+                'norm1': block.norm1,
+                'norm2': block.norm2,
+            }
+
     def get_vision_subsets_in_block(self, block):
         return [
             {

llmc/models/llama.py

@@ -43,7 +43,7 @@ def skip_layer_name(self):
     def has_bias(self):
         return False
 
-    def get_layernorms_in_block(self, block):
+    def get_layernorms_in_block(self, block, modality='language'):
         return {
             'input_layernorm': block.input_layernorm,
             'post_attention_layernorm': block.post_attention_layernorm,

llmc/models/llava.py

@@ -104,6 +104,18 @@ def find_blocks(self, modality='language'):
         elif modality == 'vision':
             self.blocks = self.vision_model.vision_model.encoder.layers
 
+    def get_layernorms_in_block(self, block, modality='language'):
+        if modality == 'language':
+            return {
+                'input_layernorm': block.input_layernorm,
+                'post_attention_layernorm': block.post_attention_layernorm,
+            }
+        elif modality == 'vision':
+            return {
+                'layer_norm1': block.layer_norm1,
+                'layer_norm2': block.layer_norm2,
+            }
+
     def get_vision_subsets_in_block(self, block):
         return [
             {

llmc/models/minicpm.py

@@ -39,7 +39,7 @@ def skip_layer_name(self):
     def has_bias(self):
         return False
 
-    def get_layernorms_in_block(self, block):
+    def get_layernorms_in_block(self, block, modality='language'):
         return {
             'input_layernorm': block.input_layernorm,
             'post_attention_layernorm': block.post_attention_layernorm,

llmc/models/mistral.py

@@ -36,7 +36,7 @@ def skip_layer_name(self):
     def has_bias(self):
         return False
 
-    def get_layernorms_in_block(self, block):
+    def get_layernorms_in_block(self, block, modality='language'):
         return {
             'input_layernorm': block.input_layernorm,
             'post_attention_layernorm': block.post_attention_layernorm,

llmc/models/mixtral.py

@@ -29,7 +29,7 @@ def skip_layer_name(self):
     def has_bias(self):
         return False
 
-    def get_layernorms_in_block(self, block):
+    def get_layernorms_in_block(self, block, modality='language'):
         return {
             'input_layernorm': block.input_layernorm,
             'post_attention_layernorm': block.post_attention_layernorm,

llmc/models/mllama.py

@@ -99,7 +99,7 @@ def single_process(self, img_qas):
         ).to(next(self.vlm_model.parameters()).dtype)
         return inputs
 
-    def get_layernorms_in_block(self, block):
+    def get_layernorms_in_block(self, block, modality='language'):
         return {
             'input_layernorm': block.input_layernorm,
             'post_attention_layernorm': block.post_attention_layernorm,

llmc/models/opt.py

@@ -44,7 +44,7 @@ def skip_layer_name(self):
     def has_bias(self):
         return True
 
-    def get_layernorms_in_block(self, block):
+    def get_layernorms_in_block(self, block, modality='language'):
         return {
             'self_attn_layer_norm': block.self_attn_layer_norm,
             'final_layer_norm': block.final_layer_norm,

llmc/models/phi.py

@@ -40,7 +40,7 @@ def skip_layer_name(self):
     def has_bias(self):
         return False
 
-    def get_layernorms_in_block(self, block):
+    def get_layernorms_in_block(self, block, modality='language'):
         return {
             'input_layernorm': block.input_layernorm
         }

llmc/models/phi3.py

@@ -35,7 +35,7 @@ def skip_layer_name(self):
     def has_bias(self):
         return False
 
-    def get_layernorms_in_block(self, block):
+    def get_layernorms_in_block(self, block, modality='language'):
         return {
             'input_layernorm': block.input_layernorm,
             'post_attention_layernorm': block.post_attention_layernorm,

llmc/models/qwen.py

@@ -39,7 +39,7 @@ def skip_layer_name(self):
     def has_bias(self):
         return False
 
-    def get_layernorms_in_block(self, block):
+    def get_layernorms_in_block(self, block, modality='language'):
         return {
             'ln_1': block.ln_1,
             'ln_2': block.ln_2,

llmc/models/qwen2.py

@@ -51,7 +51,7 @@ def skip_layer_name(self):
     def has_bias(self):
         return False
 
-    def get_layernorms_in_block(self, block):
+    def get_layernorms_in_block(self, block, modality='language'):
         return {
             'input_layernorm': block.input_layernorm,
             'post_attention_layernorm': block.post_attention_layernorm,

llmc/models/qwen2moe.py

@@ -42,7 +42,7 @@ def skip_layer_name(self):
     def has_bias(self):
         return False
 
-    def get_layernorms_in_block(self, block):
+    def get_layernorms_in_block(self, block, modality='language'):
         return {
             'input_layernorm': block.input_layernorm,
             'post_attention_layernorm': block.post_attention_layernorm,

llmc/models/smollm.py

@@ -36,7 +36,7 @@ def skip_layer_name(self):
     def has_bias(self):
         return False
 
-    def get_layernorms_in_block(self, block):
+    def get_layernorms_in_block(self, block, modality='language'):
         return {
             'input_layernorm': block.input_layernorm,
             'post_attention_layernorm': block.post_attention_layernorm,

llmc/models/stablelm.py

@@ -40,7 +40,7 @@ def skip_layer_name(self):
     def has_bias(self):
         return False
 
-    def get_layernorms_in_block(self, block):
+    def get_layernorms_in_block(self, block, modality='language'):
         return {
             'input_layernorm': block.input_layernorm,
             'post_attention_layernorm': block.post_attention_layernorm,

llmc/models/starcoder.py

@@ -35,7 +35,7 @@ def skip_layer_name(self):
     def has_bias(self):
         return True
 
-    def get_layernorms_in_block(self, block):
+    def get_layernorms_in_block(self, block, modality='language'):
         return {
             'ln_1': block.ln_1,
             'ln_2': block.ln_2,

llmc/models/vit.py

@@ -48,7 +48,7 @@ def skip_layer_name(self):
     def has_bias(self):
         return False
 
-    def get_layernorms_in_block(self, block):
+    def get_layernorms_in_block(self, block, modality='vision'):
         return {
             'layernorm_before': block.layernorm_before,
             'layernorm_after': block.layernorm_after,