#!/usr/bin/env python3 """ Analyze REAP saliency JSON from llama-imatrix (ReapTraceCollector v2), verify §2.5 invariants, generate per-layer expert keep-maps for K in {320, 256, 224} across: - reap_l2_sum: count * REAP (expected gated output magnitude per token) - reap_shrink: zero-prior damped REAP L2 mean S_e / (c_e + m) with LCB guard (protects high-norm specialists with real support while scoring 0-hit/noise experts at 0) - reap_hybrid: rank-percentile average 0.5 * pct(reap_l2_sum) + 0.5 * pct(reap_shrink) (outlier-proof across skewed layers like L47) and compute offline top-10 substitution error on a Q4RT trace file (§3.3). """ import argparse import json import math from pathlib import Path def verify_invariants(data: dict) -> dict: layers = data["layers"] assert len(layers) == 48, f"Expected 48 layers, got {len(layers)}" n_picks = sum(layers["0"]["count"]) assert n_picks > 0 and n_picks % 10 == 0, f"Invalid L0 pick count: {n_picks}" n_tokens = n_picks // 10 for l_idx in range(48): ld = layers[str(l_idx)] c_sum = sum(ld["count"]) p_sum = sum(ld["prob_sum"]) r_sum = sum(ld["reap_l2_sum"]) assert c_sum == n_picks, f"Layer {l_idx} count sum {c_sum} != {n_picks}" assert abs(p_sum - n_tokens) / max(1, n_tokens) < 1e-2, f"Layer {l_idx} prob_sum {p_sum} != {n_tokens}" assert r_sum > 0.0, f"Layer {l_idx} reap_l2_sum is zero!" for e in range(len(ld["count"])): if ld["count"][e] == 0: assert ld["reap_l2_sum"][e] == 0.0, f"Layer {l_idx} expert {e} count=0 but reap>0" return { "total_tokens_evaluated": n_tokens, "total_picks_per_layer": n_picks, "nonfinite_drops": data.get("nonfinite_drops", 0), "trace_tokens_written": data.get("trace_tokens_written", 0), } def _rank_percentiles(vals: list[float]) -> list[float]: n = len(vals) order = sorted(range(n), key=lambda i: (vals[i], i)) pct = [0.0] * n for rank, idx in enumerate(order): pct[idx] = rank / max(1, n - 1) return pct def compute_layer_scores(ld: dict, metric: str, shrink_m: float = 64.0) -> list[float]: counts = ld["count"] reap_sum = ld["reap_l2_sum"] reap_sq = ld.get("reap_l2_sq_sum", [0.0] * len(counts)) n_exp = len(counts) if metric == "reap_l2_sum": return list(reap_sum) if metric == "prob_sum": return list(ld["prob_sum"]) # Zero-prior damped mean with LCB guard: # Dead experts (c == 0) score 0.0. Rare experts require genuine sample support (c >> m) # and low variance (mean - 1.0 * SE) so 1-hit/5-hit noise never beats active experts. shrink_scores = [] for c, r, sq in zip(counts, reap_sum, reap_sq): if c <= 0: shrink_scores.append(0.0) else: mean = r / c var = max(0.0, (sq / c) - (mean * mean)) se = math.sqrt(var / c) lcb = max(0.0, mean - 1.0 * se) damped = (c / (c + shrink_m)) * lcb shrink_scores.append(damped) if metric == "reap_shrink": return shrink_scores if metric == "reap_hybrid": pct_sum = _rank_percentiles(reap_sum) pct_shr = _rank_percentiles(shrink_scores) return [0.5 * ps + 0.5 * pr for ps, pr in zip(pct_sum, pct_shr)] raise ValueError(f"Unknown metric: {metric}") def score_trace_substitution(trace_path: str, keep_sets: dict[int, set[int]], k_used: int = 10) -> dict: import struct raw = Path(trace_path).read_bytes() magic, ver, n_layers, n_experts, top_m, n_tokens = struct.unpack_from("<4sHHHH I", raw, 0) rec_bytes = top_m * 4 actual_tokens = (len(raw) - 16) // (n_layers * rec_bytes) if n_tokens == 0 or actual_tokens < n_tokens: n_tokens = actual_tokens sample_tokens = min(n_tokens, 2048) ids_fmt = f"<{top_m}H" logits_fmt = f"<{top_m}e" id_bytes = top_m * 2 total_slots = 0 replaced_slots = 0 prob_displaced = 0.0 for t in range(sample_tokens): base_t = 16 + t * n_layers * rec_bytes for l in range(n_layers): off = base_t + l * rec_bytes ids = struct.unpack_from(ids_fmt, raw, off) logits = struct.unpack_from(logits_fmt, raw, off + id_bytes) orig_top = ids[:k_used] orig_logits = logits[:k_used] max_z = orig_logits[0] exps = [math.exp(z - max_z) for z in orig_logits] z_sum = sum(exps) probs = [x / z_sum for x in exps] kset = keep_sets[l] for eid, p in zip(orig_top, probs): total_slots += 1 if eid not in kset: replaced_slots += 1 prob_displaced += p return { "slot_replacement_rate": round(replaced_slots / max(1, total_slots), 4), "mean_top10_weight_displaced": round(prob_displaced / max(1, sample_tokens * n_layers), 4), } def generate_keep_maps( saliency_path: str, out_dir: str, ks: list[int] = [320, 256, 224], metric: str = "reap_l2_sum", trace_path: str | None = None, ) -> dict: data = json.loads(Path(saliency_path).read_text()) inv = verify_invariants(data) layers = data["layers"] n_layers = len(layers) Path(out_dir).mkdir(parents=True, exist_ok=True) report = { "invariants": inv, "metric": metric, "variants": {}, } for k in ks: keep_map = {} keep_sets = {} reap_retained_per_layer = [] prob_retained_per_layer = [] count_retained_per_layer = [] min_kept_count_per_layer = [] zero_hit_kept = 0 for l_idx in range(n_layers): l_data = layers[str(l_idx)] scores = compute_layer_scores(l_data, metric=metric) reap_l2 = l_data["reap_l2_sum"] probs = l_data["prob_sum"] counts = l_data["count"] n_exp = len(scores) ranked = sorted(range(n_exp), key=lambda e: (scores[e], reap_l2[e], probs[e]), reverse=True) kept = sorted(ranked[:k]) keep_map[str(l_idx)] = kept keep_sets[l_idx] = set(kept) tot_reap = sum(reap_l2) kept_reap = sum(reap_l2[e] for e in kept) reap_retained_per_layer.append(kept_reap / tot_reap if tot_reap > 0 else 1.0) tot_prob = sum(probs) kept_prob = sum(probs[e] for e in kept) prob_retained_per_layer.append(kept_prob / tot_prob if tot_prob > 0 else 1.0) tot_cnt = sum(counts) kept_cnt = sum(counts[e] for e in kept) count_retained_per_layer.append(kept_cnt / tot_cnt if tot_cnt > 0 else 1.0) min_c = min(counts[e] for e in kept) min_kept_count_per_layer.append(min_c) zero_hit_kept += sum(1 for e in kept if counts[e] == 0) assert zero_hit_kept == 0, f"Metric {metric} at K={k} kept {zero_hit_kept} dead (0-hit) experts!" out_file = Path(out_dir) / f"keep_{metric}_k{k}.json" out_file.write_text(json.dumps(keep_map, indent=2)) v_info = { "keep_json": str(out_file), "mean_reap_l2_retained": round(sum(reap_retained_per_layer) / n_layers, 4), "min_reap_l2_retained": round(min(reap_retained_per_layer), 4), "min_reap_layer": int(min(range(n_layers), key=lambda i: reap_retained_per_layer[i])), "mean_prob_retained": round(sum(prob_retained_per_layer) / n_layers, 4), "WholeModel_slot_replacement": round(1.0 - sum(count_retained_per_layer) / n_layers, 4), "min_kept_expert_hits": int(min(min_kept_count_per_layer)), "zero_hit_experts_kept": zero_hit_kept, } if trace_path and Path(trace_path).exists(): v_info["trace_substitution"] = score_trace_substitution(trace_path, keep_sets) report["variants"][str(k)] = v_info return report def main(): p = argparse.ArgumentParser() p.add_argument("--saliency", required=True) p.add_argument("--out-dir", default="/content/keep_maps") p.add_argument("--metric", default="reap_l2_sum", choices=["reap_l2_sum", "reap_shrink", "reap_hybrid", "prob_sum"]) p.add_argument("--trace", default="/content/routing_trace_v2.q4rt") args = p.parse_args() rep = generate_keep_maps(args.saliency, args.out_dir, metric=args.metric, trace_path=args.trace) print(json.dumps(rep, indent=2)) if __name__ == "__main__": main()