add DeekSeek and Kimi-K2;support per-tensor and w4a8;support batch calibration in fp8 low_memory_run

Author: ali-88123

angelslim/compressor/quant/core/__init__.py

@@ -15,9 +15,11 @@
 from .config import *  # noqa: F401 F403
 from .hook import DiTHook, PTQHook  # noqa: F401
 from .metrics import mse_loss, snr_loss  # noqa: F401
-from .packing_utils import dequantize_gemm  # noqa: F401
+from .packing_utils import dequantize_gemm, pack_weight_to_int8  # noqa: F401
 from .quant_func import *  # noqa: F401 F403
 from .sample_func import EMASampler, MultiStepSampler  # noqa: F401
+from .save import DeepseekV3HfPTQSave  # noqa: F401
+from .save import DeepseekV3PTQSaveTRTLLM  # noqa: F401
 from .save import PTQPTMSave  # noqa: F401
 from .save import PTQSaveVllmHF  # noqa: F401
 from .save import PTQTorchSave  # noqa: F401

angelslim/compressor/quant/core/config.py

@@ -74,10 +74,24 @@ def __init__(self, config, global_config=None):
             )
             self.weight_observer = WEIGHT_OBSERVERS_CLASS[weight_quant_method]
             self.kv_cache_observer = None
-            self.quant_algo_info = {
-                "w": f"fp8_{weight_quant_method}",
-                "ignore_layers": quantization_args.ignore_layers,
-            }
+
+            if "w4a8" in self.quant_algo:
+                group_size = (
+                    128
+                    if quantization_args.quant_method["group_size"] == -1
+                    else quantization_args.quant_method["group_size"]
+                )
+                self.quant_algo_info = {
+                    "w": f"int4_{weight_quant_method}",
+                    "w_group_size": group_size,
+                    "ignore_layers": quantization_args.ignore_layers,
+                }
+            else:
+                self.quant_algo_info = {
+                    "w": f"fp8_{weight_quant_method}",
+                    "ignore_layers": quantization_args.ignore_layers,
+                }
+
             if act_quant_method is not None:
                 self.quant_algo_info["a"] = f"fp8_{act_quant_method}-{is_dynamic}"
             self.hidden_size = global_config.hidden_size

angelslim/compressor/quant/core/packing_utils.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import numpy as np
 import torch
 
 AWQ_ORDER = [0, 2, 4, 6, 1, 3, 5, 7]
@@ -115,3 +116,21 @@ def dequantize_gemm(qweight, qzeros, scales, bits, group_size):
     iweight = (iweight - izeros) * scales
 
     return iweight
+
+
+def pack_weight_to_int8(weight):
+    weight = weight.t().contiguous().cpu()
+    weight = weight.to(torch.float32).numpy().astype(np.int8)
+
+    i = 0
+    row = 0
+    packed_weight = np.zeros((weight.shape[0] // 2, weight.shape[1]), dtype=np.int8)
+    while row < packed_weight.shape[0]:
+        for j in range(i, i + (8 // 4)):
+            packed_weight[row] |= (weight[j] & 0x0F) << (4 * (j - i))
+        i += 8 // 4
+        row += 1
+
+    packed_weight = packed_weight.astype(np.int8)
+    packed_weight = torch.from_numpy(packed_weight).t().contiguous()
+    return packed_weight

angelslim/compressor/quant/core/quant_func.py

@@ -15,6 +15,8 @@
 from typing import Tuple
 
 import torch
+import triton
+import triton.language as tl
 
 from .metrics import mse_loss
 
@@ -314,3 +316,67 @@ def tensor_quant_dequant_fp8(x, scale, bits=8, mantissa_bit=3, sign_bits=1):
     )
     quant_dequant_x *= scale
     return quant_dequant_x
+
+
+# This function is copied from DeepSeek-V3 (MIT License):
+# Copyright (c) 2023 DeepSeek-AI
+# Original source: https://github.com/deepseek-ai/DeepSeek-V3
+@triton.jit
+def weight_dequant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr):
+    """
+    Dequantizes weights using the provided scaling factors and stores the result.
+
+    Args:
+        x_ptr (tl.pointer): Pointer to the quantized weights.
+        s_ptr (tl.pointer): Pointer to the scaling factors.
+        y_ptr (tl.pointer): Pointer to the output buffer for dequantized weights.
+        M (int): Number of rows in the weight matrix.
+        N (int): Number of columns in the weight matrix.
+        BLOCK_SIZE (tl.constexpr): Size of the block for tiling.
+
+    Returns:
+        None
+    """
+    pid_m = tl.program_id(axis=0)
+    pid_n = tl.program_id(axis=1)
+    n = tl.cdiv(N, BLOCK_SIZE)
+    offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    offs_n = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    offs = offs_m[:, None] * N + offs_n[None, :]
+    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
+    x = tl.load(x_ptr + offs, mask=mask).to(tl.float32)
+    s = tl.load(s_ptr + pid_m * n + pid_n)
+    y = x * s
+    tl.store(y_ptr + offs, y, mask=mask)
+
+
+# This function is copied from DeepSeek-V3 (MIT License):
+# Copyright (c) 2023 DeepSeek-AI
+# Original source: https://github.com/deepseek-ai/DeepSeek-V3
+def weight_dequant(
+    x: torch.Tensor, s: torch.Tensor, block_size: int = 128
+) -> torch.Tensor:
+    """
+    Dequantizes the given weight tensor using the provided scale tensor.
+
+    Args:
+        x (torch.Tensor): The quantized weight tensor of shape (M, N).
+        s (torch.Tensor): The scale tensor of shape (M, N).
+        block_size (int, optional): The block size to use for dequantization. Defaults to 128. # noqa: E501
+
+    Returns:
+        torch.Tensor: The dequantized weight tensor of the same shape as `x`.
+
+    Raises:
+        AssertionError: If `x` or `s` are not contiguous or if their dimensions are not 2. # noqa: E501
+    """
+    assert x.is_contiguous() and s.is_contiguous(), "Input tensors must be contiguous"
+    assert x.dim() == 2 and s.dim() == 2, "Input tensors must have 2 dimensions"
+    M, N = x.size()
+    y = torch.empty_like(x, dtype=torch.get_default_dtype())
+    grid = lambda meta: (  # noqa: E731
+        triton.cdiv(M, meta["BLOCK_SIZE"]),
+        triton.cdiv(N, meta["BLOCK_SIZE"]),
+    )
+    weight_dequant_kernel[grid](x, s, y, M, N, BLOCK_SIZE=block_size)
+    return y