Add support for GPTBigCode

auto_gptq/modeling/__init__.py

@@ -7,3 +7,4 @@ from .gptj import *
 from .llama import *
 from .moss import *
 from .opt import *
+from .gpt_bigcode import *

auto_gptq/modeling/_const.py

@@ -6,7 +6,7 @@ from transformers import __version__ as transformers_version
 CPU = device("cpu")
 CUDA_0 = device("cuda:0")
 
-SUPPORTED_MODELS = ["bloom", "gptj", "gpt2", "gpt_neox", "opt", "moss"]
+SUPPORTED_MODELS = ["bloom", "gptj", "gpt2", "gpt_neox", "opt", "moss", "gpt_bigcode"]
 if parse_version(transformers_version) >= parse_version("v4.28.0"):
     SUPPORTED_MODELS.append("llama")
 

auto_gptq/modeling/auto.py

@@ -9,6 +9,7 @@ from .gpt2 import GPT2GPTQForCausalLM
 from .llama import LlamaGPTQForCausalLM
 from .moss import MOSSGPTQForCausalLM
 from .opt import OPTGPTQForCausalLM
+from .gpt_bigcode import GPTBigCodeGPTQForCausalLM
 
 
 GPTQ_CAUSAL_LM_MODEL_MAP = {
@@ -18,7 +19,8 @@ GPTQ_CAUSAL_LM_MODEL_MAP = {
     "gpt2": GPT2GPTQForCausalLM,
     "llama": LlamaGPTQForCausalLM,
     "opt": OPTGPTQForCausalLM,
-    "moss": MOSSGPTQForCausalLM
+    "moss": MOSSGPTQForCausalLM,
+    "gpt_bigcode": GPTBigCodeGPTQForCausalLM
 }
 
 

auto_gptq/modeling/gpt_bigcode.py

@@ -0,0 +1,17 @@
+from auto_gptq.modeling import BaseGPTQForCausalLM
+
+
+class GPTBigCodeGPTQForCausalLM(BaseGPTQForCausalLM):
+    layer_type = "GPTBigCodeBlock"
+    layers_block_name = "transformer.h"
+    outside_layer_modules = [
+        "transformer.wpe", "transformer.wte", "transformer.ln_f"
+    ]
+    inside_layer_modules = [
+        ["attn.c_attn"],
+        ["attn.c_proj"],
+        ["mlp.c_fc"],
+        ["mlp.c_proj"]
+    ]
+
+__all__ = ["GPTBigCodeGPTQForCausalLM"]