58 lines
2.1 KiB
Python
58 lines
2.1 KiB
Python
|
"""Expand a `cpu_int4` checkpoint back to a plain bf16 one.
|
||
|
|
|
||
|
|
Used as a control: if the engine is applying the quantized weights faithfully,
|
||
|
|
the expanded checkpoint must score exactly the same perplexity as the
|
||
|
|
quantized one, and any difference is the runtime's, not the quantizer's.
|
||
|
|
"""
|
||
|
|
import argparse, json, shutil
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from safetensors.torch import load_file, save_file
|
||
|
|
|
||
|
|
from vllm_omni.model_executor.layers.quantization.cpu_int4 import unpack_nibbles
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
ap = argparse.ArgumentParser()
|
||
|
|
ap.add_argument("--src", required=True)
|
||
|
|
ap.add_argument("--dst", required=True)
|
||
|
|
a = ap.parse_args()
|
||
|
|
src, dst = Path(a.src), Path(a.dst)
|
||
|
|
dst.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
cfg = json.loads((src / "config.json").read_text())
|
||
|
|
group = cfg["quantization_config"]["group_size"]
|
||
|
|
cfg.pop("quantization_config", None)
|
||
|
|
|
||
|
|
state = {}
|
||
|
|
for shard in sorted(src.glob("*.safetensors")):
|
||
|
|
state.update(load_file(str(shard)))
|
||
|
|
|
||
|
|
out = {}
|
||
|
|
for name, t in state.items():
|
||
|
|
if name.endswith(".weight_packed"):
|
||
|
|
base = name[: -len(".weight_packed")]
|
||
|
|
codes = unpack_nibbles(t).float()
|
||
|
|
n, k = codes.shape
|
||
|
|
s = state[base + ".weight_scale"].float()
|
||
|
|
z = state[base + ".weight_zero"].float()
|
||
|
|
deq = ((codes.reshape(n, k // group, group) - 8.0) * s.unsqueeze(-1)
|
||
|
|
+ z.unsqueeze(-1)).reshape(n, k)
|
||
|
|
out[base + ".weight"] = deq.to(torch.bfloat16)
|
||
|
|
elif name.endswith((".weight_scale", ".weight_zero")):
|
||
|
|
continue
|
||
|
|
else:
|
||
|
|
out[name] = t
|
||
|
|
|
||
|
|
save_file(out, str(dst / "model.safetensors"), metadata={"format": "pt"})
|
||
|
|
(dst / "config.json").write_text(json.dumps(cfg, indent=2))
|
||
|
|
for f in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json",
|
||
|
|
"vocab.json", "merges.txt", "generation_config.json", "chat_template.jinja"):
|
||
|
|
if (src / f).exists():
|
||
|
|
shutil.copy(src / f, dst / f)
|
||
|
|
print(json.dumps({"tensors": len(out), "group": group, "dst": str(dst)}))
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|