308 lines
9.3 KiB
Python
308 lines
9.3 KiB
Python
"""SVD effective rank & weight statistics analysis for model checkpoints."""
|
||
|
||
import argparse
|
||
import json
|
||
from pathlib import Path
|
||
|
||
import safetensors.torch
|
||
import torch
|
||
|
||
|
||
def effective_rank_metrics(w: torch.Tensor) -> dict:
|
||
if w.ndim == 1:
|
||
return {"shape": tuple(w.shape), "is_1d": True}
|
||
|
||
w = w.float()
|
||
s = torch.linalg.svdvals(w)
|
||
s_sq = s**2
|
||
total = s_sq.sum()
|
||
cumsum = torch.cumsum(s_sq, dim=0) / total
|
||
|
||
min_dim = min(w.shape[0], w.shape[1])
|
||
er_90 = (cumsum < 0.90).sum().item() + 1
|
||
er_95 = (cumsum < 0.95).sum().item() + 1
|
||
er_99 = (cumsum < 0.99).sum().item() + 1
|
||
|
||
p = s_sq / total
|
||
p = p[p > 1e-30]
|
||
entropy = -(p * torch.log(p)).sum()
|
||
entropic_rank = torch.exp(entropy).item()
|
||
|
||
return {
|
||
"shape": tuple(w.shape),
|
||
"min_dim": min_dim,
|
||
"er_90": er_90,
|
||
"er_95": er_95,
|
||
"er_99": er_99,
|
||
"er_99_norm": er_99 / min_dim,
|
||
"er_95_norm": er_95 / min_dim,
|
||
"entropic_rank": entropic_rank,
|
||
"entropic_rank_norm": entropic_rank / min_dim,
|
||
"top1_ratio": s[0].item() / s.sum().item(),
|
||
"top5_ratio": s[:5].sum().item() / s.sum().item(),
|
||
"decay_ratio": s[-1].item() / s[0].item(),
|
||
"condition_number": s[0].item() / s[-1].item(),
|
||
"mean": w.mean().item(),
|
||
"std": w.std().item(),
|
||
"min": w.min().item(),
|
||
"max": w.max().item(),
|
||
}
|
||
|
||
|
||
def format_header(headers: list[str], widths: list[int]) -> str:
|
||
return "".join(h.ljust(w) for h, w in zip(headers, widths))
|
||
|
||
|
||
def format_row(values: list[str], widths: list[int]) -> str:
|
||
return "".join(v.ljust(w) for v, w in zip(values, widths))
|
||
|
||
|
||
def group_by_component(results: dict[str, dict]) -> dict[str, list[dict]]:
|
||
groups: dict[str, list[dict]] = {}
|
||
for key, r in results.items():
|
||
parts = key.split(".")
|
||
if parts[0] == "layers" and len(parts) >= 3:
|
||
sub = parts[2:]
|
||
if sub[0] == "attention":
|
||
comp = f"attn.{sub[1]}"
|
||
elif sub[0] == "mlp":
|
||
comp = f"mlp.{sub[1]}"
|
||
elif sub[0] == "input_norm":
|
||
comp = "input_norm"
|
||
elif sub[0] == "post_attention_norm":
|
||
comp = "post_attn_norm"
|
||
else:
|
||
comp = ".".join(sub)
|
||
else:
|
||
comp = key
|
||
groups.setdefault(comp, []).append(r)
|
||
return groups
|
||
|
||
|
||
def print_component_summary(results: dict[str, dict], title: str):
|
||
groups = group_by_component(results)
|
||
matrix_groups = {
|
||
k: [v for v in vs if not v.get("is_1d")]
|
||
for k, vs in groups.items()
|
||
if any(not v.get("is_1d") for v in vs)
|
||
}
|
||
|
||
widths = [20, 12, 12, 12, 12, 12]
|
||
print(f"\n{title}")
|
||
print(
|
||
format_header(
|
||
["Component", "N", "ER@99%", "EntRank%", "Top1 σ(%)", "Cond. Num"], widths
|
||
)
|
||
)
|
||
print("-" * sum(widths))
|
||
|
||
for name in sorted(matrix_groups.keys()):
|
||
items = matrix_groups[name]
|
||
n = len(items)
|
||
print(
|
||
format_row(
|
||
[
|
||
name,
|
||
str(n),
|
||
f"{sum(r['er_99_norm'] for r in items) / n:.4f}",
|
||
f"{sum(r['entropic_rank_norm'] for r in items) / n:.4f}",
|
||
f"{sum(r['top1_ratio'] for r in items) / n:.4f}",
|
||
f"{sum(r['condition_number'] for r in items) / n:.1f}",
|
||
],
|
||
widths,
|
||
)
|
||
)
|
||
|
||
all_er = [
|
||
r["er_99_norm"]
|
||
for vs in matrix_groups.values()
|
||
for r in vs
|
||
if "_norm" not in r or not r.get("is_1d")
|
||
]
|
||
if all_er:
|
||
m = sum(all_er) / len(all_er)
|
||
print(f"\n Overall Mean ER@99: {m:.4f} ({m * 100:.1f}% of dimension)")
|
||
if m > 0.85:
|
||
print(" → HIGH utilization: model near capacity → need more params")
|
||
elif m > 0.5:
|
||
print(" → MODERATE utilization: some headroom left")
|
||
else:
|
||
print(" → LOW utilization: significant unused capacity")
|
||
|
||
|
||
def print_layer_grid(results: dict[str, dict]):
|
||
comps = [
|
||
"attn.q_proj",
|
||
"attn.k_proj",
|
||
"attn.v_proj",
|
||
"attn.o_proj",
|
||
"mlp.up",
|
||
"mlp.gate",
|
||
"mlp.down",
|
||
]
|
||
widths = [6] + [10] * len(comps)
|
||
metric = "er_99_norm"
|
||
|
||
print(f"\n--- Per-Layer Effective Rank (99% energy) ---")
|
||
print(format_header(["Layer"] + comps, widths))
|
||
print("-" * sum(widths))
|
||
|
||
layer_data: dict[int, dict[str, dict]] = {}
|
||
for key, r in results.items():
|
||
parts = key.split(".")
|
||
if parts[0] != "layers":
|
||
continue
|
||
li = int(parts[1])
|
||
sub = parts[2:]
|
||
if sub[0] == "attention":
|
||
cname = f"attn.{sub[1]}"
|
||
elif sub[0] == "mlp":
|
||
cname = f"mlp.{sub[1]}"
|
||
else:
|
||
continue
|
||
layer_data.setdefault(li, {})[cname] = r
|
||
|
||
for li in sorted(layer_data):
|
||
values = [str(li)]
|
||
for c in comps:
|
||
v = layer_data[li].get(c, {}).get(metric, 0)
|
||
values.append(f"{v:.4f}")
|
||
print(format_row(values, widths))
|
||
|
||
|
||
def print_weight_stats(results: dict[str, dict]):
|
||
groups = group_by_component(results)
|
||
widths = [20, 12, 12, 12, 12]
|
||
print(f"\n--- Weight Value Statistics ---")
|
||
print(format_header(["Component", "Mean", "Std", "Min", "Max"], widths))
|
||
print("-" * sum(widths))
|
||
|
||
for name in sorted(groups.keys()):
|
||
items = groups[name]
|
||
means = [r.get("mean", 0) for r in items]
|
||
stds = [r.get("std", 0) for r in items]
|
||
mins = [r.get("min", 0) for r in items]
|
||
maxs = [r.get("max", 0) for r in items]
|
||
g_mean = sum(means) / len(means)
|
||
g_std = sum(stds) / len(stds)
|
||
g_min = min(mins)
|
||
g_max = max(maxs)
|
||
print(
|
||
format_row(
|
||
[
|
||
name,
|
||
f"{g_mean:.6f}",
|
||
f"{g_std:.6f}",
|
||
f"{g_min:.6f}",
|
||
f"{g_max:.6f}",
|
||
],
|
||
widths,
|
||
)
|
||
)
|
||
|
||
|
||
def print_params_summary(results: dict[str, dict]):
|
||
total_2d = sum(
|
||
r["shape"][0] * r["shape"][1] for r in results.values() if not r.get("is_1d")
|
||
)
|
||
total_1d = sum(r["shape"][0] for r in results.values() if r.get("is_1d"))
|
||
print(f"\n Total 2D params: {total_2d:,}")
|
||
print(f" Total 1D params: {total_1d:,}")
|
||
print(f" Total params: {total_2d + total_1d:,}")
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(
|
||
description="SVD effective rank & weight statistics of a model checkpoint."
|
||
)
|
||
parser.add_argument(
|
||
"--ckpt_dir",
|
||
type=str,
|
||
required=True,
|
||
help="Path to checkpoint directory (containing model.safetensors + config.json).",
|
||
)
|
||
parser.add_argument(
|
||
"--compare",
|
||
type=str,
|
||
nargs="*",
|
||
help="Additional checkpoint directories to compare against.",
|
||
)
|
||
parser.add_argument(
|
||
"--no_svd",
|
||
action="store_true",
|
||
help="Skip SVD analysis, only show weight statistics (mean/std/min/max).",
|
||
)
|
||
args = parser.parse_args()
|
||
|
||
def analyze_one(ckpt_dir: str, label: str):
|
||
ckpt_dir = Path(ckpt_dir)
|
||
weights_path = ckpt_dir / "model.safetensors"
|
||
if not weights_path.exists():
|
||
print(f"ERROR: {weights_path} not found")
|
||
return {}
|
||
|
||
meta = {}
|
||
meta_path = ckpt_dir / "meta.json"
|
||
if meta_path.exists():
|
||
with open(meta_path) as f:
|
||
meta = json.load(f)
|
||
|
||
print(f"\n{'=' * 70}")
|
||
print(f" {label}: {ckpt_dir}")
|
||
if meta:
|
||
print(
|
||
f" Iteration: {meta.get('iteration', '?')}, "
|
||
f"Strategy: {meta.get('strategy', '?')}, "
|
||
f"nprocs={meta.get('nprocs', '?')}"
|
||
)
|
||
print(f"{'=' * 70}")
|
||
|
||
print(f"Loading weights...")
|
||
sd = safetensors.torch.load_file(str(weights_path))
|
||
print(f" {len(sd)} keys loaded")
|
||
|
||
weight_keys = [
|
||
k
|
||
for k in sd
|
||
if ".weight" in k and "rotary_embedding" not in k and "freqs_cis" not in k
|
||
]
|
||
|
||
results = {}
|
||
if not args.no_svd:
|
||
print(f"Computing SVD on {len(weight_keys)} tensors...")
|
||
for i, k in enumerate(sorted(weight_keys)):
|
||
print(f" [{i + 1}/{len(weight_keys)}] {k:<60s}", end="\r")
|
||
results[k] = effective_rank_metrics(sd[k])
|
||
print()
|
||
else:
|
||
print(f"Computing stats on {len(weight_keys)} tensors (no SVD)...")
|
||
for i, k in enumerate(sorted(weight_keys)):
|
||
t = sd[k]
|
||
results[k] = {
|
||
"shape": tuple(t.shape),
|
||
"is_1d": t.ndim == 1,
|
||
"mean": t.float().mean().item(),
|
||
"std": t.float().std().item(),
|
||
"min": t.float().min().item(),
|
||
"max": t.float().max().item(),
|
||
}
|
||
|
||
print_params_summary(results)
|
||
if not args.no_svd:
|
||
print_component_summary(
|
||
results, "\n=== SVD Effective Rank by Component ==="
|
||
)
|
||
print_layer_grid(results)
|
||
print_weight_stats(results)
|
||
return results
|
||
|
||
analyze_one(args.ckpt_dir, "Primary")
|
||
|
||
if args.compare:
|
||
for cdir in args.compare:
|
||
analyze_one(cdir, "Compare")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|