File size: 6,896 Bytes
c61c435 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | from __future__ import annotations
import json
from pathlib import Path
from typing import Any
from .base import ModelComparison, TensorComparison, is_cancelled, report
from .detector import inspect_model
from .generic import _extract_state_dict, _weight_files
from .statistics import component_for_name
def _iter_named_tensors(path: Path):
files = _weight_files(path)
for file in files:
if file.suffix.casefold() == ".safetensors":
from safetensors import safe_open
with safe_open(str(file), framework="pt", device="cpu") as handle:
for key in handle.keys():
yield key, handle.get_tensor(key)
else:
import torch
payload = torch.load(str(file), map_location="cpu", weights_only=False)
state = _extract_state_dict(payload)
for key, value in state.items():
if hasattr(value, "shape"):
yield str(key), value
def _tensor_map(path: Path) -> dict[str, Any]:
return {name: tensor for name, tensor in _iter_named_tensors(path)}
def _config_differences(configs_a: dict[str, Any], configs_b: dict[str, Any], *, limit: int = 40) -> list[str]:
diffs: list[str] = []
keys = sorted(set(configs_a) | set(configs_b))
for key in keys:
if key not in configs_a:
diffs.append(f"Only B has config {key}")
elif key not in configs_b:
diffs.append(f"Only A has config {key}")
elif json.dumps(configs_a[key], sort_keys=True, default=str) != json.dumps(configs_b[key], sort_keys=True, default=str):
diffs.append(f"Config differs: {key}")
if len(diffs) >= limit:
diffs.append("Additional config differences omitted.")
break
return diffs
def compare_models(
path_a: str | Path,
path_b: str | Path,
*,
arch_a: str = "",
arch_b: str = "",
settings_a: dict[str, Any] | None = None,
settings_b: dict[str, Any] | None = None,
progress=None,
cancelled=None,
) -> ModelComparison:
report(progress, 2, "Inspecting first model")
summary_a = inspect_model(path_a, recorded_architecture=arch_a, run_settings=settings_a, progress=progress, cancelled=cancelled)
report(progress, 30, "Inspecting second model")
summary_b = inspect_model(path_b, recorded_architecture=arch_b, run_settings=settings_b, progress=progress, cancelled=cancelled)
report(progress, 55, "Loading comparable tensors")
tensors_a = _tensor_map(Path(summary_a.resolved_path))
if is_cancelled(cancelled):
raise RuntimeError("Comparison cancelled.")
tensors_b = _tensor_map(Path(summary_b.resolved_path))
names_a = set(tensors_a)
names_b = set(tensors_b)
common = sorted(names_a & names_b)
only_a = sorted(names_a - names_b)[:200]
only_b = sorted(names_b - names_a)[:200]
shape_mismatches = []
comparable = []
for name in common:
if tuple(tensors_a[name].shape) != tuple(tensors_b[name].shape):
shape_mismatches.append(name)
else:
comparable.append(name)
comparisons: list[TensorComparison] = []
import torch
for index, name in enumerate(comparable):
if is_cancelled(cancelled):
raise RuntimeError("Comparison cancelled.")
if index % 10 == 0:
report(progress, 58 + int(36 * index / max(1, len(comparable))), f"Comparing {index + 1} of {len(comparable)} tensors")
with torch.no_grad():
a = tensors_a[name].detach().to(device="cpu").float().reshape(-1)
b = tensors_b[name].detach().to(device="cpu").float().reshape(-1)
if a.numel() == 0:
continue
limit = 1_000_000
if a.numel() > limit:
stride = max(1, a.numel() // limit)
a = a[::stride][:limit]
b = b[::stride][:limit]
delta = b - a
mean_abs = delta.abs().mean().item()
base_abs = a.abs().mean().item()
relative = mean_abs / (base_abs + 1e-12)
l2 = torch.linalg.vector_norm(delta).item()
norm_a = torch.linalg.vector_norm(a).item()
norm_b = torch.linalg.vector_norm(b).item()
cosine = torch.nn.functional.cosine_similarity(a, b, dim=0).item() if norm_a and norm_b else None
drift = (norm_b - norm_a) / (norm_a + 1e-12) if norm_a else None
score = relative * 0.7 + (1 - cosine if cosine is not None else 0) * 0.3
comparisons.append(
TensorComparison(
name=name,
shape=tuple(int(dim) for dim in tensors_a[name].shape),
component=component_for_name(name),
mean_abs_difference=float(mean_abs),
relative_difference=float(relative),
cosine_similarity=float(cosine) if cosine is not None else None,
l2_distance=float(l2),
drift=float(drift) if drift is not None else None,
change_score=float(score),
)
)
comparisons.sort(key=lambda item: item.change_score or 0, reverse=True)
groups: dict[str, dict[str, float]] = {}
for item in comparisons:
group = groups.setdefault(item.component, {"tensors": 0, "mean_change_score": 0.0, "mean_abs_difference": 0.0})
group["tensors"] += 1
group["mean_change_score"] += item.change_score or 0
group["mean_abs_difference"] += item.mean_abs_difference or 0
for group in groups.values():
count = max(1, int(group["tensors"]))
group["mean_change_score"] /= count
group["mean_abs_difference"] /= count
messages = [
"Change Score is a statistical weight-change metric; it does not directly equal behavioral importance."
]
if shape_mismatches:
messages.append("Some tensor comparisons are unavailable because tensor shapes differ.")
report(progress, 100, "Comparison complete")
return ModelComparison(
path_a=summary_a.resolved_path,
path_b=summary_b.resolved_path,
architecture_a=summary_a.architecture,
architecture_b=summary_b.architecture,
architecture_match=summary_a.architecture == summary_b.architecture,
config_differences=_config_differences(summary_a.configs, summary_b.configs),
resolution_difference=(
(summary_a.resolution, summary_b.resolution)
if summary_a.resolution != summary_b.resolution else None
),
parameter_count_difference=summary_b.total_parameters - summary_a.total_parameters,
only_a=only_a,
only_b=only_b,
shape_mismatches=shape_mismatches[:200],
tensor_comparisons=comparisons[:200],
group_comparisons=groups,
messages=messages,
)
|