Added DeepSeek V3 support. (#688)

Author: LagPixelLOL

awq/models/__init__.py

@@ -24,6 +24,7 @@
 from .phi3_v import Phi3VAWQForCausalLM
 from .cohere import CohereAWQForCausalLM
 from .deepseek_v2 import DeepseekV2AWQForCausalLM
+from .deepseek_v3 import DeepseekV3AWQForCausalLM
 from .minicpm import MiniCPMAWQForCausalLM
 from .internlm2 import InternLM2AWQForCausalLM
 from .minicpm3 import MiniCPM3AWQForCausalLM

awq/models/auto.py

@@ -34,6 +34,7 @@
     "phi3_v": Phi3VAWQForCausalLM,
     "cohere": CohereAWQForCausalLM,
     "deepseek_v2": DeepseekV2AWQForCausalLM,
+    "deepseek_v3": DeepseekV3AWQForCausalLM,
     "minicpm": MiniCPMAWQForCausalLM,
     "internlm2": InternLM2AWQForCausalLM,
     "minicpm3": MiniCPM3AWQForCausalLM,

awq/models/base.py

@@ -82,6 +82,7 @@
     "phi3_v": "AutoModelForCausalLM",
     "cohere": "AutoModelForCausalLM",
     "deepseek_v2": "AutoModelForCausalLM",
+    "deepseek_v3": "AutoModelForCausalLM",
     "minicpm": "AutoModelForCausalLM",
     "minicpm3":"AutoModelForCausalLM",
     "internlm2": "AutoModelForCausalLM",

awq/models/deepseek_v3.py

@@ -0,0 +1,128 @@
+import tqdm
+from typing import List, Tuple
+from .base import BaseAWQForCausalLM
+
+
+class DeepseekV3AWQForCausalLM(BaseAWQForCausalLM):
+    layer_type = "DeepseekV3DecoderLayer"
+    max_seq_len_key = "max_position_embeddings"
+
+    @staticmethod
+    def get_model_layers(model):
+        return model.model.layers
+
+    @staticmethod
+    def get_act_for_scaling(module):
+        return dict(is_scalable=False)
+
+    @staticmethod
+    def move_embed(model, device: str):
+        model.model.embed_tokens = model.model.embed_tokens.to(device)
+
+    @staticmethod
+    def get_layers_for_scaling(
+        module, input_feat, module_kwargs
+    ):
+        layers = []
+
+        if hasattr(module.self_attn, "q_proj"):
+            # attention input
+            layers.append(
+                dict(
+                    prev_op=module.input_layernorm,
+                    layers=[
+                        module.self_attn.q_proj,
+                        module.self_attn.kv_a_proj_with_mqa,
+                    ],
+                    inp=input_feat["self_attn.q_proj"],
+                    module2inspect=module.self_attn,
+                    kwargs=module_kwargs,
+                )
+            )
+        else:
+            # attention input
+            layers.append(
+                dict(
+                    prev_op=module.input_layernorm,
+                    layers=[
+                        module.self_attn.q_a_proj,
+                        module.self_attn.kv_a_proj_with_mqa,
+                    ],
+                    inp=input_feat["self_attn.q_a_proj"],
+                    module2inspect=module.self_attn,
+                    kwargs=module_kwargs,
+                )
+            )
+            layers.append(
+                dict(
+                    prev_op=module.self_attn.q_a_layernorm,
+                    layers=[
+                        module.self_attn.q_b_proj,
+                    ],
+                    inp=input_feat["self_attn.q_b_proj"],
+                )
+            )
+
+        # kv layernorm
+        layers.append(
+            dict(
+                prev_op=module.self_attn.kv_a_layernorm,
+                layers=[
+                    module.self_attn.kv_b_proj,
+                ],
+                inp=input_feat["self_attn.kv_b_proj"],
+            )
+        )
+
+        if hasattr(module.mlp, "gate"):
+            # linear in
+            layers.append(
+                dict(
+                    prev_op=module.post_attention_layernorm,
+                    layers=[
+                        w
+                        for expert in module.mlp.experts
+                        for w in [expert.gate_proj, expert.up_proj]
+                    ] + [module.mlp.shared_experts.gate_proj, module.mlp.shared_experts.up_proj],
+                    inp=input_feat["mlp"],
+                    module2inspect=module.mlp,
+                )
+            )
+
+            # linear out
+            for i, expert in enumerate(module.mlp.experts):
+                layers.append(
+                    dict(
+                        prev_op=expert.up_proj,
+                        layers=[expert.down_proj],
+                        inp=input_feat[f"mlp.experts.{i}.down_proj"],
+                    )
+                )
+            layers.append(
+                dict(
+                    prev_op=module.mlp.shared_experts.up_proj,
+                    layers=[module.mlp.shared_experts.down_proj],
+                    inp=input_feat[f"mlp.shared_experts.down_proj"],
+                )
+            )
+        else:
+            # linear 1
+            layers.append(
+                dict(
+                    prev_op=module.post_attention_layernorm,
+                    layers=[module.mlp.gate_proj, module.mlp.up_proj],
+                    inp=input_feat["mlp.gate_proj"],
+                    module2inspect=module.mlp,
+                )
+            )
+
+            # linear 2
+            layers.append(
+                dict(
+                    prev_op=module.mlp.up_proj,
+                    layers=[module.mlp.down_proj],
+                    inp=input_feat["mlp.down_proj"],
+                )
+            )
+
+        return layers

awq/quantize/quantizer.py

@@ -73,7 +73,7 @@ def __init__(
     def pseudo_quantize_tensor(self, w: torch.Tensor):
         org_w_shape = w.shape
         if self.group_size > 0:
-            assert org_w_shape[-1] % self.group_size == 0
+            assert org_w_shape[-1] % self.group_size == 0, f"org_w_shape ({org_w_shape[-1]}) must be a multiple of group_size ({self.group_size})!"
             w = w.reshape(-1, self.group_size)
         assert w.dim() == 2
         assert torch.isnan(w).sum() == 0
@@ -338,6 +338,7 @@ def _search_best_scale(
         with torch.no_grad():
             module_kwargs = self._sanitize_kwargs(kwargs, module2inspect)
             fp16_output = self._module_forward(inp, module2inspect, module_kwargs)
+            fp16_output = fp16_output.clip(torch.finfo(fp16_output.dtype).min, torch.finfo(fp16_output.dtype).max)
 
         # [STEP 4]: Compute loss
         best_scales = self._compute_best_scale(
@@ -406,6 +407,7 @@ def _compute_best_scale(
 
             # W * X
             int_w_output = self._module_forward(x, module2inspect, kwargs)
+            int_w_output = int_w_output.clip(torch.finfo(int_w_output.dtype).min, torch.finfo(int_w_output.dtype).max)
 
             # compute mean squared error (L2 norm)
             loss = self._compute_loss(fp16_output, int_w_output, device)
@@ -623,7 +625,7 @@ def cache_input_hook(m, x, y, name, feat_dict):
                 "block_sparse_moe": layer.block_sparse_moe,
             }
 
-        if self.awq_model.model_type == "deepseek_v2":
+        if self.awq_model.model_type == "deepseek_v2" or self.awq_model.model_type == "deepseek_v3":
             named_linears = {
                 **named_linears,
                 "mlp": layer.mlp,

awq/utils/utils.py

@@ -72,8 +72,8 @@ def set_module_name(model, name, value):
 def clear_memory(weight=None):
     if weight is not None:
         del weight
-    gc.collect()
-    torch.cuda.empty_cache()
+    # gc.collect()
+    # torch.cuda.empty_cache()
 
 
 def compute_memory_used_pct(device):