|
|
| """Step 7: evaluate an extracted FastH3 Omni LoRA's error relative to TARGET weights.
|
|
|
| The extraction TSV reports ||Delta - LowRank(Delta)|| / ||Delta||. That number
|
| alone cannot quantify how far the reconstructed model is from Omni itself.
|
| This utility calculates ||Target|| and ||Base|| from the original safetensors,
|
| then combines them with the exact per-matrix residual sums in the TSV.
|
|
|
| No torch, GPU, or model checkpoint modifications. Reads in ~4 MiB chunks.
|
| """
|
|
|
| import argparse
|
| import csv
|
| import json
|
| import math
|
| import struct
|
| from collections import defaultdict
|
| from pathlib import Path
|
|
|
| import numpy as np
|
|
|
| PREFIXES = ("model.diffusion_model.", "diffusion_model.")
|
| TYPE_INFO = {"F32": ("<f4", 4), "F16": ("<f2", 2), "BF16": ("<u2", 2)}
|
|
|
|
|
| def canonical(s):
|
| for prefix in PREFIXES:
|
| if s.startswith(prefix):
|
| return s[len(prefix) :]
|
| return s
|
|
|
|
|
| class HeaderReader:
|
| def __init__(self, path):
|
| self.fp = open(path, "rb")
|
| self.path = path
|
| self.entries = {}
|
| try:
|
| header_size = struct.unpack("<Q", self.fp.read(8))[0]
|
| if header_size > 256_000_000:
|
| raise ValueError("Header is implausibly large")
|
| header = json.loads(self.fp.read(header_size))
|
| self.data_at = 8 + header_size
|
| for raw, d in header.items():
|
| if raw == "__metadata__":
|
| continue
|
| k = canonical(raw)
|
| if k in self.entries:
|
| raise ValueError(f"Duplicate canonical tensor: {k}")
|
| self.entries[k] = d
|
| except BaseException:
|
| self.fp.close()
|
| raise
|
|
|
| def chunk(self, key, start, n):
|
| e = self.entries[key]
|
| dtype = e["dtype"]
|
| if dtype not in TYPE_INFO:
|
| raise ValueError(f"Unsupported dtype for {key}: {dtype}")
|
| kind, width = TYPE_INFO[dtype]
|
| count = math.prod(e["shape"])
|
| if start < 0 or n < 0 or start + n > count:
|
| raise ValueError(f"Invalid slice for {key}")
|
| begin = e["data_offsets"][0]
|
| self.fp.seek(self.data_at + begin + start * width)
|
| raw = self.fp.read(n * width)
|
| if len(raw) != n * width:
|
| raise IOError(f"Short read of {key} in {self.path}")
|
| a = np.frombuffer(raw, dtype=kind)
|
| if dtype == "BF16":
|
| return (a.astype(np.uint32) << np.uint32(16)).view(np.float32)
|
| return a.astype(np.float32, copy=False)
|
|
|
| def close(self):
|
| self.fp.close()
|
|
|
|
|
| def family(key):
|
| if key.startswith("blocks."):
|
| if ".attn.qkv_proj." in key:
|
| return "backbone QKV"
|
| if ".attn.out_proj." in key:
|
| return "backbone attention out"
|
| if ".mlp.fc1." in key:
|
| return "backbone MLP FC1"
|
| if ".mlp.fc2." in key:
|
| return "backbone MLP FC2"
|
| return "other"
|
|
|
|
|
| def fmt(n, total, detail):
|
| if total <= 0:
|
| return "n/a"
|
| return f"{100 * math.sqrt(max(n, 0.0) / total):.5f}%"
|
|
|
|
|
| def main():
|
| parser = argparse.ArgumentParser(description=__doc__)
|
| parser.add_argument("base", help="Original Ref2VA base .safetensors")
|
| parser.add_argument("target", help="FastH3 Omni .safetensors")
|
| parser.add_argument("report", help="FastH3_Omni_r64_report.tsv")
|
| parser.add_argument("--out", default="h3_omni_r64_weight_error.tsv")
|
| parser.add_argument("--chunk-elements", type=int, default=1_000_000)
|
| args = parser.parse_args()
|
| if args.chunk_elements <= 0:
|
| parser.error("--chunk-elements must be positive")
|
| with open(args.report, encoding="utf-8", newline="") as fp:
|
| input_rows = list(csv.DictReader(fp, delimiter="\t"))
|
| if not input_rows or not all(
|
| x in input_rows[0] for x in ("key", "true_delta_norm_sq", "residual_norm_sq")
|
| ):
|
| parser.error(
|
| "Invalid extraction TSV; expected key,true_delta_norm_sq,residual_norm_sq"
|
| )
|
| if len({r["key"] for r in input_rows}) != len(input_rows):
|
| parser.error("Duplicate key in extraction TSV")
|
| b, t = HeaderReader(args.base), HeaderReader(args.target)
|
| aggregate = defaultdict(
|
| lambda: dict(num=0, base=0.0, target=0.0, delta=0.0, unreconstructed=0.0)
|
| )
|
| all_rows = []
|
| mismatched_delta_norms = []
|
| try:
|
| for i, r in enumerate(input_rows, 1):
|
| key = canonical(r["key"])
|
| if key not in b.entries or key not in t.entries:
|
| raise ValueError(f"Missing tensor in base or target: {key}")
|
| shape = b.entries[key]["shape"]
|
| if shape != t.entries[key]["shape"]:
|
| raise ValueError(f"Shape mismatch: {key}")
|
| n = math.prod(shape)
|
| bn = tn = dn = 0.0
|
| for off in range(0, n, args.chunk_elements):
|
| length = min(args.chunk_elements, n - off)
|
| bf = b.chunk(key, off, length)
|
| tf = t.chunk(key, off, length)
|
|
|
| bv = bf.astype(np.float64)
|
| tv = tf.astype(np.float64)
|
| dv = tv - bv
|
| bn += float(np.dot(bv, bv))
|
| tn += float(np.dot(tv, tv))
|
| dn += float(np.dot(dv, dv))
|
| reported_dn = float(r["true_delta_norm_sq"])
|
| error = float(r["residual_norm_sq"])
|
| if abs(dn - reported_dn) > max(1e-5, 2e-4 * reported_dn):
|
| mismatched_delta_norms.append((key, dn, reported_dn))
|
| if error < -1e-6 or reported_dn < -1e-6:
|
| raise ValueError(f"Invalid norm in report: {key}")
|
| row = dict(
|
| key=key,
|
| family=family(key),
|
| base_norm_sq=bn,
|
| target_norm_sq=tn,
|
| delta_norm_sq=dn,
|
| residual_norm_sq=error,
|
| delta_over_target_percent=(
|
| 100 * math.sqrt(dn / tn) if tn else float("nan")
|
| ),
|
| residual_over_target_percent=(
|
| 100 * math.sqrt(error / tn) if tn else float("nan")
|
| ),
|
| residual_over_delta_percent=(
|
| 100 * math.sqrt(error / dn) if dn else float("nan")
|
| ),
|
| )
|
| all_rows.append(row)
|
| a = aggregate[row["family"]]
|
| a["num"] += 1
|
| a["base"] += bn
|
| a["target"] += tn
|
| a["delta"] += dn
|
| a["unreconstructed"] += error
|
| if i % 20 == 0 or i == len(input_rows):
|
| print(f"[{i:3}/{len(input_rows)}] checked {key}", flush=True)
|
| finally:
|
| b.close()
|
| t.close()
|
| with open(args.out, "w", encoding="utf-8", newline="") as fp:
|
| fields = list(all_rows[0].keys())
|
| writer = csv.DictWriter(fp, fieldnames=fields, delimiter="\t")
|
| writer.writeheader()
|
| writer.writerows(all_rows)
|
| whole = dict(num=0, base=0.0, target=0.0, delta=0.0, unreconstructed=0.0)
|
| for a in aggregate.values():
|
| for k in whole:
|
| whole[k] += a[k]
|
| print("\n" + "=" * 76)
|
| print("STEP 7: R64 ERROR RELATIVE TO OMNI WEIGHTS")
|
| print("=" * 76)
|
| print(f"Matrices analyzed: {whole['num']}")
|
| print(
|
| f"Base -> Omni change: {fmt(whole['delta'], whole['target'], '')} of target norm"
|
| )
|
| print(
|
| f"r64 reconstruction error: {fmt(whole['unreconstructed'], whole['target'], '')} of target norm"
|
| )
|
| print(
|
| f"r64 delta residual: {fmt(whole['unreconstructed'], whole['delta'], '')} of delta norm"
|
| )
|
| print(
|
| f"Squared delta energy kept: {100*(1-whole['unreconstructed']/whole['delta']):.2f}%"
|
| )
|
| print("\nPer matrix family:")
|
| for name in sorted(aggregate):
|
| a = aggregate[name]
|
| print(
|
| f" {name:23} {a['num']:3} "
|
| + f"error/target={fmt(a['unreconstructed'],a['target'],''):>10} "
|
| + f"error/delta={fmt(a['unreconstructed'],a['delta'],''):>10}"
|
| )
|
| print("\nTop 10 matrices by residual error relative to their TARGET weight:")
|
| for r in sorted(
|
| all_rows, key=lambda x: x["residual_over_target_percent"], reverse=True
|
| )[:10]:
|
| print(f" {r['residual_over_target_percent']:9.4f}% {r['key']}")
|
| print(f"\nPer-layer TSV: {args.out}")
|
| if mismatched_delta_norms:
|
| print(
|
| f"\nWARNING: {len(mismatched_delta_norms)} delta-norm mismatches against extraction TSV."
|
| )
|
| for k, observed, reported in mismatched_delta_norms[:8]:
|
| print(f" {k}: recomputed={observed:.8e} report={reported:.8e}")
|
| print("Check the report matches the exact base and target checkpoints!")
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
