ADD GLM model support

auto_gptq/modeling/__init__.py

@@ -6,6 +6,7 @@
 from .cohere import CohereGPTQForCausalLM
 from .decilm import DeciLMGPTQForCausalLM
 from .gemma import GemmaGPTQForCausalLM
+from .chatglm import ChatGLMForCausalLM
 from .gpt2 import GPT2GPTQForCausalLM
 from .gpt_bigcode import GPTBigCodeGPTQForCausalLM
 from .gpt_neox import GPTNeoXGPTQForCausalLM

auto_gptq/modeling/_const.py

@@ -15,6 +15,7 @@
     "moss",
     "gpt_bigcode",
     "codegen",
+    "chatglm",
     "RefinedWebModel",
     "RefinedWeb",
     "baichuan",

auto_gptq/modeling/auto.py

@@ -9,6 +9,7 @@
 from .cohere import CohereGPTQForCausalLM
 from .decilm import DeciLMGPTQForCausalLM
 from .gemma import GemmaGPTQForCausalLM
+from .chatglm import ChatGLMForCausalLM
 from .gpt2 import GPT2GPTQForCausalLM
 from .gpt_bigcode import GPTBigCodeGPTQForCausalLM
 from .gpt_neox import GPTNeoXGPTQForCausalLM
@@ -39,6 +40,7 @@
     "llama": LlamaGPTQForCausalLM,
     "opt": OPTGPTQForCausalLM,
     "moss": MOSSGPTQForCausalLM,
+    "chatglm": ChatGLMForCausalLM,
     "gpt_bigcode": GPTBigCodeGPTQForCausalLM,
     "codegen": CodeGenGPTQForCausalLM,
     "cohere": CohereGPTQForCausalLM,

auto_gptq/modeling/chatglm.py

@@ -0,0 +1,15 @@
+from ._base import BaseGPTQForCausalLM
+
+
+class ChatGLMForCausalLM(BaseGPTQForCausalLM):
+    layer_type = "GLMBlock"
+    layers_block_name = "transformer.encoder.layers"
+    outside_layer_modules = ["transformer.embedding.word_embeddings", "transformer.output_layer"]
+    inside_layer_modules = [
+        ["self_attention.query_key_value"],
+        ["self_attention.dense"],
+        ["mlp.dense_h_to_4h"],
+        ["mlp.dense_4h_to_h"],
+    ]
+
+__all__ = ["ChatGLMForCausalLM"]