diff --git a/probe/issue-5/verify_artifact.py b/probe/issue-5/verify_artifact.py new file mode 100644 index 000000000..cbf5e5dee --- /dev/null +++ b/probe/issue-5/verify_artifact.py @@ -0,0 +1,63 @@ +#!/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())