""" 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 = "/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()