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

README.md

@@ -31,12 +31,11 @@
 - [技术交流](#技术交流)
 
 ## 📣最新进展
-- [25/08/04] 我们支持了`Hunyuan 0.5B/1.8B/4B/7B`模型和多模态模型`Qwen2.5VL 3B/7B/32B/72B`的FP8、INT4量化。我们还开源了`Hunyuan 1.8B/4B/7B`系列模型的Eagle3权重。
+- [25/08/06] 我们支持了`Hunyuan 0.5B/1.8B/4B/7B`和`Qwen2.5VL 3B/7B/32B/72B`的FP8、INT4量化，支持了`DeepSeek-R1/V3`和`Kimi-K2`模型的`FP8-Static`、`W4A8-FP8`量化。我们还开源了`Hunyuan 1.8B/4B/7B`系列模型的Eagle3权重。
 - [25/07/04] 我们支持了`Hunyuan/Qwen2.5/Qwen3/DeepSeek-R1-Distill-Qwen`等模型的量化，包含INT8、FP8、INT4等算法。
 我们还开源了`Qwen3`系列模型的Eagle3权重。
 
 Coming soon：
-- [ ] DeepSeek-R1的W4A8量化支持
 - [ ] 投机采样新算法发布
 
 ## 🌟主要特性
@@ -310,6 +309,34 @@ Qwen2.5VL系列模型的`BF16`、`FP8-Static`、`FP8-Dynamic`、`INT4-GPTQ`、`I
   </tbody>
 </table>
 
+#### DeepSeek系列模型
+
+DeepSeek-R1-0528模型的`FP8-Block-Wise`、`W4A8-FP8`在`GPQA Diamond`、`AIME 2024`、`SimpleQA`、`LiveCodeBench`上的评测结果如下：
+
+<table>
+  <thead>
+    <tr><th>Model</th><th>Quantization</th><th>GPQA Diamond</th><th>AIME 2024</th><th>SimpleQA</th><th>LiveCodeBench</th></tr>
+  </thead>
+  <tbody>
+    <tr><td rowspan="6">DeepSeek-R1-0528</td><td>FP8-Block-Wise</td><td>78.28</td><td>88.67</td><td>27.8</td><td>77.1</td></tr>
+    <tr><td>W4A8-FP8</td><td>77.37</td><td>88.67</td><td>26.83</td><td>78.86</td></tr>
+  </tbody>
+</table>
+
+> **备注**：
+> - 以上评测结果使用TRT-LLM框架部署测试5次求平均
+> - 评测时使用的超参如下:
+> ```json
+>{
+>  "top_k": 20,
+>  "top_p": 0.6,
+>  "temperature": 0.7,
+>  "output_seq_len": 32768,
+>  "max_input_seq_len": 16384
+>}
+>```
+
+
 #### 其他模型
 
 其他模型的`BF16`、`FP8-Static`、`FP8-Dynamic`、`INT4-GPTQ`、`INT4-AWQ`在`CEVAL`、`MMLU`、`GSM8K`上的评测结果如下：

README_en.md

@@ -31,13 +31,10 @@ Dedicated to building a more intuitive, comprehensive, and efficient LLMs compre
 - [Technical Discussion](#technical-discussion)
 
 ## 📣Latest Updates
-
-- [25/08/04] We now support quantization for `Hunyuan 0.5B/1.8B/4B/7B` and multimodal model `Qwen2.5VL 3B/7B/32B/72B`, including `FP8/INT4` algorithms. We also opensource `Hunyuan 1.8B/4B/7B` series Eagle3 model weight.
+- [25/08/06] We now support quantization for `Hunyuan 0.5B/1.8B/4B/7B` and multimodal model `Qwen2.5VL 3B/7B/32B/72B`, including `FP8/INT4` algorithms, and quantization for `DeepSeek-R1/V3` and `Kimi-K2`, including `FP8-Static` and `W4A8-FP8` algorithms. We also opensource `Hunyuan 1.8B/4B/7B` series Eagle3 model weight.
 - [25/07/04] We now support quantization for `Hunyuan/Qwen2.5/Qwen3/DeepSeek-R1-Distill-Qwen` and other models, including `INT8/FP8/INT4` algorithms. We also opensource `Qwen3` series Eagle3 model weight.
 
 Coming soon:
-
-- [ ] Support W4A8 quantization for DeepSeek-R1.
 - [ ] Release of new algorithm for speculative sampling.
 
 ## 🌟Key Features
@@ -309,6 +306,33 @@ Benchmark results for Qwen2.5VL series models with `BF16`、`FP8-Static`、`FP8-
   </tbody>
 </table>
 
+#### DeepSeek Series Models
+
+Benchmark results for DeepSeek-R1-0528 series models with `FP8-Block-Wise` and `W4A8-FP8` quantization algorithms on datasets including `GPQA Diamond`、`AIME 2024`、`SimpleQA` and `LiveCodeBench`：
+
+<table>
+  <thead>
+    <tr><th>Model</th><th>Quantization</th><th>GPQA Diamond</th><th>AIME 2024</th><th>SimpleQA</th><th>LiveCodeBench</th></tr>
+  </thead>
+  <tbody>
+    <tr><td rowspan="6">DeepSeek-R1-0528</td><td>FP8-Block-Wise</td><td>78.28</td><td>88.67</td><td>27.8</td><td>77.1</td></tr>
+    <tr><td>W4A8-FP8</td><td>77.37</td><td>88.67</td><td>26.83</td><td>78.86</td></tr>
+  </tbody>
+</table>
+
+> **Note**：
+> - The above results are based on the average of 5 test runs deployed with TRT-LLM
+> - The hyperparameters used during evaluation are as follows:
+> ```json
+>{
+>  "top_k": 20,
+>  "top_p": 0.6,
+>  "temperature": 0.7,
+>  "output_seq_len": 32768,
+>  "max_input_seq_len": 16384
+>}
+>```
+
 #### Other Models
 
 Benchmark results for other models with `FP8-Static`, `FP8-Dynamic`, `INT4-GPTQ`, and `INT4-AWQ` quantization algorithms on datasets including `CEVAL`, `MMLU` and `GSM8K`:

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