first commit

dev/functions/dequant.py

@@ -0,0 +1,156 @@
+import re
+import os
+import json
+import torch
+
+from safetensors.torch import load_file
+
+from .specs import *
+from .utils import compile_ignore, get_ignore, estimate_gpu_bytes
+from .dequant_register import DEQUANT_KERNEL_TABLE
+from .config_builder import create_quant_config_builder, resolve_quant_artifacts, to_dequant_spec
+
+
+
+from safetensors import safe_open
+
+
+class LazyTensorDict(dict):
+
+    def __init__(self, base_tensors, tensor_index, input_dir):
+        super().__init__(base_tensors)
+        self.tensor_index = tensor_index
+        self.input_dir = input_dir
+        self.file_cache = {}
+
+    def _load_tensor(self, name):
+
+        file_name = self.tensor_index[name]
+        path = os.path.join(self.input_dir, file_name)
+
+        if file_name not in self.file_cache:
+            self.file_cache[file_name] = safe_open(path, framework="pt")
+        return self.file_cache[file_name].get_tensor(name)
+
+    def __getitem__(self, key):
+        try:
+            return super().__getitem__(key)
+        except KeyError:
+            # 键不在字典中，检查是否在索引中
+            if self.tensor_index and key in self.tensor_index:
+                tensor = self._load_tensor(key)
+                self[key] = tensor  # 现在可以安全地添加
+                return tensor
+            else:
+                raise KeyError(key)
+    
+    def __contains__(self, key):
+
+        if super().__contains__(key):
+            return True
+
+        if self.tensor_index and key in self.tensor_index:
+            return True
+
+        return False
+
+
+
+def dequantize_(
+    tensors: dict[str, torch.Tensor],
+    qconfig: dict,
+    custom_dequant_config,
+    output_dtype=torch.bfloat16,
+    device: torch.device = "cpu",
+    **kwargs
+):
+    ignore_exact, ignore_regex = compile_ignore(
+        get_ignore(qconfig),
+        get_ignore(custom_dequant_config),
+    )
+
+    builder = create_quant_config_builder(
+        qconfig,
+        custom_dequant_config,
+        model_type=kwargs.get("model_type")
+    )
+
+    outputs = {}
+    gpu_cache = {}
+    gpu_cache_size = 0
+    max_gpu_cache_size = 16 * (1024**3)
+    weight_names = list(tensors.keys())
+    # for weight_name in tensors.keys():
+    for weight_name in weight_names:
+        rq = builder.resolve(weight_name)
+
+        if rq is None:  # 包含scale的一定不会在这里
+            assert "scale" not in weight_name, f"{weight_name}"
+            outputs[weight_name] = tensors[weight_name]
+            continue
+
+        assert not rq.dynamic, "Dynamic quantization not supported in offline dequant"
+
+        artifacts = resolve_quant_artifacts(weight_name, tensors, rq)
+
+        if artifacts is None:
+            # 跳过scale, shape等
+            continue
+
+        spec = to_dequant_spec(rq)
+        if spec not in DEQUANT_KERNEL_TABLE:
+            raise RuntimeError(f"No dequant kernel registered for {spec}")
+
+        kernel = DEQUANT_KERNEL_TABLE[spec]
+
+        gpu_cache_size += estimate_gpu_bytes(artifacts, rq)
+
+        try:
+            output = kernel(
+                artifacts=artifacts,
+                rq=rq,
+                output_dtype=output_dtype,
+                device=device,
+            )
+        except Exception as e:
+            print(weight_name)
+            raise e
+        # 处理名称
+        weight_name = weight_name.replace("weight_packed", "weight").replace("qweight", "weight")
+        gpu_cache[weight_name] = output
+
+        if gpu_cache_size >= max_gpu_cache_size:
+            for k, v in gpu_cache.items():
+                outputs[k] = v.cpu()
+            gpu_cache.clear()
+            gpu_cache_size = 0
+    
+    for k, v in gpu_cache.items():
+        outputs[k] = v.cpu()
+        
+    gpu_cache.clear()
+    gpu_cache_size = 0
+
+    return outputs
+
+
+def dequant(file_name: str,
+            quant_config,
+            custom_dequant_config,
+            device: torch.device,
+            model_type,
+            tensor_index=None,
+            input_dir=None):
+
+    base_tensors = load_file(file_name)
+
+    if tensor_index is not None:
+        tensors = LazyTensorDict(base_tensors, tensor_index, input_dir)
+    else:
+        tensors = base_tensors
+
+    return dequantize_(tensors, quant_config, custom_dequant_config, device=device, model_type=model_type)
+
+
+if __name__ == "__main__":
+    ...