Cohere Command-R Added (#631)

Co-authored-by: Egor Tolmachev <t333ga@gmail.com>

auto_gptq/modeling/__init__.py

@@ -3,6 +3,7 @@
 from .baichuan import BaiChuanGPTQForCausalLM
 from .bloom import BloomGPTQForCausalLM
 from .codegen import CodeGenGPTQForCausalLM
+from .cohere import CohereGPTQForCausalLM
 from .decilm import DeciLMGPTQForCausalLM
 from .gemma import GemmaGPTQForCausalLM
 from .gpt2 import GPT2GPTQForCausalLM

auto_gptq/modeling/_const.py

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

auto_gptq/modeling/auto.py

@@ -6,6 +6,7 @@
 from .baichuan import BaiChuanGPTQForCausalLM
 from .bloom import BloomGPTQForCausalLM
 from .codegen import CodeGenGPTQForCausalLM
+from .cohere import CohereGPTQForCausalLM
 from .decilm import DeciLMGPTQForCausalLM
 from .gemma import GemmaGPTQForCausalLM
 from .gpt2 import GPT2GPTQForCausalLM
@@ -40,6 +41,7 @@
     "moss": MOSSGPTQForCausalLM,
     "gpt_bigcode": GPTBigCodeGPTQForCausalLM,
     "codegen": CodeGenGPTQForCausalLM,
+    "cohere": CohereGPTQForCausalLM,
     "RefinedWebModel": RWGPTQForCausalLM,
     "RefinedWeb": RWGPTQForCausalLM,
     "falcon": RWGPTQForCausalLM,

auto_gptq/modeling/cohere.py

@@ -0,0 +1,20 @@
+from logging import getLogger
+
+from ._base import BaseGPTQForCausalLM
+
+
+logger = getLogger(__name__)
+
+class CohereGPTQForCausalLM(BaseGPTQForCausalLM):
+    layer_type = "CohereDecoderLayer"
+    layers_block_name = "model.layers"
+    outside_layer_modules = ["model.embed_tokens", "model.norm"]
+    inside_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"],
+    ]
+
+__all__ = ["CohereGPTQForCausalLM"]
+