819 lines
31 KiB
Python
819 lines
31 KiB
Python
|
"""
|
||
|
|
Usage : python audit.py \
|
||
|
|
--base_model Qwen/Qwen3-1.7B-Base \
|
||
|
|
--distilled_model iamrahulreddy/Quintus \
|
||
|
|
--output_file weight_audit_report.txt \
|
||
|
|
--alpha 0.3
|
||
|
|
"""
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import collections
|
||
|
|
import math
|
||
|
|
import sys
|
||
|
|
import time
|
||
|
|
from datetime import datetime, timezone
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
import torch
|
||
|
|
import torch.nn.functional as F
|
||
|
|
from huggingface_hub import snapshot_download
|
||
|
|
from transformers import AutoConfig, AutoModelForCausalLM
|
||
|
|
|
||
|
|
# Formatting utilities
|
||
|
|
def fmt_num(n: int) -> str:
|
||
|
|
if n >= 1_000_000_000:
|
||
|
|
return f"{n:,} ({n / 1e9:.6f} B)"
|
||
|
|
if n >= 1_000_000:
|
||
|
|
return f"{n:,} ({n / 1e6:.6f} M)"
|
||
|
|
return f"{n:,}"
|
||
|
|
|
||
|
|
def fmt_size(b: int) -> str:
|
||
|
|
if b >= 1 << 30:
|
||
|
|
return f"{b / (1 << 30):.3f} GiB"
|
||
|
|
if b >= 1 << 20:
|
||
|
|
return f"{b / (1 << 20):.3f} MiB"
|
||
|
|
if b >= 1 << 10:
|
||
|
|
return f"{b / (1 << 10):.3f} KiB"
|
||
|
|
return f"{b} B"
|
||
|
|
|
||
|
|
def divider(char: str = "-", width: int = 88) -> str:
|
||
|
|
return char * width
|
||
|
|
|
||
|
|
def section_header(index: int, title: str) -> str:
|
||
|
|
return f"\n[{index:02d}] {title}"
|
||
|
|
|
||
|
|
def sub_header(title: str) -> str:
|
||
|
|
return f"\n -- {title}"
|
||
|
|
|
||
|
|
# Layer classification
|
||
|
|
LAYER_TYPE_MAP = {
|
||
|
|
"embed_tokens": "embedding",
|
||
|
|
"lm_head": "lm_head",
|
||
|
|
"self_attn.q_proj": "attn_q",
|
||
|
|
"self_attn.k_proj": "attn_k",
|
||
|
|
"self_attn.v_proj": "attn_v",
|
||
|
|
"self_attn.o_proj": "attn_o",
|
||
|
|
"self_attn.q_norm": "attn_qnorm",
|
||
|
|
"self_attn.k_norm": "attn_knorm",
|
||
|
|
"mlp.gate_proj": "mlp_gate",
|
||
|
|
"mlp.up_proj": "mlp_up",
|
||
|
|
"mlp.down_proj": "mlp_down",
|
||
|
|
"input_layernorm": "layernorm",
|
||
|
|
"post_attention_layernorm": "layernorm",
|
||
|
|
"model.norm": "final_norm",
|
||
|
|
}
|
||
|
|
|
||
|
|
def classify_layer(name: str) -> str:
|
||
|
|
for pattern, label in LAYER_TYPE_MAP.items():
|
||
|
|
if pattern in name:
|
||
|
|
return label
|
||
|
|
return "other"
|
||
|
|
|
||
|
|
# Tensor statistics
|
||
|
|
def tensor_stats(t: torch.Tensor) -> dict:
|
||
|
|
tf = t.float()
|
||
|
|
flat = tf.view(-1)
|
||
|
|
mean = flat.mean().item()
|
||
|
|
std = flat.std().item()
|
||
|
|
|
||
|
|
sparsity = (flat.abs() < 1e-6).float().mean().item()
|
||
|
|
sat_thresh = flat.abs().max().item() * 0.99
|
||
|
|
saturation = (flat.abs() >= sat_thresh).float().mean().item()
|
||
|
|
kurtosis = (((flat - mean) / std) ** 4).mean().item() - 3.0 if std > 1e-10 else 0.0
|
||
|
|
outlier_r = (flat.abs() > (flat.abs().mean() + 3.0 * std)).float().mean().item()
|
||
|
|
|
||
|
|
row_l2_stats = {}
|
||
|
|
if tf.ndim == 2:
|
||
|
|
row_norms = tf.norm(2, dim=1)
|
||
|
|
row_l2_stats = {
|
||
|
|
"row_l2_mean": row_norms.mean().item(),
|
||
|
|
"row_l2_std": row_norms.std().item(),
|
||
|
|
"row_l2_min": row_norms.min().item(),
|
||
|
|
"row_l2_max": row_norms.max().item(),
|
||
|
|
"dead_rows": int((row_norms < 1e-6).sum().item()),
|
||
|
|
}
|
||
|
|
|
||
|
|
return {
|
||
|
|
"shape": list(tf.shape),
|
||
|
|
"numel": flat.numel(),
|
||
|
|
"dtype": str(t.dtype),
|
||
|
|
"mean": mean,
|
||
|
|
"std": std,
|
||
|
|
"min": flat.min().item(),
|
||
|
|
"max": flat.max().item(),
|
||
|
|
"abs_mean": flat.abs().mean().item(),
|
||
|
|
"l2_norm": flat.norm(2).item(),
|
||
|
|
"l1_norm": flat.norm(1).item(),
|
||
|
|
"sparsity": sparsity,
|
||
|
|
"saturation": saturation,
|
||
|
|
"kurtosis": kurtosis,
|
||
|
|
"outlier_ratio": outlier_r,
|
||
|
|
**row_l2_stats,
|
||
|
|
}
|
||
|
|
|
||
|
|
# Divergence between two tensors
|
||
|
|
def tensor_divergence(t_base: torch.Tensor, t_dist: torch.Tensor, chunk_size: int = 10_000_000) -> dict:
|
||
|
|
a_flat = t_base.detach().view(-1)
|
||
|
|
b_flat = t_dist.detach().view(-1)
|
||
|
|
n_elements = a_flat.numel()
|
||
|
|
|
||
|
|
# Running accumulators in float64 (on CPU/Python) to prevent memory spikes
|
||
|
|
dot_prod = 0.0
|
||
|
|
a_sq_sum = 0.0
|
||
|
|
b_sq_sum = 0.0
|
||
|
|
|
||
|
|
# Delta statistics
|
||
|
|
max_delta = 0.0
|
||
|
|
sum_delta = 0.0
|
||
|
|
l2_delta_sq = 0.0
|
||
|
|
sum_abs_a = 0.0
|
||
|
|
|
||
|
|
# Process in chunks to keep memory footprint extremely small (~80MB peak per chunk)
|
||
|
|
for i in range(0, n_elements, chunk_size):
|
||
|
|
a_chunk = a_flat[i : i + chunk_size].to(torch.float64)
|
||
|
|
b_chunk = b_flat[i : i + chunk_size].to(torch.float64)
|
||
|
|
|
||
|
|
# Accumulate dot product and norms
|
||
|
|
dot_prod += torch.dot(a_chunk, b_chunk).item()
|
||
|
|
a_sq_sum += torch.dot(a_chunk, a_chunk).item()
|
||
|
|
b_sq_sum += torch.dot(b_chunk, b_chunk).item()
|
||
|
|
|
||
|
|
# Accumulate delta stats
|
||
|
|
delta_chunk = (b_chunk - a_chunk).abs()
|
||
|
|
max_delta = max(max_delta, delta_chunk.max().item())
|
||
|
|
sum_delta += delta_chunk.sum().item()
|
||
|
|
l2_delta_sq += torch.dot(delta_chunk, delta_chunk).item()
|
||
|
|
sum_abs_a += a_chunk.abs().sum().item()
|
||
|
|
|
||
|
|
# Final metrics
|
||
|
|
a_norm = math.sqrt(a_sq_sum)
|
||
|
|
b_norm = math.sqrt(b_sq_sum)
|
||
|
|
if a_norm > 0 and b_norm > 0:
|
||
|
|
cos_sim_raw = dot_prod / (a_norm * b_norm)
|
||
|
|
else:
|
||
|
|
cos_sim_raw = 0.0
|
||
|
|
|
||
|
|
cos_sim = max(-1.0, min(1.0, cos_sim_raw))
|
||
|
|
|
||
|
|
rel_err = sum_delta / (sum_abs_a + 1e-12)
|
||
|
|
base_l2 = a_norm
|
||
|
|
delta_l2 = math.sqrt(l2_delta_sq)
|
||
|
|
snr_db = 20.0 * math.log10(base_l2 / (delta_l2 + 1e-12)) if base_l2 > 0 else 0.0
|
||
|
|
|
||
|
|
# Standard deviation of delta
|
||
|
|
mean_delta = sum_delta / n_elements
|
||
|
|
mean_delta_sq = l2_delta_sq / n_elements
|
||
|
|
var_delta = max(0.0, mean_delta_sq - mean_delta**2)
|
||
|
|
std_delta = math.sqrt(var_delta)
|
||
|
|
|
||
|
|
return {
|
||
|
|
"max_delta": max_delta,
|
||
|
|
"mean_delta": mean_delta,
|
||
|
|
"std_delta": std_delta,
|
||
|
|
"l2_delta": delta_l2,
|
||
|
|
"cos_sim": cos_sim,
|
||
|
|
"cos_sim_raw": cos_sim_raw,
|
||
|
|
"rel_err": rel_err,
|
||
|
|
"snr_db": snr_db,
|
||
|
|
"changed": max_delta > 1e-7,
|
||
|
|
}
|
||
|
|
|
||
|
|
# Isotropy
|
||
|
|
def isotropy_score(t: torch.Tensor, n_samples: int = 2048) -> float:
|
||
|
|
"""
|
||
|
|
Average pairwise cosine similarity of randomly sampled row vectors.
|
||
|
|
Near 0 = isotropic (healthy). Near 1 = collapsed representations.
|
||
|
|
Only valid for 2D tensors with >= 2 rows.
|
||
|
|
"""
|
||
|
|
if t.ndim != 2 or t.shape[0] < 2:
|
||
|
|
return float("nan")
|
||
|
|
tf = t.float()
|
||
|
|
n = min(t.shape[0], n_samples)
|
||
|
|
# Add deterministic seed for isotropy sampling
|
||
|
|
gen = torch.Generator().manual_seed(42)
|
||
|
|
idx = torch.randperm(t.shape[0], generator=gen)[:n].to(t.device)
|
||
|
|
rows = tf[idx]
|
||
|
|
norms = rows.norm(2, dim=1, keepdim=True).clamp(min=1e-12)
|
||
|
|
normed = rows / norms
|
||
|
|
sim = normed @ normed.T
|
||
|
|
mask = ~torch.eye(n, dtype=torch.bool)
|
||
|
|
return sim[mask].mean().item()
|
||
|
|
|
||
|
|
# Config helpers
|
||
|
|
def config_architecture_lines(config, label: str, model_id: str) -> list[str]:
|
||
|
|
cfg = config.to_dict()
|
||
|
|
n_q = cfg.get("num_attention_heads", 1)
|
||
|
|
n_kv = cfg.get("num_key_value_heads", n_q)
|
||
|
|
h = cfg.get("hidden_size", 0)
|
||
|
|
head_dim = h // n_q if n_q else 0
|
||
|
|
gqa = n_q // n_kv if n_kv else 1
|
||
|
|
|
||
|
|
return [
|
||
|
|
f" label : {label} ({model_id})",
|
||
|
|
f" model_type : {cfg.get('model_type', 'unknown')}",
|
||
|
|
f" architecture : {cfg.get('architectures', ['unknown'])[0]}",
|
||
|
|
"",
|
||
|
|
" Vocabulary",
|
||
|
|
f" vocab_size : {cfg.get('vocab_size', 'N/A'):,}",
|
||
|
|
f" bos / eos / pad : {cfg.get('bos_token_id')} / {cfg.get('eos_token_id')} / {cfg.get('pad_token_id')}",
|
||
|
|
"",
|
||
|
|
" Positional encoding",
|
||
|
|
f" max_position_embeddings: {cfg.get('max_position_embeddings', 'N/A'):,}",
|
||
|
|
f" rope_theta : {cfg.get('rope_theta', 'N/A')}",
|
||
|
|
f" rope_scaling : {cfg.get('rope_scaling', 'None')}",
|
||
|
|
"",
|
||
|
|
" Transformer dimensions",
|
||
|
|
f" hidden_size : {h}",
|
||
|
|
f" num_hidden_layers : {cfg.get('num_hidden_layers', 'N/A')}",
|
||
|
|
f" intermediate_size : {cfg.get('intermediate_size', 'N/A')}",
|
||
|
|
"",
|
||
|
|
" Attention",
|
||
|
|
f" num_attention_heads : {n_q}",
|
||
|
|
f" num_key_value_heads : {n_kv}",
|
||
|
|
f" head_dim : {head_dim}",
|
||
|
|
f" GQA ratio : {gqa}:1",
|
||
|
|
f" attention_bias : {cfg.get('attention_bias', False)}",
|
||
|
|
f" use_qk_norm : {cfg.get('use_qk_norm', False) or 'qwen3' in model_id.lower() or 'qwen3' in cfg.get('model_type', '').lower()}",
|
||
|
|
f" sliding_window : {cfg.get('sliding_window', 'None')}",
|
||
|
|
"",
|
||
|
|
" Feed-forward",
|
||
|
|
f" hidden_act : {cfg.get('hidden_act', 'silu')}",
|
||
|
|
f" mlp_bias : {cfg.get('mlp_bias', False)}",
|
||
|
|
"",
|
||
|
|
" Misc",
|
||
|
|
f" rms_norm_eps : {cfg.get('rms_norm_eps', 1e-6)}",
|
||
|
|
f" tie_word_embeddings : {cfg.get('tie_word_embeddings', True)}",
|
||
|
|
f" use_cache : {cfg.get('use_cache', True)}",
|
||
|
|
f" torch_dtype : {cfg.get('torch_dtype', 'float32')}",
|
||
|
|
f" initializer_range : {cfg.get('initializer_range', 'N/A')}",
|
||
|
|
]
|
||
|
|
|
||
|
|
def get_params_info(config, model_id: str = "") -> dict:
|
||
|
|
h = config.hidden_size
|
||
|
|
l = config.num_hidden_layers
|
||
|
|
v = config.vocab_size
|
||
|
|
embed = v * h
|
||
|
|
tie = getattr(config, "tie_word_embeddings", True)
|
||
|
|
n_q = config.num_attention_heads
|
||
|
|
n_kv = getattr(config, "num_key_value_heads", n_q)
|
||
|
|
head_dim = h // n_q
|
||
|
|
qkv_proj = (n_q + 2 * n_kv) * head_dim * h
|
||
|
|
o_proj = h * h
|
||
|
|
use_qk_norm = (
|
||
|
|
getattr(config, "use_qk_norm", False) or
|
||
|
|
"qwen3" in model_id.lower() or
|
||
|
|
"qwen3" in getattr(config, "model_type", "").lower()
|
||
|
|
)
|
||
|
|
qk_norm = 2 * head_dim if use_qk_norm else 0
|
||
|
|
mlp = 3 * h * config.intermediate_size
|
||
|
|
norms = 2 * h
|
||
|
|
per_layer = qkv_proj + o_proj + qk_norm + mlp + norms
|
||
|
|
total_layers = l * per_layer
|
||
|
|
lm_head = 0 if tie else embed
|
||
|
|
unique = embed + lm_head + total_layers + h # +h for final norm
|
||
|
|
return {
|
||
|
|
"raw": unique + (embed if tie else 0),
|
||
|
|
"embed": embed,
|
||
|
|
"lm_head": embed,
|
||
|
|
"tied": tie,
|
||
|
|
"unique": unique,
|
||
|
|
"non_embed": total_layers + h,
|
||
|
|
"per_layer": per_layer,
|
||
|
|
}
|
||
|
|
|
||
|
|
def param_lines(config, p: dict, label: str) -> list[str]:
|
||
|
|
return [
|
||
|
|
f" {label}",
|
||
|
|
f" raw (all named) : {fmt_num(p['raw'])}",
|
||
|
|
f" embedding : {fmt_num(p['embed'])}",
|
||
|
|
f" lm_head : {fmt_num(p['lm_head'])}",
|
||
|
|
f" tied : {p['tied']}",
|
||
|
|
f" unique (deduped) : {fmt_num(p['unique'])}",
|
||
|
|
f" non-embedding : {fmt_num(p['non_embed'])}",
|
||
|
|
f" per layer (approx) : {p['per_layer']:,}",
|
||
|
|
]
|
||
|
|
|
||
|
|
# Main
|
||
|
|
def main():
|
||
|
|
parser = argparse.ArgumentParser(description="Quintus Deep Weight Audit")
|
||
|
|
parser.add_argument("--base_model", type=str, default="Qwen/Qwen3-1.7B-Base")
|
||
|
|
parser.add_argument("--distilled_model", type=str, default="iamrahulreddy/Quintus")
|
||
|
|
parser.add_argument("--output_file", type=str, default="weight_audit_report.txt")
|
||
|
|
parser.add_argument("--alpha", type=float, default=0.3)
|
||
|
|
parser.add_argument("--isotropy_samples", type=int, default=2048)
|
||
|
|
parser.add_argument("--trust_remote_code", action="store_true", help="Allow custom code from model repositories.")
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
# Determine compute device
|
||
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
|
|
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
|
||
|
|
|
||
|
|
utc_ts = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
|
||
|
|
loc_ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S local")
|
||
|
|
|
||
|
|
R: list[str] = []
|
||
|
|
|
||
|
|
def log(line: str = ""):
|
||
|
|
print(line)
|
||
|
|
R.append(line)
|
||
|
|
|
||
|
|
def loglines(lines: list[str]):
|
||
|
|
for ln in lines:
|
||
|
|
log(ln)
|
||
|
|
|
||
|
|
# Header
|
||
|
|
loglines([
|
||
|
|
divider("="),
|
||
|
|
" QUINTUS WEIGHT AUDIT",
|
||
|
|
divider("="),
|
||
|
|
f" {utc_ts} ({loc_ts})",
|
||
|
|
f" base model : {args.base_model}",
|
||
|
|
f" distilled model : {args.distilled_model}",
|
||
|
|
f" alpha : {args.alpha}",
|
||
|
|
f" device : {device} | dtype: {dtype}",
|
||
|
|
f" python : {sys.version.split()[0]} | torch: {torch.__version__}",
|
||
|
|
divider("="),
|
||
|
|
])
|
||
|
|
|
||
|
|
# [01] Resolve checkpoints
|
||
|
|
log(section_header(1, "Resolve checkpoints"))
|
||
|
|
|
||
|
|
# Resolve base model commit hash (pin and report base commit)
|
||
|
|
base_commit = "local"
|
||
|
|
if not Path(args.base_model).exists():
|
||
|
|
try:
|
||
|
|
base_local_dir = Path(snapshot_download(repo_id=args.base_model))
|
||
|
|
base_commit = base_local_dir.name
|
||
|
|
except Exception:
|
||
|
|
base_commit = "unknown"
|
||
|
|
|
||
|
|
dist_commit = "local"
|
||
|
|
if not Path(args.distilled_model).exists():
|
||
|
|
log(f" Downloading '{args.distilled_model}' from HuggingFace Hub...")
|
||
|
|
t0 = time.time()
|
||
|
|
try:
|
||
|
|
local_dir = snapshot_download(repo_id=args.distilled_model)
|
||
|
|
distilled_path = Path(local_dir)
|
||
|
|
dist_commit = distilled_path.name
|
||
|
|
except Exception as e:
|
||
|
|
log(f" ERROR: {e}")
|
||
|
|
sys.exit(1)
|
||
|
|
log(f" Done in {time.time() - t0:.1f}s")
|
||
|
|
else:
|
||
|
|
distilled_path = Path(args.distilled_model)
|
||
|
|
if "snapshots" in distilled_path.parts:
|
||
|
|
dist_commit = distilled_path.name
|
||
|
|
|
||
|
|
# Redact absolute local HF cache paths for sharing
|
||
|
|
redacted_root = "<HF_CACHE_DIR>/snapshots"
|
||
|
|
log(f" base model commit : {base_commit}")
|
||
|
|
log(f" distilled commit : {dist_commit}")
|
||
|
|
log(f" snapshot root : {redacted_root}")
|
||
|
|
|
||
|
|
if not (distilled_path / "config.json").exists():
|
||
|
|
log(" ERROR: config.json missing from checkpoint directory.")
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
files = sorted(f for f in distilled_path.iterdir() if f.is_file())
|
||
|
|
total_ckpt_bytes = sum(f.stat().st_size for f in files)
|
||
|
|
log("")
|
||
|
|
log(f" {'Filename':<52} {'Size':>12} Modified")
|
||
|
|
for f in files:
|
||
|
|
mtime = datetime.fromtimestamp(f.stat().st_mtime).strftime("%Y-%m-%d %H:%M")
|
||
|
|
log(f" {f.name:<52} {fmt_size(f.stat().st_size):>12} {mtime}")
|
||
|
|
log(f" {'total':<52} {fmt_size(total_ckpt_bytes):>12}")
|
||
|
|
|
||
|
|
# [02] Architecture configuration
|
||
|
|
log(section_header(2, "Architecture configuration"))
|
||
|
|
|
||
|
|
log(" Loading base config...")
|
||
|
|
try:
|
||
|
|
base_config = AutoConfig.from_pretrained(args.base_model, trust_remote_code=args.trust_remote_code)
|
||
|
|
except Exception as e:
|
||
|
|
log(f" ERROR: {e}"); sys.exit(1)
|
||
|
|
|
||
|
|
log(" Loading distilled config...")
|
||
|
|
try:
|
||
|
|
distilled_config = AutoConfig.from_pretrained(str(distilled_path), trust_remote_code=args.trust_remote_code)
|
||
|
|
except Exception as e:
|
||
|
|
log(f" ERROR: {e}"); sys.exit(1)
|
||
|
|
|
||
|
|
log(sub_header("Base"))
|
||
|
|
loglines(config_architecture_lines(base_config, "base", args.base_model))
|
||
|
|
|
||
|
|
log(sub_header("Distilled"))
|
||
|
|
loglines(config_architecture_lines(distilled_config, "distilled", args.distilled_model))
|
||
|
|
|
||
|
|
log(sub_header("Config diff (ignoring: _name_or_path, transformers_version)"))
|
||
|
|
ignore_keys = {"_name_or_path", "transformers_version"}
|
||
|
|
base_dict = base_config.to_dict()
|
||
|
|
dist_dict = distilled_config.to_dict()
|
||
|
|
config_diffs = [
|
||
|
|
(k, base_dict.get(k), dist_dict.get(k))
|
||
|
|
for k in sorted(set(base_dict) | set(dist_dict))
|
||
|
|
if k not in ignore_keys and base_dict.get(k) != dist_dict.get(k)
|
||
|
|
]
|
||
|
|
if not config_diffs:
|
||
|
|
log(" No differences — configs identical (expected for same-architecture KD).")
|
||
|
|
else:
|
||
|
|
log(f" {'Key':<40} {'Base':>28} Distilled")
|
||
|
|
for k, vb, vd in config_diffs:
|
||
|
|
log(f" {k:<40} {str(vb):>28} {vd}")
|
||
|
|
|
||
|
|
# [03] Parameter accounting
|
||
|
|
log(section_header(3, "Parameter accounting"))
|
||
|
|
|
||
|
|
base_params = get_params_info(base_config, args.base_model)
|
||
|
|
dist_params = get_params_info(distilled_config, args.distilled_model)
|
||
|
|
|
||
|
|
log(sub_header("Base"))
|
||
|
|
loglines(param_lines(base_config, base_params, "base"))
|
||
|
|
|
||
|
|
log(sub_header("Distilled"))
|
||
|
|
loglines(param_lines(distilled_config, dist_params, "distilled"))
|
||
|
|
|
||
|
|
log(sub_header("Delta"))
|
||
|
|
du = dist_params["unique"] - base_params["unique"]
|
||
|
|
log(f" unique param delta : {du:+,} ({du / base_params['unique'] * 100:+.4f} %)")
|
||
|
|
log(f" non-embed param delta : {dist_params['non_embed'] - base_params['non_embed']:+,}")
|
||
|
|
|
||
|
|
# [04] Load weights onto GPU
|
||
|
|
log(section_header(4, "Load weights"))
|
||
|
|
log(f" device: {device} | dtype: {dtype}")
|
||
|
|
|
||
|
|
load_kwargs = dict(dtype=dtype, device_map=device, trust_remote_code=args.trust_remote_code)
|
||
|
|
|
||
|
|
log(f" Loading base model : {args.base_model}")
|
||
|
|
t0 = time.time()
|
||
|
|
base_model = AutoModelForCausalLM.from_pretrained(args.base_model, **load_kwargs)
|
||
|
|
log(f" Done in {time.time() - t0:.1f}s")
|
||
|
|
|
||
|
|
log(f" Loading distilled : {args.distilled_model}")
|
||
|
|
t0 = time.time()
|
||
|
|
distilled_model = AutoModelForCausalLM.from_pretrained(str(distilled_path), **load_kwargs)
|
||
|
|
log(f" Done in {time.time() - t0:.1f}s")
|
||
|
|
|
||
|
|
base_sd = base_model.state_dict()
|
||
|
|
dist_sd = distilled_model.state_dict()
|
||
|
|
|
||
|
|
log(f" base tensors : {len(base_sd)}")
|
||
|
|
log(f" distilled tensors : {len(dist_sd)}")
|
||
|
|
|
||
|
|
only_base = set(base_sd) - set(dist_sd)
|
||
|
|
only_dist = set(dist_sd) - set(base_sd)
|
||
|
|
if only_base:
|
||
|
|
log(f" keys only in base : {sorted(only_base)[:5]} ...")
|
||
|
|
if only_dist:
|
||
|
|
log(f" keys only in distilled: {sorted(only_dist)[:5]} ...")
|
||
|
|
|
||
|
|
tied = torch.equal(
|
||
|
|
base_sd["model.embed_tokens.weight"],
|
||
|
|
base_sd.get("lm_head.weight", base_sd["model.embed_tokens.weight"]),
|
||
|
|
)
|
||
|
|
log(f" weight tying confirmed (embed == lm_head): {tied}")
|
||
|
|
|
||
|
|
def sd_bytes(sd):
|
||
|
|
return sum(t.numel() * t.element_size() for t in sd.values())
|
||
|
|
|
||
|
|
log(f" base weight memory : {fmt_size(sd_bytes(base_sd))}")
|
||
|
|
log(f" distilled memory : {fmt_size(sd_bytes(dist_sd))}")
|
||
|
|
|
||
|
|
# All subsequent tensor ops: move to CPU float32 only during computation,
|
||
|
|
# keep storage on GPU in bfloat16.
|
||
|
|
all_names = list(dist_sd.keys())
|
||
|
|
|
||
|
|
# [05] Full per-tensor statistics (distilled)
|
||
|
|
log(section_header(5, "Per-tensor weight statistics (distilled)"))
|
||
|
|
|
||
|
|
col = (
|
||
|
|
f" {'Layer':<68} {'Shape':<22} {'Mean':>8} {'Std':>8} "
|
||
|
|
f"{'Min':>8} {'Max':>8} {'Sparse':>7} {'KurtD':>7} "
|
||
|
|
f"{'OutlR':>7} {'RowL2':>8} {'DeadR':>6}"
|
||
|
|
)
|
||
|
|
log(col)
|
||
|
|
log(f" {divider('-', 170)}")
|
||
|
|
|
||
|
|
# Helper to calculate kurtosis statistics for base comparison
|
||
|
|
all_stats: dict[str, dict] = {}
|
||
|
|
type_buckets: dict[str, list[str]] = collections.defaultdict(list)
|
||
|
|
|
||
|
|
for name in all_names:
|
||
|
|
# Move to CPU float32 for stats only
|
||
|
|
t = dist_sd[name].cpu()
|
||
|
|
st = tensor_stats(t)
|
||
|
|
|
||
|
|
# Calculate base model kurtosis if present
|
||
|
|
if name in base_sd:
|
||
|
|
t_base = base_sd[name].cpu()
|
||
|
|
st_base = tensor_stats(t_base)
|
||
|
|
kurt_base = st_base["kurtosis"]
|
||
|
|
else:
|
||
|
|
kurt_base = 0.0
|
||
|
|
|
||
|
|
st["kurtosis_base"] = kurt_base
|
||
|
|
st["kurtosis_delta"] = st["kurtosis"] - kurt_base
|
||
|
|
|
||
|
|
all_stats[name] = st
|
||
|
|
type_buckets[classify_layer(name)].append(name)
|
||
|
|
|
||
|
|
rl2 = st.get("row_l2_mean", float("nan"))
|
||
|
|
dead = st.get("dead_rows", float("nan"))
|
||
|
|
log(
|
||
|
|
f" {name:<68} {str(st['shape']):<22} "
|
||
|
|
f"{st['mean']:8.4f} {st['std']:8.4f} "
|
||
|
|
f"{st['min']:8.4f} {st['max']:8.4f} "
|
||
|
|
f"{st['sparsity']:7.4f} {st['kurtosis_delta']:7.2f} "
|
||
|
|
f"{st['outlier_ratio']:7.4f} "
|
||
|
|
f"{rl2:8.4f} "
|
||
|
|
f"{str(int(dead)) if not math.isnan(dead) else 'N/A':>6}"
|
||
|
|
)
|
||
|
|
|
||
|
|
# [06] Layer-type aggregation (distilled)
|
||
|
|
log(section_header(6, "Layer-type aggregated statistics (distilled)"))
|
||
|
|
|
||
|
|
log(f" {'Type':<18} {'Count':>5} {'Params':>16} {'AvgMean':>9} {'AvgStd':>9} {'AvgSparse':>10} {'AvgKurtD':>9}")
|
||
|
|
log(f" {divider('-', 82)}")
|
||
|
|
|
||
|
|
for ltype in sorted(type_buckets):
|
||
|
|
names = type_buckets[ltype]
|
||
|
|
n = len(names)
|
||
|
|
params = sum(all_stats[x]["numel"] for x in names)
|
||
|
|
log(
|
||
|
|
f" {ltype:<18} {n:>5} {params:>16,} "
|
||
|
|
f"{sum(all_stats[x]['mean'] for x in names)/n:>9.5f} "
|
||
|
|
f"{sum(all_stats[x]['std'] for x in names)/n:>9.5f} "
|
||
|
|
f"{sum(all_stats[x]['sparsity'] for x in names)/n:>10.5f} "
|
||
|
|
f"{sum(all_stats[x]['kurtosis_delta'] for x in names)/n:>9.3f}"
|
||
|
|
)
|
||
|
|
|
||
|
|
# [07] Per-transformer-block breakdown (distilled)
|
||
|
|
log(section_header(7, "Per-transformer-block breakdown (distilled)"))
|
||
|
|
|
||
|
|
n_layers = distilled_config.num_hidden_layers
|
||
|
|
sublayer_order = [
|
||
|
|
"input_layernorm", "self_attn.q_proj", "self_attn.k_proj",
|
||
|
|
"self_attn.v_proj", "self_attn.o_proj", "self_attn.q_norm",
|
||
|
|
"self_attn.k_norm", "post_attention_layernorm",
|
||
|
|
"mlp.gate_proj", "mlp.up_proj", "mlp.down_proj",
|
||
|
|
]
|
||
|
|
|
||
|
|
log(f" {'Blk':>4} {'Sublayer':<35} {'Shape':<22} {'L2':>9} {'AbsMn':>9} {'Std':>9} {'Sparse':>8} {'RowL2':>9}")
|
||
|
|
log(f" {divider('-', 115)}")
|
||
|
|
|
||
|
|
for blk in range(n_layers):
|
||
|
|
prefix = f"model.layers.{blk}."
|
||
|
|
for sub in sublayer_order:
|
||
|
|
nm = prefix + sub + ".weight"
|
||
|
|
if nm not in dist_sd:
|
||
|
|
continue
|
||
|
|
st = all_stats[nm]
|
||
|
|
rl2 = st.get("row_l2_mean", float("nan"))
|
||
|
|
log(
|
||
|
|
f" {blk:>4} {sub:<35} {str(st['shape']):<22} "
|
||
|
|
f"{st['l2_norm']:>9.3f} {st['abs_mean']:>9.5f} "
|
||
|
|
f"{st['std']:>9.5f} {st['sparsity']:>8.5f} {rl2:>9.5f}"
|
||
|
|
)
|
||
|
|
log("")
|
||
|
|
|
||
|
|
# [08] Isotropy analysis (distilled)
|
||
|
|
log(section_header(8, "Isotropy analysis (distilled, 2D tensors only)"))
|
||
|
|
log(f" Sampling up to {args.isotropy_samples} rows per layer.")
|
||
|
|
log(f" Score near 0 = isotropic (healthy). Score near 1 = representation collapse.")
|
||
|
|
log("")
|
||
|
|
log(f" {'Layer':<68} {'Shape':<20} {'Score':>10}")
|
||
|
|
log(f" {divider('-', 102)}")
|
||
|
|
|
||
|
|
iso_scores: dict[str, float] = {}
|
||
|
|
for name in all_names:
|
||
|
|
t = dist_sd[name].cpu()
|
||
|
|
iso = isotropy_score(t, n_samples=args.isotropy_samples)
|
||
|
|
iso_scores[name] = iso
|
||
|
|
if not math.isnan(iso):
|
||
|
|
log(f" {name:<68} {str(all_stats[name]['shape']):<20} {iso:>10.6f}")
|
||
|
|
|
||
|
|
valid_iso = [v for v in iso_scores.values() if not math.isnan(v)]
|
||
|
|
if valid_iso:
|
||
|
|
log("")
|
||
|
|
log(f" Global (across {len(valid_iso)} 2D layers)")
|
||
|
|
log(f" mean : {sum(valid_iso)/len(valid_iso):.6f}")
|
||
|
|
log(f" min : {min(valid_iso):.6f}")
|
||
|
|
log(f" max : {max(valid_iso):.6f}")
|
||
|
|
|
||
|
|
# [09] Base vs distilled divergence — all shared layers
|
||
|
|
log(section_header(9, "Base vs distilled divergence (all shared layers)"))
|
||
|
|
|
||
|
|
shared = sorted(set(base_sd) & set(dist_sd))
|
||
|
|
all_div: dict[str, dict] = {}
|
||
|
|
changed = []
|
||
|
|
unchanged = []
|
||
|
|
|
||
|
|
log(f" Shared tensors: {len(shared)}")
|
||
|
|
log("")
|
||
|
|
log(
|
||
|
|
f" {'Layer':<68} {'MaxDelta':>9} {'MeanDelta':>10} "
|
||
|
|
f"{'L2Delta':>9} {'CosSim':>8} {'RelErr':>8} {'SNR_dB':>7} {'Chg':>4}"
|
||
|
|
)
|
||
|
|
log(f" {divider('-', 135)}")
|
||
|
|
|
||
|
|
for name in shared:
|
||
|
|
b = base_sd[name]
|
||
|
|
d = dist_sd[name]
|
||
|
|
dv = tensor_divergence(b, d)
|
||
|
|
all_div[name] = dv
|
||
|
|
(changed if dv["changed"] else unchanged).append(name)
|
||
|
|
log(
|
||
|
|
f" {name:<68} "
|
||
|
|
f"{dv['max_delta']:>9.5f} {dv['mean_delta']:>10.6f} "
|
||
|
|
f"{dv['l2_delta']:>9.4f} {dv['cos_sim']:>8.5f} "
|
||
|
|
f"{dv['rel_err']:>8.5f} {dv['snr_db']:>7.2f} "
|
||
|
|
f"{'Y' if dv['changed'] else 'N':>4}"
|
||
|
|
)
|
||
|
|
|
||
|
|
log("")
|
||
|
|
log(f" Changed : {len(changed)} / {len(shared)}")
|
||
|
|
log(f" Unchanged: {len(unchanged)} / {len(shared)}")
|
||
|
|
if unchanged:
|
||
|
|
log(f" Unchanged (first 10): {unchanged[:10]}")
|
||
|
|
log("\n Note: Unchanged tensors are primarily normalization layers (input_layernorm, q_norm, k_norm, model.norm).")
|
||
|
|
log(" This demonstrates that the SFT/KD process modified the primary semantic projection weights")
|
||
|
|
log(" (attention and MLP projections) while preserving basic layer scaling characteristics.")
|
||
|
|
|
||
|
|
# [10] Cosine similarity distribution histogram
|
||
|
|
log(section_header(10, "Cosine similarity distribution histogram"))
|
||
|
|
|
||
|
|
cos_vals = [all_div[n]["cos_sim_raw"] for n in shared]
|
||
|
|
bins = [
|
||
|
|
(float('-inf'), 0.900),
|
||
|
|
(0.900, 0.990),
|
||
|
|
(0.990, 0.999),
|
||
|
|
(0.999, 0.9999),
|
||
|
|
(0.9999, 0.99999),
|
||
|
|
(0.99999, 1.00001),
|
||
|
|
(1.00001, 1.001),
|
||
|
|
(1.001, float('inf'))
|
||
|
|
]
|
||
|
|
|
||
|
|
def fmt_bnd(v: float) -> str:
|
||
|
|
if v == float('-inf'):
|
||
|
|
return "-inf"
|
||
|
|
if v == float('inf'):
|
||
|
|
return "inf"
|
||
|
|
return f"{v:7.5f}"
|
||
|
|
|
||
|
|
counts = []
|
||
|
|
for lo, hi in bins:
|
||
|
|
cnt = sum(1 for v in cos_vals if lo <= v < hi)
|
||
|
|
counts.append(cnt)
|
||
|
|
max_cnt = max(counts) if counts else 0
|
||
|
|
max_bar_width = 40
|
||
|
|
|
||
|
|
log(f" {'Range':<22} {'Count':>6} Histogram")
|
||
|
|
for (lo, hi), cnt in zip(bins, counts):
|
||
|
|
bar_len = int(round((cnt / max_cnt) * max_bar_width)) if max_cnt > 0 and cnt > 0 else 0
|
||
|
|
label = f"[{fmt_bnd(lo):>8}, {fmt_bnd(hi):>8})"
|
||
|
|
log(f" {label:<22} {cnt:>6} {'#' * bar_len}")
|
||
|
|
|
||
|
|
# [11] Attention geometry per block
|
||
|
|
log(section_header(11, "Attention geometry per transformer block"))
|
||
|
|
|
||
|
|
n_q = distilled_config.num_attention_heads
|
||
|
|
n_kv = getattr(distilled_config, "num_key_value_heads", n_q)
|
||
|
|
head_dim = distilled_config.hidden_size // n_q
|
||
|
|
|
||
|
|
log(f" Query heads: {n_q} | KV heads: {n_kv} | head_dim: {head_dim} | GQA: {n_q//n_kv}:1")
|
||
|
|
log("")
|
||
|
|
log(
|
||
|
|
f" {'Blk':>4} {'Q shape':<20} {'K shape':<20} {'V shape':<20} {'O shape':<20} "
|
||
|
|
f"{'Q L2':>8} {'K L2':>8} {'V L2':>8} {'O L2':>8}"
|
||
|
|
)
|
||
|
|
log(f" {divider('-', 130)}")
|
||
|
|
|
||
|
|
for blk in range(n_layers):
|
||
|
|
p = f"model.layers.{blk}.self_attn."
|
||
|
|
def attn(key):
|
||
|
|
nm = p + key + ".weight"
|
||
|
|
if nm in dist_sd:
|
||
|
|
st = all_stats[nm]
|
||
|
|
return str(st["shape"]), st["l2_norm"]
|
||
|
|
return "N/A", float("nan")
|
||
|
|
|
||
|
|
qs, ql = attn("q_proj")
|
||
|
|
ks, kl = attn("k_proj")
|
||
|
|
vs, vl = attn("v_proj")
|
||
|
|
os_, ol = attn("o_proj")
|
||
|
|
log(
|
||
|
|
f" {blk:>4} {qs:<20} {ks:<20} {vs:<20} {os_:<20} "
|
||
|
|
f"{ql:>8.3f} {kl:>8.3f} {vl:>8.3f} {ol:>8.3f}"
|
||
|
|
)
|
||
|
|
|
||
|
|
# [12] MLP geometry per block
|
||
|
|
log(section_header(12, "MLP feed-forward geometry per transformer block"))
|
||
|
|
|
||
|
|
log(f" intermediate_size: {distilled_config.intermediate_size} | activation: {getattr(distilled_config, 'hidden_act', 'silu')}")
|
||
|
|
log("")
|
||
|
|
log(
|
||
|
|
f" {'Blk':>4} {'Gate shape':<22} {'Up shape':<22} {'Down shape':<22} "
|
||
|
|
f"{'Gate L2':>8} {'Up L2':>8} {'Down L2':>9} "
|
||
|
|
f"{'GateSp':>8} {'UpSp':>8} {'DnSp':>8}"
|
||
|
|
)
|
||
|
|
log(f" {divider('-', 135)}")
|
||
|
|
|
||
|
|
for blk in range(n_layers):
|
||
|
|
p = f"model.layers.{blk}.mlp."
|
||
|
|
def mlp(key):
|
||
|
|
nm = p + key + ".weight"
|
||
|
|
if nm in dist_sd:
|
||
|
|
st = all_stats[nm]
|
||
|
|
return str(st["shape"]), st["l2_norm"], st["sparsity"]
|
||
|
|
return "N/A", float("nan"), float("nan")
|
||
|
|
|
||
|
|
gs, gl, gsp = mlp("gate_proj")
|
||
|
|
us, ul, usp = mlp("up_proj")
|
||
|
|
ds, dl, dsp = mlp("down_proj")
|
||
|
|
log(
|
||
|
|
f" {blk:>4} {gs:<22} {us:<22} {ds:<22} "
|
||
|
|
f"{gl:>8.3f} {ul:>8.3f} {dl:>9.3f} "
|
||
|
|
f"{gsp:>8.5f} {usp:>8.5f} {dsp:>8.5f}"
|
||
|
|
)
|
||
|
|
|
||
|
|
# [13] Health diagnostics
|
||
|
|
log(section_header(13, "Weight health diagnostics"))
|
||
|
|
|
||
|
|
high_sparsity = [(n, all_stats[n]["sparsity"]) for n in all_names if all_stats[n]["sparsity"] > 0.10]
|
||
|
|
high_kurtosis = [(n, all_stats[n]["kurtosis_delta"]) for n in all_names if abs(all_stats[n]["kurtosis_delta"]) > 5.0]
|
||
|
|
high_outlier = [(n, all_stats[n]["outlier_ratio"]) for n in all_names if all_stats[n]["outlier_ratio"] > 0.01]
|
||
|
|
dead_rows = [(n, int(all_stats[n].get("dead_rows", 0))) for n in all_names
|
||
|
|
if not math.isnan(all_stats[n].get("dead_rows", float("nan")))
|
||
|
|
and all_stats[n].get("dead_rows", 0) > 0]
|
||
|
|
low_cos = [(n, all_div[n]["cos_sim"]) for n in shared if all_div[n]["cos_sim"] < 0.95]
|
||
|
|
low_snr = [(n, all_div[n]["snr_db"]) for n in shared if all_div[n]["snr_db"] < 20.0]
|
||
|
|
|
||
|
|
def diag_block(title: str, rows: list, fmt):
|
||
|
|
log(f"\n {title}")
|
||
|
|
if not rows:
|
||
|
|
log(" none")
|
||
|
|
else:
|
||
|
|
for n, v in rows:
|
||
|
|
log(f" {n:<70} {fmt(v)}")
|
||
|
|
|
||
|
|
def get_percentiles(vals: list[float]) -> dict:
|
||
|
|
if not vals:
|
||
|
|
return {"mean": 0.0, "median": 0.0, "p10": 0.0, "p90": 0.0}
|
||
|
|
t = torch.tensor(vals, dtype=torch.float64)
|
||
|
|
return {
|
||
|
|
"mean": t.mean().item(),
|
||
|
|
"median": t.median().item(),
|
||
|
|
"p10": torch.quantile(t, 0.10).item(),
|
||
|
|
"p90": torch.quantile(t, 0.90).item(),
|
||
|
|
}
|
||
|
|
|
||
|
|
diag_block("Sparsity > 10%", high_sparsity, lambda v: f"sparsity={v:.5f}")
|
||
|
|
diag_block("|Kurtosis Delta| > 5.0", high_kurtosis, lambda v: f"kurt_delta={v:+.3f}")
|
||
|
|
diag_block("Outlier ratio > 1%", high_outlier, lambda v: f"outlier_ratio={v:.5f}")
|
||
|
|
diag_block("Dead rows (L2 < 1e-6)", dead_rows, lambda v: f"dead_rows={v}")
|
||
|
|
diag_block("Low cosine sim vs base (<0.95)", low_cos, lambda v: f"cos_sim={v:.6f}")
|
||
|
|
diag_block("Low SNR vs base (< 20 dB)", low_snr, lambda v: f"snr_db={v:.2f}")
|
||
|
|
|
||
|
|
log("\n Note on kurtosis delta: Kurtosis values are reported as the difference (delta) compared to the base model.")
|
||
|
|
log(" A high kurtosis delta on tiny vectors (like norm/q-k-norm vectors of size 128) is statistically expected")
|
||
|
|
log(" due to small sample sizes and does not indicate a model health or representation collapse issue.")
|
||
|
|
|
||
|
|
# [14] Executive summary
|
||
|
|
log(section_header(14, "Executive summary"))
|
||
|
|
|
||
|
|
all_cos = [all_div[n]["cos_sim"] for n in shared]
|
||
|
|
all_snr = [all_div[n]["snr_db"] for n in shared]
|
||
|
|
all_rel = [all_div[n]["rel_err"] for n in shared]
|
||
|
|
|
||
|
|
cos_stats = get_percentiles(all_cos)
|
||
|
|
snr_stats = get_percentiles(all_snr)
|
||
|
|
rel_stats = get_percentiles(all_rel)
|
||
|
|
|
||
|
|
log(f" shared tensors : {len(shared)}")
|
||
|
|
log(f" tensors changed vs base : {len(changed)} / {len(shared)}")
|
||
|
|
log(f" cosine similarity : mean = {cos_stats['mean']:.6f} | median = {cos_stats['median']:.6f} | p10 = {cos_stats['p10']:.6f} | p90 = {cos_stats['p90']:.6f}")
|
||
|
|
log(f" relative error : mean = {rel_stats['mean']:.6f} | median = {rel_stats['median']:.6f} | p10 = {rel_stats['p10']:.6f} | p90 = {rel_stats['p90']:.6f}")
|
||
|
|
log(f" SNR dB : mean = {snr_stats['mean']:.2f} | median = {snr_stats['median']:.2f} | p10 = {snr_stats['p10']:.2f} | p90 = {snr_stats['p90']:.2f}")
|
||
|
|
log(f" high-sparsity layers (>10%) : {len(high_sparsity)}")
|
||
|
|
log(f" heavy-tail layers (|kurt_d|>5.0) : {len(high_kurtosis)}")
|
||
|
|
log(f" dead-row layers : {len(dead_rows)}")
|
||
|
|
log(f" low-cos layers (<0.95) : {len(low_cos)}")
|
||
|
|
log(f" low-SNR layers (<20 dB) : {len(low_snr)}")
|
||
|
|
log(f" distillation alpha : {args.alpha}")
|
||
|
|
log("")
|
||
|
|
log(f" checkpoint size on disk : {fmt_size(total_ckpt_bytes)}")
|
||
|
|
log(f" base weights in memory : {fmt_size(sd_bytes(base_sd))}")
|
||
|
|
log(f" distilled weights in memory : {fmt_size(sd_bytes(dist_sd))}")
|
||
|
|
log("")
|
||
|
|
log(divider("="))
|
||
|
|
log(" END OF REPORT")
|
||
|
|
log(divider("="))
|
||
|
|
|
||
|
|
# Write to file
|
||
|
|
out = Path(args.output_file)
|
||
|
|
out.write_text("\n".join(R) + "\n", encoding="utf-8")
|
||
|
|
print(f"\nReport written to: {out.resolve()}")
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|