Upload overfit_e2e.py with huggingface_hub

overfit_e2e.py

@@ -0,0 +1,185 @@
+#!/usr/bin/env python3
+"""
+END-TO-END OVERFIT: can a SINGLE learnable recurrent step() keep KL low across a whole
+sequence?  Train the hypernet end-to-end (KL through the recurrence, FULL BPTT over all
+chunks) on a few FIXED sequences; log per-chunk KL (esp. deep chunks) over training.
+
+This is a SYMMETRY-INVARIANT existence test (output/KL space, never touches a specific
+SP* vector):
+  deep-chunk KL -> low  => a recurrent SP trajectory EXISTS and is findable by gradient
+                           descent (the floor was amortization/generalization -> fixable).
+  deep-chunk KL stuck   => the recurrent structure genuinely cannot carry the history.
+"""
+import sys
+sys.path.insert(0, "/workspace")
+import argparse, json, time
+import torch
+import torch.nn as nn
+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 load_samples(path, tok, cfg, n_want, skip, min_chars, ans_len):
+    out = []
+    with open(path) as f:
+        for line in f:
+            if len(out) >= skip + 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=ans_len, truncation=True, add_special_tokens=False).input_ids
+            if len(ai) < ans_len:        # need full-length sequences (no padding/masking)
+                continue
+            out.append((qi, ai))
+    return out[skip:skip + n_want]
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--init_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_seq", type=int, default=2)
+    p.add_argument("--skip", type=int, default=400)
+    p.add_argument("--min_chars", type=int, default=2500)
+    p.add_argument("--ans_len", type=int, default=512)
+    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=400)
+    p.add_argument("--lr", type=float, default=1e-3)
+    p.add_argument("--kl_temperature", type=float, default=1.0)
+    p.add_argument("--from_scratch", action="store_true")
+    args = p.parse_args()
+
+    device = torch.device("cuda"); dtype = torch.bfloat16
+    torch.backends.cuda.matmul.allow_tf32 = True
+    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 = False
+    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)
+    if not args.from_scratch:
+        ckd = torch.load(args.init_ckpt, map_location="cpu", weights_only=False)
+        hn.load_state_dict(ckd["hypernet"], strict=False)
+        print(f"init from {args.init_ckpt}", flush=True)
+    else:
+        print("init from scratch", flush=True)
+    hn.train()
+    opt = torch.optim.AdamW([pp for pp in hn.parameters() if pp.requires_grad], lr=args.lr, weight_decay=0.0)
+
+    samples = load_samples(args.data, tok, cfg, args.n_seq, args.skip, args.min_chars, args.ans_len)
+    B = len(samples)
+    max_q = max(len(s[0]) for s in samples)
+    max_a = args.ans_len
+    q_ids = torch.zeros(B, max_q, dtype=torch.long, device=device)
+    a_ids = torch.zeros(B, max_a, dtype=torch.long, device=device)
+    q_lens = torch.zeros(B, dtype=torch.long, device=device)
+    for i, (qi, ai) in enumerate(samples):
+        q_ids[i, :len(qi)] = torch.tensor(qi, device=device); q_lens[i] = len(qi)
+        a_ids[i, :max_a] = torch.tensor(ai[:max_a], device=device)
+    pos_q = torch.arange(max_q, device=device).unsqueeze(0).expand(B, -1)
+    valid_q = (pos_q < q_lens.unsqueeze(1)).long()
+    n_chunks = (max_a + C - 1) // C
+    print(f"overfit B={B} seqs | max_q={max_q} max_a={max_a} n_chunks={n_chunks} | lr={args.lr} steps={args.steps}\n", flush=True)
+
+    # ---- teacher (frozen, full real context), precompute once ----
+    with torch.no_grad():
+        cache_t = DynamicCache()
+        q_emb = (embed(q_ids) * valid_q.unsqueeze(-1)).to(dtype)
+        out_q = llm(inputs_embeds=q_emb, attention_mask=valid_q, past_key_values=cache_t,
+                    use_cache=True, cache_position=torch.arange(max_q, device=device))
+        last_q = (q_lens - 1).clamp(min=0)
+        t0 = out_q.logits[torch.arange(B, device=device), last_q, :]
+        a_emb_t = embed(a_ids).to(dtype)
+        attn_qa = torch.cat([valid_q, torch.ones(B, max_a, dtype=torch.long, device=device)], dim=1)
+        pos_a = torch.arange(max_q, max_q + max_a, device=device)
+        out_a = llm(inputs_embeds=a_emb_t, attention_mask=attn_qa, past_key_values=cache_t,
+                    position_ids=pos_a.unsqueeze(0).expand(B, -1), use_cache=True, cache_position=pos_a)
+        V = out_a.logits.size(-1)
+        teacher = torch.empty(B, max_a, V, dtype=out_a.logits.dtype, device=device)
+        teacher[:, 0, :] = t0; teacher[:, 1:, :] = out_a.logits[:, :max_a - 1, :]
+        t_p_all = F.softmax(teacher.float() / T, dim=-1)       # (B,max_a,V)
+        del cache_t, out_q, out_a, teacher
+    torch.cuda.empty_cache()
+
+    def student_step(record_kl=False):
+        cache_s = DynamicCache()
+        with torch.no_grad():
+            llm(inputs_embeds=(embed(q_ids) * valid_q.unsqueeze(-1)).to(dtype), attention_mask=valid_q,
+                past_key_values=cache_s, use_cache=True, cache_position=torch.arange(max_q, device=device))
+        sp = hn.init_sp.expand(B, -1, -1).contiguous()
+        total = 0.0; per_chunk = []
+        for j in range(n_chunks):
+            c0 = j * C; c1 = min(c0 + C, max_a); cur_C = c1 - c0
+            with torch.no_grad():
+                nrm = sp.norm(dim=-1, keepdim=True).clamp(min=1e-6)
+                sc = torch.where(nrm > max_norm, max_norm / nrm, torch.ones_like(nrm))
+            sp_c = (sp * sc).to(dtype)
+            R = min(c0, args.raw_window)
+            chunk_emb = embed(a_ids[:, c0:c1]).to(dtype)
+            parts = [sp_c] + ([embed(a_ids[:, c0 - R:c0]).to(dtype)] if R > 0 else []) + [chunk_emb]
+            llm_in = torch.cat(parts, dim=1); 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)], dim=1)
+            out = llm(inputs_embeds=llm_in, attention_mask=attn, past_key_values=cache_s,
+                      position_ids=cpos.unsqueeze(0).expand(B, -1), use_cache=True, cache_position=cpos)
+            cache_s.crop(max_q)
+            for layer in cache_s.layers:
+                layer.keys = layer.keys.detach().contiguous(); layer.values = layer.values.detach().contiguous()
+            s_logp = F.log_softmax(out.logits[:, S - 1 + R: S - 1 + R + cur_C, :].float() / T, dim=-1)
+            t_p = t_p_all[:, c0:c1, :]
+            kl = (t_p * (t_p.clamp_min(1e-9).log() - s_logp)).sum(-1).mean() * (T * T)
+            total = total + kl + soft_prompt_stability_loss(sp, cfg)
+            if record_kl:
+                per_chunk.append(kl.item())
+            if j < n_chunks - 1:
+                sp = hn.step(sp, chunk_emb.float(), None)       # FULL BPTT (no detach)
+        return total, per_chunk
+
+    t0t = time.time()
+    for st in range(1, args.steps + 1):
+        opt.zero_grad(set_to_none=True)
+        rec = (st % 25 == 0 or st == 1)
+        total, per_chunk = student_step(record_kl=rec)
+        total.backward()
+        nn.utils.clip_grad_norm_([pp for pp in hn.parameters() if pp.requires_grad], 1.0)
+        opt.step()
+        torch.cuda.empty_cache()
+        if rec:
+            pcs = " ".join(f"{k:.3f}" for k in per_chunk)
+            deep = sum(per_chunk[3:]) / max(len(per_chunk[3:]), 1)
+            print(f"step {st:4d} | meanKL={sum(per_chunk)/len(per_chunk):.4f} deepKL(c0>=192)={deep:.4f} "
+                  f"| per-chunk[{pcs}] | {time.time()-t0t:.0f}s", flush=True)
+    print("\nDONE. If deepKL dropped low, a recurrent SP trajectory EXISTS & is learnable.")
+
+
+if __name__ == "__main__":
+    main()