157 lines
4.2 KiB
Python
157 lines
4.2 KiB
Python
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__":
|
|
...
|