Upload consistency_canon.py with huggingface_hub

consistency_canon.py

@@ -0,0 +1,163 @@
+#!/usr/bin/env python3
+"""
+Is approach (A) [canonicalize oracle SP* then MSE-regress] viable?
+Test: for the SAME context, optimize the SP from TWO different random inits -> sp_a, sp_b
+(both reach low KL, but symmetry makes them differ). Then PROCRUSTES-align sp_b to sp_a.
+  raw_dist  = ||sp_a - sp_b||^2 / token            (how different the two valid SPs are)
+  canon_dist= ||sp_a - sp_b@R||^2 / token          (after removing an orthogonal rotation)
+If canon_dist << raw_dist  => the symmetry is (mostly) a rotation -> canonicalization makes
+   the target unique -> MSE-regression viable => GO (A).
+If canon_dist ~ raw_dist   => symmetry is deeper than rotation -> MSE won't work => GO (B, KL).
+Also try per-token permutation alignment (the 128 soft tokens may be reorderable).
+"""
+import sys
+sys.path.insert(0, "/workspace")
+import argparse, json
+import torch
+import torch.nn.functional as F
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers.cache_utils import DynamicCache
+from train_qwen_distill import (HyperNetwork, Config, extract_qa, CJK_RE, TOOLCALL_RE,
+                                soft_prompt_stability_loss)
+
+
+@torch.no_grad()
+def teacher_dist(llm, embed, q, a, mq, ma, dev, dt):
+    cache = DynamicCache()
+    oq = llm(inputs_embeds=embed(q[:, :mq]).to(dt), attention_mask=torch.ones(1, mq, device=dev),
+             past_key_values=cache, use_cache=True, cache_position=torch.arange(mq, device=dev))
+    t0 = oq.logits[:, -1, :]
+    pos = torch.arange(mq, mq + ma, device=dev)
+    oa = llm(inputs_embeds=embed(a[:, :ma]).to(dt), attention_mask=torch.ones(1, mq + ma, device=dev),
+             past_key_values=cache, position_ids=pos.unsqueeze(0), use_cache=True, cache_position=pos)
+    V = oa.logits.size(-1); T = torch.empty(1, ma, V, dtype=oa.logits.dtype, device=dev)
+    T[:, 0] = t0; T[:, 1:] = oa.logits[:, :ma - 1]; return T
+
+
+@torch.no_grad()
+def prefill(llm, embed, q, mq, dev, dt):
+    c = DynamicCache()
+    llm(inputs_embeds=embed(q[:, :mq]).to(dt), attention_mask=torch.ones(1, mq, device=dev),
+        past_key_values=c, use_cache=True, cache_position=torch.arange(mq, device=dev))
+    return c
+
+
+def opt_sp(llm, embed, q, a, teacher, c0, c1, mq, dev, dt, cfg, mn, rw, steps, lr, S, T, seed):
+    cur = c1 - c0; R = min(c0, rw)
+    raw = embed(a[:, c0 - R:c0]).to(dt) if R > 0 else None
+    chunk = embed(a[:, c0:c1]).to(dt); n = S + R + cur
+    cp = torch.arange(mq, mq + n, device=dev)
+    tp = F.softmax(teacher[:, c0:c1].float() / T, dim=-1)
+    g = torch.Generator(device=dev).manual_seed(seed)
+    sp = torch.randn(1, S, cfg.hidden_dim, generator=g, device=dev)
+    sp = (sp / sp.norm(dim=-1, keepdim=True).clamp(min=1e-6) * cfg.target_norm).requires_grad_(True)
+    opt = torch.optim.Adam([sp], lr=lr); best = float("inf"); bv = None
+    for _ in range(steps):
+        opt.zero_grad(set_to_none=True)
+        nm = sp.norm(dim=-1, keepdim=True).clamp(min=1e-6); sc = torch.where(nm > mn, mn / nm, torch.ones_like(nm))
+        spc = (sp * sc).to(dt)
+        x = torch.cat([spc, raw, chunk], 1) if R > 0 else torch.cat([spc, chunk], 1)
+        cache = prefill(llm, embed, q, mq, dev, dt)
+        o = llm(inputs_embeds=x, attention_mask=torch.ones(1, mq + n, device=dev), past_key_values=cache,
+                position_ids=cp.unsqueeze(0), use_cache=True, cache_position=cp)
+        lp = F.log_softmax(o.logits[:, S - 1 + R:S - 1 + R + cur].float() / T, dim=-1)
+        kl = (tp * (tp.clamp_min(1e-9).log() - lp)).sum(-1).mean() * (T * T)
+        (kl + soft_prompt_stability_loss(sp, cfg)).backward(); opt.step()
+        if kl.item() < best:
+            best = kl.item(); bv = (sp.detach() * sc).clone()
+    return bv[0], best        # (S,H), kl
+
+
+def procrustes(Bm, Am):       # align Bm to Am: return Bm@R, R orthogonal minimizing ||Am - Bm R||
+    M = Bm.transpose(-1, -2) @ Am
+    U, S, Vh = torch.linalg.svd(M)
+    return Bm @ (U @ Vh)
+
+
+def perm_align(Bm, Am):       # align rows of Bm to Am by greedy nearest (permutation of 128 tokens)
+    # cost[i,j] = ||Am[i]-Bm[j]||^2 ; greedy match
+    d = torch.cdist(Am, Bm)   # (S,S)
+    used = torch.zeros(Bm.size(0), dtype=torch.bool, device=Bm.device)
+    out = torch.empty_like(Bm)
+    for i in range(Am.size(0)):
+        row = d[i].clone(); row[used] = float("inf"); j = int(row.argmin()); used[j] = True; out[i] = Bm[j]
+    return out
+
+
+def load_samples(path, tok, cfg, n, mc, mal, mt):
+    out = []
+    with open(path) as f:
+        for line in f:
+            if len(out) >= n: break
+            line = line.strip()
+            if not line: continue
+            try: row = json.loads(line)
+            except Exception: continue
+            q, a = extract_qa(row, cfg)
+            if not q or not a or len(a) < mc: continue
+            if CJK_RE.search(a) or CJK_RE.search(q) or TOOLCALL_RE.search(a) or TOOLCALL_RE.search(q): continue
+            qi = tok(q, max_length=cfg.max_query_len, truncation=True, add_special_tokens=True).input_ids
+            ai = tok(a, max_length=mal, truncation=True, add_special_tokens=False).input_ids
+            if len(ai) < mt: continue
+            out.append((qi, ai))
+    return out[300:300 + n] if len(out) > 300 + n else out
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--ckpt", default="/workspace/hypernet_qwen/hn_step7750.pt")
+    p.add_argument("--base_model", default="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B")
+    p.add_argument("--data", default="/workspace/dolphin_subset.jsonl")
+    p.add_argument("--n", type=int, default=12)
+    p.add_argument("--depths", default="128,256,384")
+    p.add_argument("--chunk_size", type=int, default=64)
+    p.add_argument("--raw_window", type=int, default=32)
+    p.add_argument("--steps", type=int, default=120)
+    p.add_argument("--lr", type=float, default=0.03)
+    args = p.parse_args()
+    dev = torch.device("cuda"); dt = torch.bfloat16
+    cfg = Config(); cfg.base_model = args.base_model
+    C = args.chunk_size; S = cfg.num_soft_tokens; T = 1.0
+    print("Loading frozen base...", flush=True)
+    tok = AutoTokenizer.from_pretrained(cfg.base_model)
+    if tok.pad_token is None: tok.pad_token = tok.eos_token
+    llm = AutoModelForCausalLM.from_pretrained(cfg.base_model, dtype=dt, device_map="cuda", attn_implementation="sdpa")
+    llm.config.use_cache = True
+    for prm in llm.parameters(): prm.requires_grad_(False)
+    llm.eval(); embed = llm.get_input_embeddings(); cfg.hidden_dim = llm.config.hidden_size
+    with torch.no_grad():
+        ids = torch.randint(0, embed.weight.size(0), (512,), device=dev)
+        cfg.target_norm = embed(ids).float().norm(dim=-1).mean().item()
+    mn = cfg.target_norm * 3.0
+    hn = HyperNetwork(cfg).to(dtype=torch.float32, device=dev); hn.eval()
+    ckd = torch.load(args.ckpt, map_location="cpu", weights_only=False); hn.load_state_dict(ckd["hypernet"], strict=False)
+    depths = [int(x) for x in args.depths.split(",")]
+    samples = load_samples(args.data, tok, cfg, args.n, 1500, 512, 400)
+    raw_l, rot_l, prm_l, kl_l = [], [], [], []
+    for si, (qi, ai) in enumerate(samples):
+        q = torch.tensor([qi], device=dev); a = torch.tensor([ai], device=dev); mq = q.size(1); ma = a.size(1)
+        teacher = teacher_dist(llm, embed, q, a, mq, ma, dev, dt)
+        for c0 in depths:
+            c1 = min(c0 + C, ma)
+            if c1 - c0 < 4 or c0 + 1 >= ma: continue
+            spa, ka = opt_sp(llm, embed, q, a, teacher, c0, c1, mq, dev, dt, cfg, mn, args.raw_window, args.steps, args.lr, S, T, 111)
+            spb, kb = opt_sp(llm, embed, q, a, teacher, c0, c1, mq, dev, dt, cfg, mn, args.raw_window, args.steps, args.lr, S, T, 999)
+            raw = ((spa - spb) ** 2).sum(-1).mean().item()
+            rot = ((spa - procrustes(spb, spa)) ** 2).sum(-1).mean().item()
+            prm = ((spa - perm_align(spb, spa)) ** 2).sum(-1).mean().item()
+            raw_l.append(raw); rot_l.append(rot); prm_l.append(prm); kl_l.append(0.5 * (ka + kb))
+            print(f"  s{si+1} c0={c0}: KL~{0.5*(ka+kb):.4f} raw={raw:.3f} afterRotation={rot:.3f} afterPerm={prm:.3f}", flush=True)
+        del teacher; torch.cuda.empty_cache()
+    R = sum(raw_l)/len(raw_l); RO = sum(rot_l)/len(rot_l); PR = sum(prm_l)/len(prm_l)
+    print("\n" + "=" * 64)
+    print(f"mean KL={sum(kl_l)/len(kl_l):.4f}  (both SPs valid)")
+    print(f"raw disagreement      = {R:.3f}")
+    print(f"after rotation align  = {RO:.3f}   ({100*RO/R:.0f}% of raw)")
+    print(f"after permutation align = {PR:.3f} ({100*PR/R:.0f}% of raw)")
+    print("=" * 64)
+    print("if rotation/perm align << raw (e.g. <40%): symmetry is removable -> (A) canonicalize+MSE VIABLE")
+    print("if still ~raw: symmetry deeper than rotation/perm -> go (B) output-space KL")
+
+
+if __name__ == "__main__":
+    main()