|
|
| """Compare two MiniMax H3 Safetensors checkpoints without torch or safetensors bindings.
|
|
|
| Only Python's JSON/file I/O and NumPy are used. One small block from each
|
| checkpoint is read at a time. This avoids torch native crashes and holding
|
| large transformer weights in RAM. Original checkpoints are never modified.
|
| """
|
|
|
| import argparse
|
| import json
|
| import math
|
| import os
|
| import struct
|
| from collections import Counter
|
| from pathlib import Path
|
|
|
| import numpy as np
|
|
|
| PREFIXES = ("model.diffusion_model.", "diffusion_model.")
|
|
|
|
|
| DTYPES = {
|
| "F64": ("<f8", 8),
|
| "F32": ("<f4", 4),
|
| "F16": ("<f2", 2),
|
| "BF16": ("<u2", 2),
|
| "I8": ("i1", 1),
|
| "U8": ("u1", 1),
|
| "BOOL": ("?", 1),
|
| "I16": ("<i2", 2),
|
| "U16": ("<u2", 2),
|
| "I32": ("<i4", 4),
|
| "U32": ("<u4", 4),
|
| "I64": ("<i8", 8),
|
| "U64": ("<u8", 8),
|
| }
|
|
|
|
|
| def normalize_key(key):
|
| for prefix in PREFIXES:
|
| if key.startswith(prefix):
|
| return key[len(prefix) :]
|
| return key
|
|
|
|
|
| class Checkpoint:
|
| def __init__(self, path):
|
| self.path = Path(path)
|
| self.file = self.path.open("rb")
|
| try:
|
| raw_len = self.file.read(8)
|
| if len(raw_len) != 8:
|
| raise ValueError(f"Invalid Safetensors header in {path}")
|
| hlen = struct.unpack("<Q", raw_len)[0]
|
| if hlen > 512 * 1024 * 1024:
|
| raise ValueError("Header unexpectedly large")
|
| raw_header = self.file.read(hlen)
|
| if len(raw_header) != hlen:
|
| raise ValueError(f"Truncated Safetensors header in {path}")
|
| header = json.loads(raw_header)
|
| self.data_start = 8 + hlen
|
| self.entries = {}
|
| max_bytes = self.path.stat().st_size - self.data_start
|
| for raw_key, info in header.items():
|
| if raw_key == "__metadata__":
|
| continue
|
| key = normalize_key(raw_key)
|
| if key in self.entries:
|
| raise ValueError(f"Duplicate normalized tensor key: {key}")
|
| dtype = info["dtype"]
|
| if dtype not in DTYPES:
|
| raise ValueError(f"Unsupported dtype {dtype} for {key}")
|
| shape = tuple(info["shape"])
|
| start, end = info["data_offsets"]
|
| count = math.prod(shape)
|
| nbytes = count * DTYPES[dtype][1]
|
| if start < 0 or end < start or end > max_bytes or end - start != nbytes:
|
| raise ValueError(f"Invalid offsets or shape for {key}")
|
| self.entries[key] = {
|
| "raw_key": raw_key,
|
| "dtype": dtype,
|
| "shape": shape,
|
| "start": start,
|
| "count": count,
|
| }
|
| except BaseException:
|
| self.file.close()
|
| raise
|
|
|
| def close(self):
|
| self.file.close()
|
|
|
| def read_chunk(self, key, element_offset, n):
|
| entry = self.entries[key]
|
| dtype = entry["dtype"]
|
| width = DTYPES[dtype][1]
|
| offset = self.data_start + entry["start"] + element_offset * width
|
| self.file.seek(offset)
|
| data = self.file.read(n * width)
|
| if len(data) != n * width:
|
| raise IOError(f"Unexpected EOF reading {key}")
|
| return data
|
|
|
| def __enter__(self):
|
| return self
|
|
|
| def __exit__(self, *_):
|
| self.close()
|
|
|
|
|
| def to_float32(raw, dtype):
|
| values = np.frombuffer(raw, dtype=np.dtype(DTYPES[dtype][0]))
|
| if dtype == "BF16":
|
|
|
| return np.left_shift(values.astype(np.uint32), 16).view(np.float32)
|
| return values.astype(np.float32, copy=False)
|
|
|
|
|
| def cast_f32_to_dtype(values, dtype):
|
| """Float32 representation after rounding values to a storage dtype."""
|
| if dtype == "F32":
|
| return values
|
| if dtype == "F16":
|
| return values.astype(np.float16).astype(np.float32)
|
| if dtype == "BF16":
|
|
|
| bits = values.astype(np.float32, copy=False).view(np.uint32)
|
| rounded = bits + (np.uint32(0x7FFF) + ((bits >> 16) & 1))
|
| return (rounded & np.uint32(0xFFFF0000)).view(np.float32)
|
| if dtype == "F64":
|
| return values
|
| np_dtype = np.dtype(DTYPES[dtype][0])
|
| return values.astype(np_dtype).astype(np.float32)
|
|
|
|
|
| def array_equal(a, b):
|
| return np.array_equal(a, b, equal_nan=True)
|
|
|
|
|
| def compare_tensor(base, target, key, chunk_elements):
|
| be = base.entries[key]
|
| te = target.entries[key]
|
| if be["shape"] != te["shape"]:
|
| raise ValueError(f"Shape mismatch: {key}")
|
|
|
| identical_bytes = be["dtype"] == te["dtype"]
|
| base_casts_to_target = True
|
| target_casts_to_base = True
|
| diff_sq = base_sq = target_sq = dot = 0.0
|
| max_abs = 0.0
|
| processed = 0
|
| for start in range(0, be["count"], chunk_elements):
|
| n = min(chunk_elements, be["count"] - start)
|
| braw = base.read_chunk(key, start, n)
|
| traw = target.read_chunk(key, start, n)
|
|
|
| if identical_bytes and braw != traw:
|
| identical_bytes = False
|
| b = to_float32(braw, be["dtype"])
|
| t = to_float32(traw, te["dtype"])
|
|
|
| if not np.isfinite(b).all() or not np.isfinite(t).all():
|
| raise ValueError(f"Non-finite values in tensor {key} at offset {start}")
|
| if base_casts_to_target and not array_equal(
|
| cast_f32_to_dtype(b, te["dtype"]), t
|
| ):
|
| base_casts_to_target = False
|
| if target_casts_to_base and not array_equal(
|
| cast_f32_to_dtype(t, be["dtype"]), b
|
| ):
|
| target_casts_to_base = False
|
|
|
| diff = t - b
|
| diff_sq += float(np.einsum("i,i->", diff, diff, dtype=np.float64))
|
| base_sq += float(np.einsum("i,i->", b, b, dtype=np.float64))
|
| target_sq += float(np.einsum("i,i->", t, t, dtype=np.float64))
|
| dot += float(np.einsum("i,i->", b, t, dtype=np.float64))
|
| max_abs = max(max_abs, float(np.max(np.abs(diff))))
|
| processed += n
|
|
|
| dnorm = math.sqrt(diff_sq)
|
| bnorm = math.sqrt(base_sq)
|
| relerr = dnorm / bnorm if bnorm else (0.0 if not dnorm else float("inf"))
|
| denom = math.sqrt(base_sq * target_sq)
|
| cosine = dot / denom if denom else (1.0 if not diff_sq else 0.0)
|
| cosine = min(1.0, max(-1.0, cosine))
|
| if identical_bytes:
|
| classification = "IDENTICAL"
|
| elif be["dtype"] != te["dtype"] and (base_casts_to_target or target_casts_to_base):
|
| classification = "CAST_EQ"
|
| else:
|
| classification = "CHANGED"
|
| return {
|
| "key": key,
|
| "shape": be["shape"],
|
| "base_dtype": be["dtype"],
|
| "target_dtype": te["dtype"],
|
| "classification": classification,
|
| "relerr": relerr,
|
| "cosine": cosine,
|
| "max_abs": max_abs,
|
| "elements": processed,
|
| }
|
|
|
|
|
| def main():
|
| parser = argparse.ArgumentParser(
|
| description="Compare MiniMax H3 checkpoint values without PyTorch or CUDA"
|
| )
|
| parser.add_argument("base", help="Base Ref2VA safetensors checkpoint")
|
| parser.add_argument("target", help="FastH3 Omni safetensors checkpoint")
|
| parser.add_argument("--report", default="h3_omni_value_comparison.tsv")
|
| parser.add_argument("--chunk-elements", type=int, default=1_000_000)
|
| args = parser.parse_args()
|
| if args.chunk_elements < 1:
|
| parser.error("--chunk-elements must be >= 1")
|
|
|
| with Checkpoint(args.base) as base, Checkpoint(args.target) as target:
|
| shared = sorted(set(base.entries) & set(target.entries))
|
| print(f"Shared tensors: {len(shared)}", flush=True)
|
| print("Using NumPy + direct Safetensors reads (no torch)\n", flush=True)
|
| results = []
|
| skipped = []
|
| for i, key in enumerate(shared, 1):
|
| bshape = base.entries[key]["shape"]
|
| tshape = target.entries[key]["shape"]
|
| if bshape != tshape:
|
| skipped.append(key)
|
| print(
|
| f"[{i:3}/{len(shared)}] SHAPE {key}: {bshape} -> {tshape}",
|
| flush=True,
|
| )
|
| continue
|
| print(f"[{i:3}/{len(shared)}] START {key}", flush=True)
|
| result = compare_tensor(base, target, key, args.chunk_elements)
|
| results.append(result)
|
| print(
|
| f"[{i:3}/{len(shared)}] {result['classification']:9} "
|
| f"rel={result['relerr']*100:9.5f}% "
|
| f"cos={result['cosine']:.8f} {key}",
|
| flush=True,
|
| )
|
|
|
| counter = Counter(r["classification"] for r in results)
|
| changed = sorted(
|
| (r for r in results if r["classification"] == "CHANGED"),
|
| key=lambda r: r["relerr"],
|
| reverse=True,
|
| )
|
| print("\n" + "=" * 80)
|
| print("SUMMARY")
|
| print("=" * 80)
|
| print(f"Comparable tensors: {len(results)}")
|
| print(f"Shape mismatches: {len(skipped)}")
|
| print(f"Identical: {counter['IDENTICAL']}")
|
| print(f"Dtype/cast only: {counter['CAST_EQ']}")
|
| print(f"Actually changed: {counter['CHANGED']}")
|
| print("\nTOP 30 CHANGED")
|
| for item in changed[:30]:
|
| print(f"{item['relerr']*100:10.6f}% cos {item['cosine']:.8f} {item['key']}")
|
|
|
| report = Path(args.report)
|
| with report.open("w", encoding="utf-8", newline="") as output:
|
| output.write(
|
| "classification\trelerr_percent\tcosine\tmax_abs\tbase_dtype\ttarget_dtype\tshape\tkey\n"
|
| )
|
| for item in sorted(results, key=lambda x: x["relerr"], reverse=True):
|
| output.write(
|
| f"{item['classification']}\t{item['relerr']*100:.10f}\t"
|
| f"{item['cosine']:.10f}\t{item['max_abs']:.10e}\t"
|
| f"{item['base_dtype']}\t{item['target_dtype']}\t"
|
| f"{item['shape']}\t{item['key']}\n"
|
| )
|
| print(f"\nFull report: {report.resolve()}")
|
| print(
|
| "Note: CAST_EQ means identical after a dtype cast, not proof of no training change."
|
| )
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
