veriloop-coder-e2-nvfp4 / scripts /validate_nvfp4.py
rodrigoramosrs's picture
Upload scripts/validate_nvfp4.py with huggingface_hub
76376a0 verified
Raw
History Blame Contribute Delete
1.72 kB
"""Structural validation for a modelopt NVFP4 checkpoint.
Checks that every tensor in model.safetensors.index.json exists in the
shards (and vice versa) and prints the dtype layout + quant config.
Usage:
python validate_nvfp4.py [MODEL_DIR] (default: current directory)
"""
import collections
import json
import os
import struct
import sys
DD = sys.argv[1] if len(sys.argv) > 1 else "."
idx = json.load(open(os.path.join(DD, "model.safetensors.index.json")))
wm = idx["weight_map"]
# headers of each shard (no tensor data is read)
have = {}
for fname in set(wm.values()):
with open(os.path.join(DD, fname), "rb") as f:
n = struct.unpack("<Q", f.read(8))[0]
header = json.loads(f.read(n))
for t in header:
if t != "__metadata__":
have[t] = (fname, header[t]["dtype"], header[t]["shape"])
missing = [t for t in wm if t not in have]
extra = [t for t in have if t not in wm]
print("tensors in index:", len(wm))
print("tensors in shards:", len(have))
print("missing:", len(missing), missing[:5])
print("extra:", len(extra), extra[:5])
dtypes = collections.Counter(v[1] for v in have.values())
print("dtypes:", dict(dtypes))
# packed NVFP4 weights show up as U8 with per-block scale tensors
suff = collections.Counter()
for t in have:
s = t.split(".")[-1]
suff[s] += 1
print("common suffixes:", suff.most_common(12))
q = json.load(open(os.path.join(DD, "hf_quant_config.json")))
print("quant_algo:", q["quantization"]["quant_algo"],
"| group_size:", q["quantization"]["group_size"])
assert not missing and not extra, "shard/index mismatch!"
assert q["quantization"]["quant_algo"] == "NVFP4", "not an NVFP4 checkpoint!"
print("OK - valid NVFP4 checkpoint")