[MODEL] add OVIS support

gptqmodel/models/_const.py

@@ -63,6 +63,7 @@ def get_device_by_type(type_value: str):
     "mobilellm",
     "hymba",
     "olmo2",
+    "ovis",
 ]
 
 EXLLAMA_DEFAULT_MAX_INPUT_LENGTH = 2048

gptqmodel/models/auto.py

@@ -56,6 +56,7 @@
 from .definitions.yi import YiGPTQ
 from .definitions.hymba import  HymbaGPTQ
 from .definitions.olmo2 import Olmo2GPTQ
+from .definitions.ovis import OvisGPTQ
 
 logger = setup_logger()
 
@@ -106,6 +107,7 @@
     "mobilellm": MobileLLMGPTQ,
     "hymba": HymbaGPTQ,
     "olmo2": Olmo2GPTQ,
+    "ovis": OvisGPTQ,
 }
 
 

gptqmodel/models/base.py

@@ -251,10 +251,7 @@ def quantize(
         for row in calibration_dataset:
             input_ids = row["input_ids"]
             if isinstance(input_ids, torch.Tensor):
-                if input_ids.dim() == 1:
-                    input_ids_length = input_ids.shape[0]
-                else:
-                    raise ValueError("Expected a 1-dimensional tensor for 'input_ids', but got a tensor with {0} dimensions.".format(input_ids.dim()))
+                input_ids_length = input_ids.numel()
             else:
                 input_ids_length = len(input_ids)
 
@@ -428,11 +425,20 @@ def store_input_hook(_, args, kwargs):
         handle = layers[0].register_forward_pre_hook(store_input_hook, with_kwargs=True)
         for example in calibration_dataset:
             for k, v in example.items():
-                if len(v.shape) == 1:
-                    v = v.unsqueeze(0)
-                example[k] = move_to(v, cur_layer_device)
+                if isinstance(v, list):
+                    for i in range(len(v)):
+                        if len(v[i].shape) == 1:
+                            v[i] = v[i].unsqueeze(0)
+                        v[i] = move_to(v[i], cur_layer_device)
+                else:
+                    if len(v.shape) == 1:
+                        v = v.unsqueeze(0)
+                    example[k] = move_to(v, cur_layer_device)
             try:
-                self.model(**example)
+                if self.__class__.__name__ == "OvisGPTQ":
+                    self.generate(**example)
+                else:
+                    self.model(**example)
             except ValueError:
                 pass
         handle.remove()

gptqmodel/models/definitions/__init__.py

@@ -40,4 +40,5 @@
 from .xverse import XverseGPTQ
 from .yi import YiGPTQ
 from .hymba import HymbaGPTQ
-from .olmo2 import Olmo2GPTQ
+from .olmo2 import Olmo2GPTQ
+from .ovis import OvisGPTQ

gptqmodel/models/definitions/ovis.py

@@ -0,0 +1,27 @@
+from ..base import BaseGPTQModel
+import torch
+
+class OvisGPTQ(BaseGPTQModel):
+    base_modules = ["llm.model.embed_tokens", "llm.model.norm", "visual_tokenizer", "vte"]
+
+    layers_node = "llm.model.layers"
+    layer_type = ["LlamaDecoderLayer", "Gemma2DecoderLayer"]
+    layer_modules = [
+        ["self_attn.k_proj", "self_attn.v_proj", "self_attn.q_proj"],
+        ["self_attn.o_proj"],
+        ["mlp.up_proj", "mlp.gate_proj"],
+        ["mlp.down_proj"],
+    ]
+
+    # hack so one can prepare examples outside
+    def _prepare_dataset_for_quantization(
+            self,
+            calibration_dataset,
+            batch_size: int = 1,
+            tokenizer=None, ):
+        return calibration_dataset
+
+    def generate(self, **kwargs):
+        """shortcut for model.generate"""
+        with torch.inference_mode(), torch.amp.autocast(device_type=self.device.type):
+            return self.model.generate(inputs=kwargs.pop("input_ids"), **kwargs)

gptqmodel/models/loader.py

@@ -103,7 +103,7 @@ def skip(*args, **kwargs):
         model = cls.loader.from_pretrained(pretrained_model_id_or_path, **model_init_kwargs)
 
         model_config = model.config.to_dict()
-        seq_len_keys = ["max_position_embeddings", "seq_length", "n_positions"]
+        seq_len_keys = ["max_position_embeddings", "seq_length", "n_positions", "multimodal_max_length"]
         if any(k in model_config for k in seq_len_keys):
             for key in seq_len_keys:
                 if key in model_config:
@@ -502,7 +502,7 @@ def skip(*args, **kwargs):
 
         # == step4: set seqlen == #
         model_config = model.config.to_dict()
-        seq_len_keys = ["max_position_embeddings", "seq_length", "n_positions"]
+        seq_len_keys = ["max_position_embeddings", "seq_length", "n_positions", "multimodal_max_length"]
         if any(k in model_config for k in seq_len_keys):
             for key in seq_len_keys:
                 if key in model_config:

tests/models/model_test.py

@@ -3,6 +3,7 @@
 import os
 
 from gptqmodel.utils.lm_eval import lm_eval
+from ovis_calibration_dataset import get_calib_dataset
 
 os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
 # -- end do not touch
@@ -93,8 +94,8 @@ def quantModel(self, model_id_or_path, trust_remote_code=False, torch_dtype="aut
         )
 
         tokenizer = self.load_tokenizer(model_id_or_path, trust_remote_code=trust_remote_code)
-
-        calibration_dataset = self.load_dataset(tokenizer)
+        is_ovis_model = "Ovis" in model_id_or_path
+        calibration_dataset = self.load_dataset(tokenizer) if not is_ovis_model else get_calib_dataset(model)
 
         # mpt model need
         if not model.config.pad_token_id:

tests/models/test_ovis_1_6_llama.py

@@ -0,0 +1,13 @@
+from model_test import ModelTest
+
+class TestOvis1_6_Llama(ModelTest):
+    NATIVE_MODEL_ID = "/monster/data/model/Ovis1.6-Llama3.2-3B"
+    NATIVE_ARC_CHALLENGE_ACC = 0.2739
+    NATIVE_ARC_CHALLENGE_ACC_NORM = 0.3055
+
+    TRUST_REMOTE_CODE = True
+    APPLY_CHAT_TEMPLATE = False
+    BATCH_SIZE = 1
+
+    def test_ovis_1_6(self):
+        self.quant_lm_eval()