Upload consistency_sp.py with huggingface_hub

consistency_sp.py

@@ -0,0 +1,198 @@
+#!/usr/bin/env python3
+"""
+Target-consistency test (the RIGHT test, replacing the loose overfit argument).
+
+For each (sample, depth c0) we optimize the target SP* TWICE from two DIFFERENT
+random initializations (constrained: clamp to target_norm*3 + stability), to the
+same KL objective. Then we measure:
+  - KL_a, KL_b : both should reach a good low KL (the SP* are both "correct")
+  - cross-seed per-token L2 ||SP*_a - SP*_b||^2 : if BOTH are low-KL yet this L2 is
+    LARGE, the optimum is NON-UNIQUE (soft-prompt symmetry) -> MSE-to-a-fixed-vector
+    has an irreducible floor ~ this L2, and we should distill in OUTPUT space (KL).
+  - also the L2 between each SP* and their midpoint, and norms, for context.
+
+Compare cross-seed L2 to the ~1.1 regression floor we observed. If they match,
+symmetry is the cause.
+"""
+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_ids, a_ids, max_q, max_a, device, dtype):
+    cache = DynamicCache()
+    q_emb = embed(q_ids[:, :max_q]).to(dtype)
+    out_q = llm(inputs_embeds=q_emb, attention_mask=torch.ones(1, max_q, device=device),
+                past_key_values=cache, use_cache=True, cache_position=torch.arange(max_q, device=device))
+    t0 = out_q.logits[:, -1, :]
+    a_emb = embed(a_ids[:, :max_a]).to(dtype)
+    pos_a = torch.arange(max_q, max_q + max_a, device=device)
+    out_a = llm(inputs_embeds=a_emb, attention_mask=torch.ones(1, max_q + max_a, device=device),
+                past_key_values=cache, position_ids=pos_a.unsqueeze(0), use_cache=True, cache_position=pos_a)
+    V = out_a.logits.size(-1)
+    teacher = torch.empty(1, max_a, V, dtype=out_a.logits.dtype, device=device)
+    teacher[:, 0, :] = t0
+    teacher[:, 1:, :] = out_a.logits[:, :max_a - 1, :]
+    return teacher
+
+
+@torch.no_grad()
+def prefill_query(llm, embed, q_ids, max_q, device, dtype):
+    cache = DynamicCache()
+    llm(inputs_embeds=embed(q_ids[:, :max_q]).to(dtype),
+        attention_mask=torch.ones(1, max_q, device=device), past_key_values=cache,
+        use_cache=True, cache_position=torch.arange(max_q, device=device))
+    return cache
+
+
+def optimize_sp(llm, embed, q_ids, a_ids, teacher, c0, c1, max_q, device, dtype, cfg,
+                max_norm, raw_window, steps, lr, seed, S, T):
+    g = torch.Generator(device=device).manual_seed(seed)
+    cur_C = c1 - c0
+    R = min(c0, raw_window)
+    raw_emb = embed(a_ids[:, c0 - R:c0]).to(dtype) if R > 0 else None
+    chunk_emb = embed(a_ids[:, c0:c1]).to(dtype)
+    n_new = S + R + cur_C
+    cache_pos = torch.arange(max_q, max_q + n_new, device=device)
+    t_p = F.softmax(teacher[:, c0:c1, :].float() / T, dim=-1)
+    # random init scaled to ~ target_norm
+    sp = torch.randn(1, S, cfg.hidden_dim, generator=g, device=device)
+    sp = sp / sp.norm(dim=-1, keepdim=True).clamp(min=1e-6) * cfg.target_norm
+    sp = sp.requires_grad_(True)
+    opt = torch.optim.Adam([sp], lr=lr)
+    best_kl = float("inf"); best_v = None
+    for _m in range(steps):
+        opt.zero_grad(set_to_none=True)
+        vn = sp.norm(dim=-1, keepdim=True).clamp(min=1e-6)
+        v_c = (sp * torch.where(vn > max_norm, max_norm / vn, torch.ones_like(vn))).to(dtype)
+        v_in = torch.cat([v_c, raw_emb, chunk_emb], dim=1) if R > 0 else torch.cat([v_c, chunk_emb], dim=1)
+        cache = prefill_query(llm, embed, q_ids, max_q, device, dtype)
+        out = llm(inputs_embeds=v_in, attention_mask=torch.ones(1, max_q + n_new, device=device),
+                  past_key_values=cache, position_ids=cache_pos.unsqueeze(0), use_cache=True,
+                  cache_position=cache_pos)
+        v_logp = F.log_softmax(out.logits[:, S - 1 + R: S - 1 + R + cur_C, :].float() / T, dim=-1)
+        kl = (t_p * (t_p.clamp_min(1e-9).log() - v_logp)).sum(-1).mean() * (T * T)
+        (kl + soft_prompt_stability_loss(sp, cfg)).backward()
+        opt.step()
+        if kl.item() < best_kl:
+            best_kl = kl.item()
+            best_v = (sp.detach() * torch.where(vn > max_norm, max_norm / vn, torch.ones_like(vn))).clone()
+    return best_kl, best_v  # best_v = clamped effective SP* (1,S,H)
+
+
+def load_samples(path, tok, cfg, n_want, min_chars, max_ans_len, min_tok):
+    out = []
+    with open(path) as f:
+        for line in f:
+            if len(out) >= n_want:
+                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=max_ans_len, 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=6)
+    p.add_argument("--skip_samples", type=int, default=500)
+    p.add_argument("--min_chars", type=int, default=1500)
+    p.add_argument("--min_tok", type=int, default=400)
+    p.add_argument("--max_ans_len", type=int, default=512)
+    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=150)
+    p.add_argument("--lr", type=float, default=0.03)
+    p.add_argument("--kl_temperature", type=float, default=1.0)
+    args = p.parse_args()
+
+    device = torch.device("cuda"); dtype = torch.bfloat16
+    cfg = Config(); cfg.base_model = args.base_model; cfg.kl_temperature = args.kl_temperature
+    T = args.kl_temperature; C = args.chunk_size; S = cfg.num_soft_tokens
+    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
+    print(f"target_norm={cfg.target_norm:.3f} max_norm={max_norm:.3f}", flush=True)
+
+    depths = [int(x) for x in args.depths.split(",")]
+    alls = load_samples(args.data, tok, cfg, args.skip_samples + args.n_samples, args.min_chars,
+                        args.max_ans_len, args.min_tok)
+    samples = alls[args.skip_samples: args.skip_samples + args.n_samples]
+    print(f"samples={len(samples)} depths={depths} steps={args.steps} lr={args.lr}\n", flush=True)
+
+    sumL2 = 0.0; sumKL = 0.0; cnt = 0; sum_selfE = 0.0
+    for si, (qi, ai) in enumerate(samples):
+        q = torch.tensor([qi], device=device); a = torch.tensor([ai], device=device)
+        max_q = q.size(1); max_a = a.size(1)
+        teacher = teacher_dist(llm, embed, q, a, max_q, max_a, device, dtype)
+        for c0 in depths:
+            c1 = min(c0 + C, max_a)
+            if c1 - c0 < 4 or c0 + 1 >= max_a:
+                continue
+            kl_a, sp_a = optimize_sp(llm, embed, q, a, teacher, c0, c1, max_q, device, dtype, cfg,
+                                     max_norm, args.raw_window, args.steps, args.lr, 111, S, T)
+            kl_b, sp_b = optimize_sp(llm, embed, q, a, teacher, c0, c1, max_q, device, dtype, cfg,
+                                     max_norm, args.raw_window, args.steps, args.lr, 999, S, T)
+            l2 = ((sp_a - sp_b) ** 2).sum(-1).mean().item()      # per-token cross-seed L2
+            selfE = (sp_a ** 2).sum(-1).mean().item()            # per-token energy of a target
+            na = sp_a.norm(dim=-1).mean().item(); nb = sp_b.norm(dim=-1).mean().item()
+            cos = F.cosine_similarity(sp_a, sp_b, dim=-1).mean().item()
+            print(f"  s{si+1} c0={c0}: KL_a={kl_a:.4f} KL_b={kl_b:.4f} | "
+                  f"crossL2={l2:.4f} (||tgt||^2~{selfE:.2f}) cos={cos:.3f} norms=({na:.2f},{nb:.2f})",
+                  flush=True)
+            sumL2 += l2; sumKL += 0.5 * (kl_a + kl_b); sum_selfE += selfE; cnt += 1
+        del teacher
+        torch.cuda.empty_cache()
+
+    if cnt:
+        print("\n" + "=" * 70)
+        print(f"mean KL (both seeds reach): {sumKL/cnt:.4f}   <- should be LOW (both correct)")
+        print(f"mean cross-seed L2:         {sumL2/cnt:.4f}   <- the MSE floor from non-uniqueness")
+        print(f"mean ||target||^2/token:    {sum_selfE/cnt:.4f}")
+        print("=" * 70)
+        print("If KL low AND crossL2 ~ the regression floor (~1.1): targets are NON-UNIQUE")
+        print("  (soft-prompt symmetry) -> MSE-to-vector is the wrong loss; use KL distillation.")
+        print("If crossL2 ~ 0: targets unique -> regression floor is capacity/optimization.")
+
+
+if __name__ == "__main__":
+    main()