Add support for MPT (#73)

* Add initial support for MPT

* Ruff and stuff

auto_gptq/modeling/__init__.py

@@ -15,6 +15,7 @@
 from .mistral import MistralGPTQForCausalLM
 from .mixtral import MixtralGPTQForCausalLM
 from .moss import MOSSGPTQForCausalLM
+from .mpt import MPTGPTQForCausalLM
 from .opt import OPTGPTQForCausalLM
 from .phi import PhiGPTQForCausalLM
 from .qwen import QwenGPTQForCausalLM

auto_gptq/modeling/_const.py

@@ -23,6 +23,7 @@
     "xverse",
     "deci",
     "stablelm_epoch",
+    "mpt",
 ]
 if compare_transformers_version("v4.28.0", op="ge"):
     SUPPORTED_MODELS.append("llama")

auto_gptq/modeling/auto.py

@@ -18,6 +18,7 @@
 from .mistral import MistralGPTQForCausalLM
 from .mixtral import MixtralGPTQForCausalLM
 from .moss import MOSSGPTQForCausalLM
+from .mpt import MPTGPTQForCausalLM
 from .opt import OPTGPTQForCausalLM
 from .phi import PhiGPTQForCausalLM
 from .qwen import QwenGPTQForCausalLM
@@ -56,6 +57,7 @@
     "longllama": LongLlamaGPTQForCausalLM,
     "gemma": GemmaGPTQForCausalLM,
     "phi": PhiGPTQForCausalLM,
+    "mpt": MPTGPTQForCausalLM,
 }
 
 

auto_gptq/modeling/mpt.py

@@ -0,0 +1,19 @@
+from auto_gptq.modeling import BaseGPTQForCausalLM
+
+
+class MPTGPTQForCausalLM(BaseGPTQForCausalLM):
+    layer_type = "MPTBlock"
+    layers_block_name = "transformer.blocks"
+    outside_layer_modules = [
+        "transformer.wte",  "transformer.norm_f"
+    ]
+
+    inside_layer_modules = [
+        ["attn.Wqkv"],
+        ["attn.out_proj"],
+        ["ffn.up_proj"],
+        ["ffn.down_proj"]
+    ]
+
+
+__all__ = ["MPTGPTQForCausalLM"]