probe: artifact round-trip verifier (unpack+dequant vs sealed BF16)

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())