File size: 7,621 Bytes
9e551a1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | # 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()
|