diff --git a/gptqmodel/models/__init__.py b/gptqmodel/models/__init__.py index 3d559f56f..32601ab22 100644 --- a/gptqmodel/models/__init__.py +++ b/gptqmodel/models/__init__.py @@ -33,3 +33,4 @@ from .starcoder2 import Starcoder2GPTQ from .xverse import XverseGPTQ from .yi import YiGPTQ +from .internlm2 import InternLM2GPTQ diff --git a/gptqmodel/models/_const.py b/gptqmodel/models/_const.py index 9e0f681b5..1725b4865 100644 --- a/gptqmodel/models/_const.py +++ b/gptqmodel/models/_const.py @@ -31,6 +31,7 @@ def get_device_by_type(type_value: str): "RefinedWeb", "baichuan", "internlm", + "internlm2", "qwen", "xverse", "deci", diff --git a/gptqmodel/models/auto.py b/gptqmodel/models/auto.py index 3850196b3..6c4dbca46 100644 --- a/gptqmodel/models/auto.py +++ b/gptqmodel/models/auto.py @@ -37,6 +37,7 @@ from .starcoder2 import Starcoder2GPTQ from .xverse import XverseGPTQ from .yi import YiGPTQ +from .internlm2 import InternLM2GPTQ MODEL_MAP = { "bloom": BloomGPTQ, @@ -55,6 +56,7 @@ "falcon": RWGPTQ, "baichuan": BaiChuanGPTQ, "internlm": InternLMGPTQ, + "internlm2": InternLM2GPTQ, "qwen": QwenGPTQ, "mistral": MistralGPTQ, "Yi": YiGPTQ, diff --git a/gptqmodel/models/internlm2.py b/gptqmodel/models/internlm2.py new file mode 100644 index 000000000..ffb46127a --- /dev/null +++ b/gptqmodel/models/internlm2.py @@ -0,0 +1,14 @@ +from .base import BaseGPTQModel + + +class InternLM2GPTQ(BaseGPTQModel): + base_modules = ["model.tok_embeddings", "model.norm"] + + layers_node = "model.layers" + layer_type = "InternLM2DecoderLayer" + layer_modules = [ + ["attention.wqkv", "attention.wo"], + + ["feed_forward.w1", "feed_forward.w3"], + ["feed_forward.w2"], + ] \ No newline at end of file