add granite support

Author: LRL-ModelCloud

gptqmodel/models/__init__.py

@@ -16,6 +16,7 @@
 from .gpt_bigcode import GPTBigCodeGPTQ
 from .gpt_neox import GPTNeoXGPTQ
 from .gptj import GPTJGPTQ
+from .granite import GraniteGPTQ
 from .grinmoe import GrinMOEGPTQ
 from .internlm import InternLMGPTQ
 from .internlm2 import InternLM2GPTQ
@@ -24,6 +25,7 @@
 from .minicpm3 import MiniCPM3GPTQ
 from .mistral import MistralGPTQ
 from .mixtral import MixtralGPTQ
+from .mllama import MLlamaGPTQ
 from .moss import MOSSGPTQ
 from .mpt import MPTGPTQ
 from .opt import OPTGPTQ
@@ -37,4 +39,3 @@
 from .starcoder2 import Starcoder2GPTQ
 from .xverse import XverseGPTQ
 from .yi import YiGPTQ
-from .mllama import MLlamaGPTQ

gptqmodel/models/_const.py

@@ -58,6 +58,7 @@ def get_device_by_type(type_value: str):
     "exaone",
     "grinmoe",
     "mllama",
+    "granite",
 ]
 
 EXLLAMA_DEFAULT_MAX_INPUT_LENGTH = 2048

gptqmodel/models/auto.py

@@ -21,6 +21,7 @@
 from .gpt_bigcode import GPTBigCodeGPTQ
 from .gpt_neox import GPTNeoXGPTQ
 from .gptj import GPTJGPTQ
+from .granite import GraniteGPTQ
 from .grinmoe import GrinMOEGPTQ
 from .internlm import InternLMGPTQ
 from .internlm2 import InternLM2GPTQ
@@ -87,6 +88,7 @@
     "exaone": ExaoneGPTQ,
     "grinmoe": GrinMOEGPTQ,
     "mllama": MLlamaGPTQ,
+    "granite": GraniteGPTQ,
 }
 
 

gptqmodel/models/mllama.py

@@ -1,6 +1,8 @@
-from .base import BaseGPTQModel
 from transformers import AutoModelForPreTraining
 
+from .base import BaseGPTQModel
+
+
 # TODO FIXME: we currently do not support quantizing cross attention layer (pixel_values)
 class MLlamaGPTQ(BaseGPTQModel):
     # AutoModelForPreTraining return a correct MLlamaForConditionalGeneration for mllama.

gptqmodel/nn_modules/qlinear/__init__.py

@@ -1,4 +1,4 @@
-from typing import Tuple, Optional
+from typing import Optional, Tuple
 
 import torch.nn as nn
 

gptqmodel/nn_modules/qlinear/qlinear_qbits.py

@@ -23,7 +23,7 @@
 def qbits_dtype() -> torch.dtype:
     try:
         from intel_extension_for_transformers import qbits
-    except Exception as e:
+    except Exception:
         raise ImportError("intel_extension_for_transformers not installed. "
                           "Please install via via 'pip install intel_extension_for_transformers")
 
@@ -112,7 +112,7 @@ def post_init(self, quantize_config):
 
         try:
             from intel_extension_for_transformers import qbits
-        except Exception as e:
+        except Exception:
             raise ImportError("intel_extension_for_transformers not installed. "
                               "Please install via via 'pip install intel_extension_for_transformers")
 
@@ -257,7 +257,7 @@ def pack(self, linear, scales, zeros, g_idx=None):
     def forward(self, x: torch.Tensor):
         try:
             from intel_extension_for_transformers import qbits
-        except Exception as e:
+        except Exception:
             raise ImportError("intel_extension_for_transformers not installed. "
                               "Please install via via 'pip install intel_extension_for_transformers")
 

setup.py

@@ -110,9 +110,9 @@
     extra_compile_args = {
         "cxx": [
             "-O3",
-            "-std=c++17", 
-            "-fopenmp", 
-            "-lgomp", 
+            "-std=c++17",
+            "-fopenmp",
+            "-lgomp",
             "-DENABLE_BF16"
             "-Wno-switch-bool",
         ],

tests/test_sharded.py

@@ -11,7 +11,7 @@
 import unittest  # noqa: E402
 
 import torch  # noqa: E402
-from gptqmodel import BACKEND, GPTQModel  # noqa: E402
+from gptqmodel import GPTQModel  # noqa: E402
 from transformers import AutoTokenizer  # noqa: E402