Upload dcor_recurrence.py with huggingface_hub

dcor_recurrence.py

@@ -0,0 +1,214 @@
+#!/usr/bin/env python3
+"""
+Is the recurrence input sp^m related to its output sp^{m+1}?  (= is sp^m -> sp^{m+1}
+a learnable, amortizable map?)  Each sp^j is optimized INDEPENDENTLY from the FIXED
+init_sp anchor (deterministic; NO chaining -> no warm-start proximity artifact), so the
+dCor reflects genuine shared-history structure, not closeness-by-construction.
+
+Measured over consecutive pairs pooled across many samples (cross-sample = amortization angle):
+  dCor(sp^m , sp^{m+1})            -- does the prev SP predict the next SP at all?
+  dCor([sp^m;chunk_m], sp^{m+1})   -- does the actual recurrence INPUT predict the output?
+vs shuffled null. PCA-reduced (N>>dim -> clean null).
+"""
+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)
+
+
+def dcor(X, Y):
+    a = torch.cdist(X, X); b = torch.cdist(Y, Y)
+    A = a - a.mean(0, keepdim=True) - a.mean(1, keepdim=True) + a.mean()
+    B = b - b.mean(0, keepdim=True) - b.mean(1, keepdim=True) + b.mean()
+    dcov2 = (A * B).mean(); dvx = (A * A).mean(); dvy = (B * B).mean()
+    return float((dcov2.clamp(min=0) / (dvx.sqrt() * dvy.sqrt()).clamp(min=1e-12)).sqrt().item())
+
+
+def zscore(X): return (X - X.mean(0, keepdim=True)) / X.std(0, keepdim=True).clamp(min=1e-6)
+
+
+def pca(X, k):
+    Xc = X - X.mean(0, keepdim=True)
+    U, S, Vh = torch.linalg.svd(Xc, full_matrices=False)
+    return Xc @ Vh[:min(k, Vh.size(0))].T
+
+
+@torch.no_grad()
+def teacher_dist(llm, embed, q, valid_q, a, max_q, max_a, q_lens, device, dtype):
+    B = q.size(0); cache = DynamicCache()
+    qe = (embed(q[:, :max_q]) * valid_q.unsqueeze(-1)).to(dtype)
+    oq = llm(inputs_embeds=qe, attention_mask=valid_q, past_key_values=cache, use_cache=True,
+             cache_position=torch.arange(max_q, device=device))
+    t0 = oq.logits[torch.arange(B, device=device), (q_lens - 1).clamp(min=0), :]
+    ae = embed(a[:, :max_a]).to(dtype)
+    pos = torch.arange(max_q, max_q + max_a, device=device)
+    attn = torch.cat([valid_q, torch.ones(B, max_a, dtype=torch.long, device=device)], 1)
+    oa = llm(inputs_embeds=ae, attention_mask=attn, past_key_values=cache,
+             position_ids=pos.unsqueeze(0).expand(B, -1), use_cache=True, cache_position=pos)
+    V = oa.logits.size(-1); T = torch.empty(B, max_a, V, dtype=oa.logits.dtype, device=device)
+    T[:, 0] = t0; T[:, 1:] = oa.logits[:, :max_a - 1]; return T
+
+
+@torch.no_grad()
+def prefill(llm, embed, q, valid_q, max_q, device, dtype):
+    cache = DynamicCache()
+    llm(inputs_embeds=(embed(q[:, :max_q]) * valid_q.unsqueeze(-1)).to(dtype),
+        attention_mask=valid_q, past_key_values=cache, use_cache=True,
+        cache_position=torch.arange(max_q, device=device))
+    return cache
+
+
+def opt_chunk_sp(llm, embed, q, valid_q, a, a_lens, teacher, init_sp, c0, c1, max_q, device,
+                 dtype, cfg, max_norm, raw_window, steps, lr, S, T):
+    """fixed-init independent optimization of the SP for predicting chunk [c0:c1]."""
+    B = q.size(0); cur_C = c1 - c0; R = min(c0, raw_window)
+    raw = embed(a[:, c0 - R:c0]).to(dtype) if R > 0 else None
+    chunk = embed(a[:, c0:c1]).to(dtype)
+    n_new = S + R + cur_C
+    cpos = torch.arange(max_q, max_q + n_new, device=device)
+    attn = torch.cat([valid_q, torch.ones(B, n_new, dtype=torch.long, device=device)], 1)
+    t_p = F.softmax(teacher[:, c0:c1].float() / T, dim=-1)
+    posv = torch.arange(cur_C, device=device)
+    vf = ((c0 + posv).unsqueeze(0) < a_lens.unsqueeze(1)).float()
+    valid = (c0 < a_lens)
+    sp = init_sp.expand(B, -1, -1).clone().requires_grad_(True)
+    opt = torch.optim.Adam([sp], lr=lr)
+    best_kl = torch.full((B,), float("inf"), device=device); best = sp.detach().clone()
+    for _ in range(steps):
+        opt.zero_grad(set_to_none=True)
+        n = sp.norm(dim=-1, keepdim=True).clamp(min=1e-6)
+        sc = torch.where(n > max_norm, max_norm / n, torch.ones_like(n))
+        spc = (sp * sc).to(dtype)
+        x = torch.cat([spc, raw, chunk], 1) if R > 0 else torch.cat([spc, chunk], 1)
+        cache = prefill(llm, embed, q, valid_q, max_q, device, dtype)
+        o = llm(inputs_embeds=x, attention_mask=attn, past_key_values=cache,
+                position_ids=cpos.unsqueeze(0).expand(B, -1), use_cache=True, cache_position=cpos)
+        lp = F.log_softmax(o.logits[:, S - 1 + R:S - 1 + R + cur_C].float() / T, dim=-1)
+        kl = ((t_p * (t_p.clamp_min(1e-9).log() - lp)).sum(-1) * vf).sum(1) / vf.sum(1).clamp(min=1) * (T * T)
+        loss = (kl * valid.float()).sum() / valid.float().sum().clamp(min=1)
+        (loss + soft_prompt_stability_loss(sp, cfg)).backward(); opt.step()
+        with torch.no_grad():
+            imp = (kl < best_kl) & valid
+            if imp.any():
+                n2 = sp.norm(dim=-1, keepdim=True).clamp(min=1e-6)
+                sc2 = torch.where(n2 > max_norm, max_norm / n2, torch.ones_like(n2))
+                best[imp] = (sp * sc2)[imp].detach(); best_kl[imp] = kl[imp]
+    return best.detach(), valid, best_kl
+
+
+def load_samples(path, tok, cfg, n, min_chars, mal, min_tok):
+    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) < min_chars: 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) < min_tok: continue
+            out.append((qi, ai))
+    out.sort(key=lambda x: len(x[1])); return 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_samples", type=int, default=24)
+    p.add_argument("--skip", type=int, default=300)
+    p.add_argument("--max_ans_len", type=int, default=512)
+    p.add_argument("--batch", type=int, default=12)
+    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=100)
+    p.add_argument("--lr", type=float, default=0.03)
+    p.add_argument("--pca_k", type=int, default=24)
+    args = p.parse_args()
+    device = torch.device("cuda"); dtype = 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=dtype, 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=device)
+        cfg.target_norm = embed(ids).float().norm(dim=-1).mean().item()
+    max_norm = cfg.target_norm * 3.0
+    hn = HyperNetwork(cfg).to(dtype=torch.float32, device=device); hn.eval()
+    ckd = torch.load(args.ckpt, map_location="cpu", weights_only=False)
+    hn.load_state_dict(ckd["hypernet"], strict=False)
+    init_sp = hn.init_sp.detach().clone()
+    alls = load_samples(args.data, tok, cfg, args.skip + args.n_samples, 1500, args.max_ans_len, 400)
+    samples = alls[args.skip:args.skip + args.n_samples]
+    print(f"  {len(samples)} samples, fixed-init independent SP per chunk, steps={args.steps}\n", flush=True)
+
+    spm, spm1, chunkm = [], [], []   # sp^m, sp^{m+1}, chunk_m (pooled)
+    B = args.batch; kl_all = []
+    for bi in range(0, len(samples), B):
+        chunk_samples = samples[bi:bi + B]; b = len(chunk_samples)
+        max_q = max(len(s[0]) for s in chunk_samples); max_a = max(len(s[1]) for s in chunk_samples)
+        q = torch.zeros(b, max_q, dtype=torch.long, device=device)
+        a = torch.zeros(b, max_a, dtype=torch.long, device=device)
+        ql = torch.zeros(b, dtype=torch.long, device=device); al = torch.zeros(b, dtype=torch.long, device=device)
+        for i, (qi, ai) in enumerate(chunk_samples):
+            q[i, :len(qi)] = torch.tensor(qi, device=device); ql[i] = len(qi)
+            a[i, :len(ai)] = torch.tensor(ai, device=device); al[i] = len(ai)
+        valid_q = (torch.arange(max_q, device=device).unsqueeze(0) < ql.unsqueeze(1)).long()
+        teacher = teacher_dist(llm, embed, q, valid_q, a, max_q, max_a, ql, device, dtype)
+        n_chunks = (max_a + C - 1) // C
+        sp_by_chunk = {}
+        for j in range(1, n_chunks):
+            c0 = j * C; c1 = min(c0 + C, max_a)
+            if c1 - c0 < 4: break
+            sp_j, valid, klj = opt_chunk_sp(llm, embed, q, valid_q, a, al, teacher, init_sp,
+                                            c0, c1, max_q, device, dtype, cfg, max_norm,
+                                            args.raw_window, args.steps, args.lr, S, T)
+            sp_by_chunk[j] = (sp_j, valid)
+            kl_all.append(klj[valid].mean().item() if valid.any() else float('nan'))
+        # consecutive pairs
+        for j in sorted(sp_by_chunk):
+            if j + 1 not in sp_by_chunk: continue
+            spj, vj = sp_by_chunk[j]; spj1, vj1 = sp_by_chunk[j + 1]
+            c0 = j * C; c1 = min(c0 + C, max_a)
+            ce = embed(a[:, c0:c1]).float().mean(1)        # chunk_m pooled
+            both = vj & vj1
+            for i in range(b):
+                if bool(both[i]):
+                    spm.append(spj[i].mean(0).float().cpu())
+                    spm1.append(spj1[i].mean(0).float().cpu())
+                    chunkm.append(ce[i].cpu())
+        del teacher; torch.cuda.empty_cache()
+        print(f"  batch {bi//B+1} done, pairs so far={len(spm)}", flush=True)
+
+    Xm = zscore(torch.stack(spm).to(device)); Xm1 = zscore(torch.stack(spm1).to(device))
+    Xc = zscore(torch.stack(chunkm).to(device))
+    N = Xm.size(0); K = args.pca_k; perm = torch.randperm(N, device=device)
+    Xm_r, Xm1_r, Xc_r = pca(Xm, K), pca(Xm1, K), pca(Xc, K)
+    Xin_r = pca(torch.cat([Xm, Xc], 1), K)
+    print(f"\nN consecutive pairs = {N} | mean opt KL = {sum(x for x in kl_all if x==x)/max(len(kl_all),1):.4f}")
+    print("=" * 64)
+    print(f"dCor(sp^m,        sp^m+1) = {dcor(Xm_r, Xm1_r):.4f}  (null {dcor(Xm_r, Xm1_r[perm]):.4f})")
+    print(f"dCor([sp^m;chunk],sp^m+1) = {dcor(Xin_r, Xm1_r):.4f}  (null {dcor(Xin_r, Xm1_r[perm]):.4f})")
+    print(f"dCor(chunk_m,     sp^m+1) = {dcor(Xc_r, Xm1_r):.4f}  (null {dcor(Xc_r, Xm1_r[perm]):.4f})")
+    print("=" * 64)
+    print(">>null => the recurrence INPUT predicts the next SP => recurrence is learnable/amortizable.")
+
+
+if __name__ == "__main__":
+    main()