NVFP4 (ModelOpt) quantization of sensenova/SenseNova-Loop-9B-Preview @27abea77 with BF16 MTP head
a57aa9f verified | #!/usr/bin/env python3 | |
| """Independent structural verification of an exported NVFP4 checkpoint (reads files only). | |
| Checks against the upstream checkpoint: every upstream tensor is present (as BF16/FP32 or as an | |
| NVFP4 packed weight + scales), packed shapes are consistent (uint8 [out, in/2], fp8 e4m3 block scales | |
| [out, in/16], fp32 global scales), all scales finite/positive, mHC routers FP32, exactly one copy of | |
| each physical (shared, looped) layer, MTP tensors present, config/quant metadata consistent. | |
| usage: python -I scripts/30_verify_export.py EXPORT_DIR UPSTREAM_DIR OUT_JSON | |
| """ | |
| import json, os, sys, collections | |
| import torch | |
| from safetensors import safe_open | |
| exp, up, out = sys.argv[1:4] | |
| def index(d): | |
| idx = os.path.join(d, "model.safetensors.index.json") | |
| if os.path.exists(idx): | |
| wm = json.load(open(idx))["weight_map"] | |
| else: | |
| wm = {} | |
| for f in os.listdir(d): | |
| if f.endswith(".safetensors"): | |
| with safe_open(os.path.join(d, f), "pt") as fh: | |
| wm.update({k: f for k in fh.keys()}) | |
| return wm | |
| def header(d, wm): | |
| meta = {} | |
| for f in set(wm.values()): | |
| with safe_open(os.path.join(d, f), "pt") as fh: | |
| for k in fh.keys(): | |
| sl = fh.get_slice(k) | |
| meta[k] = (sl.get_dtype(), tuple(sl.get_shape()), f) | |
| return meta | |
| ewm, uwm = index(exp), index(up) | |
| E, U = header(exp, ewm), header(up, uwm) | |
| qcfg = json.load(open(os.path.join(exp, "hf_quant_config.json")))["quantization"] | |
| cfg = json.load(open(os.path.join(exp, "config.json"))) | |
| rep = {"export_dir": exp, "quant_algo": qcfg.get("quant_algo"), "group_size": qcfg.get("group_size"), | |
| "kv_cache_quant_algo": qcfg.get("kv_cache_quant_algo"), | |
| "config_has_quantization_config": "quantization_config" in cfg, | |
| "n_tensors_export": len(E), "n_tensors_upstream": len(U)} | |
| problems = [] | |
| packed = [k[: -len(".weight")] for k, (dt, _, _) in E.items() if k.endswith(".weight") and dt == "U8"] | |
| rep["nvfp4_linear_modules"] = len(packed) | |
| q_params = 0 | |
| for m in packed: | |
| dt, shp, _ = E[m + ".weight"] | |
| ref = U.get(m + ".weight") | |
| if ref is None: | |
| problems.append(f"packed {m} not in upstream"); continue | |
| o, i = ref[1] | |
| q_params += o * i | |
| if shp != (o, i // 2): | |
| problems.append(f"{m}: packed shape {shp} != {(o, i // 2)}") | |
| s = E.get(m + ".weight_scale") | |
| if s is None or s[0] != "F8_E4M3" or s[1] != (o, i // qcfg["group_size"]): | |
| problems.append(f"{m}: weight_scale {s}") | |
| for k in ("weight_scale_2", "input_scale"): | |
| t = E.get(f"{m}.{k}") | |
| if t is None or t[0] != "F32" or (len(t[1]) and t[1] != (1,)): | |
| problems.append(f"{m}: {k} {t}") | |
| # Load all scales and check finite / positive. | |
| bad_scales, n_scales = 0, 0 | |
| by_file = collections.defaultdict(list) | |
| for k in E: | |
| if k.endswith(("weight_scale", "weight_scale_2", "input_scale")): | |
| by_file[E[k][2]].append(k) | |
| for f, keys in by_file.items(): | |
| with safe_open(os.path.join(exp, f), "pt") as fh: | |
| for k in keys: | |
| t = fh.get_tensor(k).float() | |
| n_scales += 1 | |
| if not torch.isfinite(t).all() or (t[t != 0] <= 0).any() or (k.endswith("input_scale") and (t <= 0).any()): | |
| bad_scales += 1 | |
| problems.append(f"bad scale {k}") | |
| rep["scale_tensors_checked"] = n_scales | |
| rep["bad_scale_tensors"] = bad_scales | |
| # Upstream coverage: every upstream tensor present in export (either same name or packed). | |
| missing = [k for k in U if k not in E] | |
| rep["upstream_tensors_missing_in_export"] = missing[:20] | |
| rep["n_upstream_missing"] = len(missing) | |
| extra = [k for k in E if k not in U and not k.endswith(("weight_scale", "weight_scale_2", "input_scale"))] | |
| rep["unexpected_extra_tensors"] = extra[:20] | |
| mhc = {k: E[k][0] for k in E if ".mhc_" in k} | |
| rep["mhc_tensors"] = len(mhc) | |
| rep["mhc_all_fp32"] = all(v == "F32" for v in mhc.values()) and len(mhc) == 96 | |
| mtp = [k for k in E if k.startswith("mtp.")] | |
| rep["mtp_tensors"] = len(mtp) | |
| rep["mtp_dtypes"] = sorted({E[k][0] for k in mtp}) | |
| # Shared weights: layer ids present must be exactly 0..31 (no duplicated per-visit copies). | |
| lids = sorted({int(k.split(".layers.")[1].split(".")[0]) for k in E if "language_model.layers." in k}) | |
| rep["physical_layer_ids"] = f"{lids[0]}..{lids[-1]} ({len(lids)})" | |
| if lids != list(range(32)): | |
| problems.append(f"layer ids {lids}") | |
| dtypes = collections.Counter() | |
| nbytes = collections.Counter() | |
| size = {"U8": 1, "F8_E4M3": 1, "BF16": 2, "F16": 2, "F32": 4} | |
| import math | |
| for k, (dt, shp, _) in E.items(): | |
| n = math.prod(shp) | |
| dtypes[dt] += n | |
| nbytes[dt] += n * size.get(dt, 0) | |
| rep["elements_by_dtype"] = dict(dtypes) | |
| rep["GB_by_dtype"] = {k: round(v / 1e9, 3) for k, v in nbytes.items()} | |
| rep["total_file_GB"] = round(sum(os.path.getsize(os.path.join(exp, f)) for f in set(ewm.values())) / 1e9, 3) | |
| total_up = sum(math.prod(s) for (_, s, _) in U.values()) | |
| rep["upstream_params_incl_mtp"] = total_up | |
| rep["nvfp4_params"] = q_params | |
| rep["fraction_params_nvfp4"] = round(q_params / total_up, 4) | |
| rep["problems"] = problems | |
| rep["ok"] = not problems and not missing and rep["mhc_all_fp32"] and len(mtp) > 0 | |
| json.dump(rep, open(out, "w"), indent=2) | |
| print(json.dumps(rep, indent=2)) | |