Something went wrong. Try again.
strata but optimized for m5 (and maybe more later)
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114import json, os, re, subprocess, timefrom pathlib import Path
PPL_BIN = "/content/llama.cpp/build/bin/llama-perplexity"PPL_CPP = Path("/content/llama.cpp/tools/perplexity/perplexity.cpp")assert "common_tokenize(ctx, params.prompt, true, true)" in PPL_CPP.read_text(), "perplexity.cpp missing parse_special=true patch!"SHARD1 = "/content/models/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-00001-of-00002.gguf"SHARD2 = "/content/models/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-00002-of-00002.gguf"
EVAL_DOMAINS = { "ops_tools": ("/content/eval_ops_tools.txt", "/content/base_logits_ops_tools.bin"), "chat_json": ("/content/eval_chat_json.txt", "/content/base_logits_chat_json.bin"), "ja": ("/content/eval_ja.txt", "/content/base_logits_ja.bin"),}
def grab(pattern: str, text: str) -> float: m = re.search(pattern, text) if not m: raise RuntimeError(f"Failed to match {pattern} in output:\n{text[-1500:]}") return float(m.group(1))
results_path = Path("/content/phase1c_kld_results.json")all_res = json.loads(results_path.read_text()) if results_path.exists() else {}
if "K512_unpruned" not in all_res: print("=== Step 1: Evaluating K512_unpruned baseline & saving teacher logits ===", flush=True) base_res = {} for dom, (fpath, lfile) in EVAL_DOMAINS.items(): t0 = time.monotonic() cmd = [ PPL_BIN, "-m", SHARD1, "-f", fpath, "-c", "2048", "-b", "512", "-ngl", "99", "--chunks", "4", "--save-all-logits", lfile, ] p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) if p.returncode != 0: raise RuntimeError(f"Baseline PPL failed on {dom} (exit {p.returncode}):\n{p.stdout[-1500:]}") dt = time.monotonic() - t0 ppl = grab(r"Final estimate: PPL = ([0-9.]+)", p.stdout) ppl_se = grab(r"Final estimate: PPL = [0-9.]+\s*\+/-\s*([0-9.]+)", p.stdout) base_res[dom] = {"ppl": ppl, "ppl_se": ppl_se, "sec": round(dt, 1)} print(f"[K512_unpruned] {dom}: PPL = {ppl:.4f} +/- {ppl_se:.5f} ({dt:.1f}s)", flush=True) all_res["K512_unpruned"] = base_res results_path.write_text(json.dumps(all_res, indent=2))
CANDIDATES = [ ("K320_reap_l2_sum", "/content/keep_maps/keep_reap_l2_sum_k320.json"), ("K320_reap_hybrid", "/content/keep_maps/keep_reap_hybrid_k320.json"), ("K256_reap_l2_sum", "/content/keep_maps/keep_reap_l2_sum_k256.json"), ("K256_reap_hybrid", "/content/keep_maps/keep_reap_hybrid_k256.json"), ("K256_reap_shrink", "/content/keep_maps/keep_reap_shrink_k256.json"), ("K224_reap_l2_sum", "/content/keep_maps/keep_reap_l2_sum_k224.json"),]
os.makedirs("/dev/shm/pruned", exist_ok=True)for label, km_path in CANDIDATES: if label in all_res: continue out1 = f"/dev/shm/pruned/{label}-00001-of-00002.gguf" out2 = f"/dev/shm/pruned/{label}-00002-of-00002.gguf" t0 = time.monotonic() print(f"=== Slicing {label} ===", flush=True) subprocess.run(["python3", "/content/slice_qwen4exp_gguf.py", "--in-gguf", SHARD1, "--out-gguf", out1, "--keep-json", km_path], check=True) subprocess.run(["python3", "/content/slice_qwen4exp_gguf.py", "--in-gguf", SHARD2, "--out-gguf", out2, "--keep-json", km_path], check=True) print(f"Sliced {label} in {time.monotonic()-t0:.1f}s", flush=True)
cand_res = {} for dom, (fpath, lfile) in EVAL_DOMAINS.items(): t1 = time.monotonic() cmd = [ PPL_BIN, "-m", out1, "-f", fpath, "-c", "2048", "-b", "512", "-ngl", "99", "--chunks", "4", "--kl-divergence", "--kl-divergence-base", lfile, ] p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) if p.returncode != 0: raise RuntimeError(f"KLD eval failed for {label} on {dom} (exit {p.returncode}):\n{p.stdout[-1500:]}") dt = time.monotonic() - t1 ppl_q = grab(r"Mean PPL\(Q\)\s*:\s*([0-9.]+)", p.stdout) ppl_base = grab(r"Mean PPL\(base\)\s*:\s*([0-9.]+)", p.stdout) ppl_diff = grab(r"Mean PPL\(Q\)-PPL\(base\)\s*:\s*([-+0-9.]+)", p.stdout) mean_kld = grab(r"Mean\s+KLD:\s*([-+0-9.]+)", p.stdout) mean_kld_se = grab(r"Mean\s+KLD:\s*[-+0-9.]+\s*±\s*([0-9.]+)", p.stdout) med_kld = grab(r"Median\s+KLD:\s*([-+0-9.]+)", p.stdout) p90_kld = grab(r"90\.0%\s+KLD:\s*([-+0-9.]+)", p.stdout) p99_kld = grab(r"99\.0%\s+KLD:\s*([-+0-9.]+)", p.stdout) rms_dp = grab(r"RMS Δp\s*:\s*([0-9.]+)", p.stdout) same_top = grab(r"Same top p:\s*([0-9.]+)", p.stdout) base_fp32 = all_res["K512_unpruned"][dom]["ppl"] cand_res[dom] = { "ppl": ppl_q, "ppl_base": ppl_base, "delta_ppl": ppl_diff, "exact_delta_ppl": round(ppl_q - base_fp32, 4), "mean_kld": mean_kld, "mean_kld_se": mean_kld_se, "median_kld": med_kld, "p90_kld": p90_kld, "p99_kld": p99_kld, "rms_delta_p_pct": rms_dp, "same_top1_pct": same_top, "sec": round(dt, 1), } print(f"[{label}] {dom}: PPL={ppl_q:.4f} (Δ={ppl_diff:+.4f}) | KLD={mean_kld:.5f}±{mean_kld_se:.5f} | Top1={same_top:.2f}% ({dt:.1f}s)", flush=True)
all_res[label] = cand_res results_path.write_text(json.dumps(all_res, indent=2)) os.remove(out1) os.remove(out2)
print("=== ALL PHASE 1C TIER-1 KLD + TOP-1 + PPL EVALS COMPLETE! ===", flush=True)