144 lines
4.1 KiB
Python
144 lines
4.1 KiB
Python
import os
|
|
import json
|
|
import shutil
|
|
|
|
|
|
def load_config(input_dir: str) -> dict:
|
|
for fname in ("config.json", "params.json"):
|
|
path = os.path.join(input_dir, fname)
|
|
if os.path.exists(path):
|
|
with open(path, "r") as f:
|
|
info = json.load(f)
|
|
return info, fname
|
|
raise ValueError(f"Cannot find quantization config in {input_dir}")
|
|
|
|
|
|
def load_quant_config(input_dir: str) -> dict:
|
|
info = load_config(input_dir)[0]
|
|
qc = info.get("quantization_config", None)
|
|
if qc is None:
|
|
if "text_config" in info:
|
|
qc = info['text_config']["quantization_config"]
|
|
else:
|
|
raise ValueError("Can not find the quantization config")
|
|
|
|
return qc
|
|
|
|
|
|
def get_ignore(quantization_config: dict):
|
|
ignore = []
|
|
|
|
ignore.extend(quantization_config.get("modules_to_not_convert", []))
|
|
ignore.extend(quantization_config.get("ignore", []))
|
|
|
|
return ignore
|
|
|
|
|
|
def copy_files(input_dir, output_dir):
|
|
for file in os.listdir(input_dir):
|
|
src_path = os.path.join(input_dir, file)
|
|
|
|
if os.path.isdir(src_path):
|
|
continue
|
|
|
|
if file.endswith(".safetensors"):
|
|
continue
|
|
|
|
dst_path = os.path.join(output_dir, file)
|
|
shutil.copy(src_path, dst_path)
|
|
|
|
|
|
def process_and_save_info(
|
|
input_dir,
|
|
output_dir,
|
|
qms_output_dir,
|
|
weight_map: dict,
|
|
quant: bool = False,
|
|
symmetric: bool = True,
|
|
ignore: list | None = None,
|
|
group_size: int = -1,
|
|
):
|
|
# 保存weight_map
|
|
candidate_suffixes = [
|
|
"model.safetensors.index.json",
|
|
"consolidated.safetensors.index.json",
|
|
]
|
|
|
|
for cs in candidate_suffixes:
|
|
input_index_path = os.path.join(input_dir, cs)
|
|
output_index_path = os.path.join(output_dir, cs)
|
|
if os.path.exists(input_index_path):
|
|
break
|
|
else:
|
|
raise ValueError(f"Can not find index file, {input_dir}")
|
|
|
|
with open(input_index_path, "r") as f:
|
|
ori_model_index = json.load(f)
|
|
|
|
ori_model_index["weight_map"] = weight_map
|
|
|
|
with open(output_index_path, "w") as f:
|
|
json.dump(ori_model_index, f, indent=2, sort_keys=True)
|
|
|
|
config, fname = load_config(output_dir)
|
|
file_name = os.path.join(output_dir, fname)
|
|
|
|
key = ""
|
|
|
|
try:
|
|
config.pop("quantization_config")
|
|
except Exception as e:
|
|
key = "compression_config"
|
|
pass
|
|
|
|
try:
|
|
config.pop("compression_config")
|
|
except Exception as e:
|
|
key = "quantization_config"
|
|
pass
|
|
|
|
if quant:
|
|
config[key] = {
|
|
"config_groups": {
|
|
"group_0": {
|
|
"input_activations": {
|
|
"block_structure": None,
|
|
"dynamic": True,
|
|
"group_size": None,
|
|
"num_bits": 8,
|
|
"observer": "memoryless",
|
|
"observer_kwargs": {},
|
|
"strategy": "token",
|
|
"symmetric": True,
|
|
"type": "int",
|
|
},
|
|
"output_activations": None,
|
|
"targets": ["Linear"],
|
|
"weights": {
|
|
"block_structure": None,
|
|
"dynamic": False,
|
|
"group_size": group_size if group_size > 0 else None,
|
|
"num_bits": 8,
|
|
"observer": "minmax",
|
|
"observer_kwargs": {},
|
|
"strategy": "group" if group_size > 0 else "channel",
|
|
"symmetric": bool(symmetric),
|
|
"type": "int",
|
|
},
|
|
}
|
|
},
|
|
"format": "int-quantized",
|
|
"global_compression_ratio": 1.0,
|
|
"ignore": ignore,
|
|
"kv_cache_scheme": None,
|
|
"quant_method": "compressed-tensors",
|
|
"quantization_status": "frozen",
|
|
}
|
|
|
|
with open(file_name, "w") as f:
|
|
json.dump(config, f, indent=2)
|
|
|
|
qms_file_name = os.path.join(qms_output_dir, fname)
|
|
with open(qms_file_name, "w") as f:
|
|
json.dump(config, f, indent=2)
|