# Tensor Roll — Recursive CUDA-Q Model Quantizer # Copyright (C) 2026 SnapKitty Collective # SPDX-License-Identifier: AGPL-3.0-or-later """F5 resolution experiment: diagnose the 3/24 token agreement, attempt a principled fix, measure the result. Protocol (all measured, fixed seeds, no cherry-picking): V0: baseline demo --fast pipeline (control). Diagnose: divergence trace + first divergence + component ablations on V0. Section-6 check: core.select_plan vs fidelity.fidelity_aware_select on the same roll leaves — do the decisions differ, and does the fidelity-aware utility change anything measurable? V1: same init, distillation stages 2 (hidden-state) and 3 (logit) rerun at 3x steps — tests whether distillation was underpowered vs whether capacity/init is the bottleneck. Report: agreement/24 and mean teacher-KL for V0 and V1. Whatever the outcome — improvement, null, or regression — is reported. Failed experiments are valid output. """ import argparse import json import os import sys import time import numpy as np sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from tensor_roll import cli as CLI from tensor_roll import core as C from tensor_roll import fidelity as F from tensor_roll import model as M from tensor_roll import planner as P from tensor_roll import search as S from tensor_roll.util import write_json PROMPTS = [[3, 1, 4, 1], [7, 7, 2, 9], [0, 5, 5, 0]] N_TOKENS = 24 def ns(**kw): a = argparse.Namespace() for k, v in kw.items(): setattr(a, k, v) return a def measure_agreement_kl(t_params, t_cfg, s_params, s_cfg): agrees, n, kls = 0, 0, [] for pr in PROMPTS: tt, _, _ = M.generate(t_params, t_cfg, pr, N_TOKENS) st, _, _ = M.generate(s_params, s_cfg, pr, N_TOKENS) agrees += sum(a == b for a, b in zip(tt, st)) n += len(tt) rng = np.random.default_rng(4242) for _ in range(2): ids = M.gen_data(rng, 32, s_cfg) tl, _ = M.forward(t_params, t_cfg, ids[:, :-1]) sl, _ = M.forward(s_params, s_cfg, ids[:, :-1]) kls.append(float(M.kl_divergence_mean(tl, sl))) return {"agreement": f"{agrees}/{n}", "agree_frac": agrees / n, "mean_kl": float(np.mean(kls))} def main(): ap = argparse.ArgumentParser() ap.add_argument("--workdir", required=True) ap.add_argument("--steps-t", type=int, default=250) ap.add_argument("--steps-d", type=int, default=40) args = ap.parse_args() wdir = args.workdir os.makedirs(wdir, exist_ok=True) rep = {"protocol": "V0 baseline vs V1 focused-distillation", "prompts": PROMPTS, "n_tokens": N_TOKENS} # ---- teacher ------------------------------------------------------ print("== teacher ==", flush=True) CLI.cmd_ingest(ns(workdir=wdir, scale="tiny", steps=args.steps_t, batch=128, lr=3e-3, seed=11)) t_params, t_cfg = CLI._teacher(wdir) # ---- roll: section-6 comparison ----------------------------------- print("== roll: select_plan vs fidelity_aware_select ==", flush=True) stats = M.measure_tensor_stats(t_params, t_cfg) trees, leaves = S.run_roll_analysis( t_params, depth=2, quantum_policy="explore", stats=stats) budget = M.count_params(t_params) * S.BUDGET_RATIO plan0 = C.select_plan(leaves, budget) sens_rows = [{"tensor": l.name, "shape": [1], "dtype": "float32", "params": 1, "bytes": 4, "shard": "x", "is_matrix": False} for l in leaves] sens = {l.name: float(np.clip( l.metrics.get("spectral", 0.5), 0, 1)) for l in leaves} plan1 = F.fidelity_aware_select(leaves, budget, sens, lam=1.0) diff = sum(1 for l in leaves if plan0[l.name]["decision"] != plan1[l.name]["decision"]) rep["section6"] = { "n_leaves": len(leaves), "baseline_decisions": {d: sum(1 for l in leaves if plan0[l.name]["decision"] == d) for d in set(plan0[l.name]["decision"] for l in leaves)}, "fidelity_aware_decisions": {d: sum( 1 for l in leaves if plan1[l.name]["decision"] == d) for d in set(plan1[l.name]["decision"] for l in leaves)}, "n_decision_diffs": diff, "baseline_cost": plan0["_total_cost_params"], "fidelity_cost": plan1["_total_cost_params"], } print(f" leaves={len(leaves)} decision_diffs={diff} " f"cost0={plan0['_total_cost_params']} " f"cost1={plan1['_total_cost_params']}", flush=True) for l in leaves: l.decision = plan0[l.name] # baseline drives the pipeline # ---- compress + train + finetune (V0) ------------------------------ print("== compress ==", flush=True) CLI.cmd_compress(ns(workdir=wdir, budget_ratio=S.BUDGET_RATIO, seed=11)) print("== train (V0) ==", flush=True) CLI.cmd_train(ns(workdir=wdir, steps=args.steps_d, batch=128, lr=3e-3, seed=12, stage="all")) print("== finetune (V0) ==", flush=True) CLI.cmd_finetune(ns(workdir=wdir, steps=args.steps_d, batch=128, seed=13)) s_params = CLI._load_params(os.path.join(wdir, "student_final.npz")) s_cfg = M.Config.from_dict( CLI.read_json(os.path.join(wdir, "student_cfg.json"))) rep["V0"] = measure_agreement_kl(t_params, t_cfg, s_params, s_cfg) print(f" V0: {rep['V0']}", flush=True) # ---- diagnose V0 ---------------------------------------------------- print("== divergence trace (V0) ==", flush=True) tr = F.divergence_trace(t_params, t_cfg, s_params, s_cfg, PROMPTS[0], max_tokens=N_TOKENS) first = F.first_divergence(tr) rep["trace_V0"] = { "agreement": tr["token_agreement"], "first_divergence": first, "mean_kl": float(np.mean([p["kl_div"] for p in tr["positions"]])), "mean_js": float(np.mean([p["js_div"] for p in tr["positions"]])), } print(f" first divergence: position={first.get('position')} " f"layer={first.get('first_diverging_layer')} " f"KL={first.get('kl_div'):.3f}", flush=True) print("== ablations (V0) ==", flush=True) abl = F.ablation_report(t_params, t_cfg, s_params, s_cfg, PROMPTS, max_tokens=16) rep["ablation_V0"] = abl["ranking"] for r in abl["ranking"][:4]: print(f" {r['component']:24s} d_agree={r['delta_agree']:+.3f} " f"d_kl={r.get('delta_kl', 0):+.3f}", flush=True) # ---- V1: focused extra distillation --------------------------------- print("== V1: extra hidden/logit distillation ==", flush=True) CLI.cmd_train(ns(workdir=wdir, steps=args.steps_d * 3, batch=128, lr=3e-3, seed=12, stage="2")) CLI.cmd_train(ns(workdir=wdir, steps=args.steps_d * 3, batch=128, lr=3e-3, seed=12, stage="3")) s_params1 = CLI._load_params(os.path.join(wdir, "student_distilled.npz")) rep["V1"] = measure_agreement_kl(t_params, t_cfg, s_params1, s_cfg) print(f" V1: {rep['V1']}", flush=True) rep["delta_agree"] = rep["V1"]["agree_frac"] - rep["V0"]["agree_frac"] rep["delta_kl"] = rep["V1"]["mean_kl"] - rep["V0"]["mean_kl"] out = os.path.join(wdir, "fidelity_resolution_report.json") write_json(out, rep) print(f"wrote {out}", flush=True) print(json.dumps({k: rep[k] for k in ("V0", "V1", "section6") if k in rep}, indent=2)) if __name__ == "__main__": main()