64 lines
2.5 KiB
Python
64 lines
2.5 KiB
Python
#!/usr/bin/env python
|
|
"""Issue 5: verify the int8 weight-only artifact round-trips to the BF16 original.
|
|
|
|
Unpacks ``weight_packed`` with compressed-tensors' own unpacker, dequantizes
|
|
with the stored per-channel scales, and compares against the sealed snapshot's
|
|
BF16 tensors. Expected rel-L2 ~1% (the quantizer's measured error); a large
|
|
error means the artifact itself is corrupt, a small error means any serve-path
|
|
collapse lives downstream of the checkpoint (loading/repack/kernel).
|
|
|
|
Usage: verify_artifact.py --artifact ARTIFACT_DIR --snapshot SNAPSHOT_DIR [--tensors N]
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import pathlib
|
|
|
|
import torch
|
|
from compressed_tensors.compressors.pack_quantized.helpers import unpack_from_int32
|
|
from safetensors.torch import load_file
|
|
|
|
QUANTIZED_KEY = "weight_packed" # checkpoint suffix; stripping it yields the base ".weight" name
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--artifact", type=pathlib.Path, required=True)
|
|
parser.add_argument("--snapshot", type=pathlib.Path, required=True)
|
|
parser.add_argument("--tensors", type=int, default=6, help="number of quantized tensors to check")
|
|
args = parser.parse_args()
|
|
|
|
art = load_file(str(args.artifact / "transformer" / "diffusion_pytorch_model.safetensors"))
|
|
index = json.loads(
|
|
(args.snapshot / "transformer" / "diffusion_pytorch_model.safetensors.index.json").read_text()
|
|
)["weight_map"]
|
|
shard_cache: dict[str, dict] = {}
|
|
|
|
checked = 0
|
|
for name in sorted(k.removesuffix("_packed") for k in art if k.endswith(f".{QUANTIZED_KEY}")):
|
|
if checked >= args.tensors:
|
|
break
|
|
base_name = name
|
|
shard = index[base_name]
|
|
if shard not in shard_cache:
|
|
shard_cache.clear() # keep at most one ~7 GB shard resident
|
|
shard_cache[shard] = load_file(str(args.snapshot / "transformer" / shard))
|
|
ref = shard_cache[shard][base_name].to(torch.float32)
|
|
shape = tuple(art[f"{base_name}_shape"].tolist())
|
|
q = unpack_from_int32(art[f"{base_name}_packed"], 8, shape)
|
|
dequant = q.to(torch.float32) * art[f"{base_name}_scale"]
|
|
rel = float((dequant - ref).norm() / ref.norm())
|
|
print(
|
|
f"{base_name}: shape {shape} rel_l2={rel:.5f} "
|
|
f"scale {tuple(art[f'{base_name}_scale'].shape)}/{art[f'{base_name}_scale'].dtype}"
|
|
)
|
|
checked += 1
|
|
print(f"checked {checked} tensors")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|