Upload v24_grpo_chain_v4.py with huggingface_hub

v24_grpo_chain_v4.py

@@ -0,0 +1,364 @@
+import argparse, json, math, os, random, re
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from datasets import load_dataset
+from transformers import AutoTokenizer
+
+from scugnizz_core import ScugnizzDecoder, preset_config, load_checkpoint
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--model-size", default="1.7b")
+    p.add_argument("--tokenizer", required=True)
+    p.add_argument("--checkpoint", required=True)
+    p.add_argument("--dataset", required=True)
+    p.add_argument("--max-steps", type=int, default=30)
+    p.add_argument("--group-size", type=int, default=4)
+    p.add_argument("--max-new", type=int, default=64)
+    p.add_argument("--temperature", type=float, default=0.8)
+    p.add_argument("--top-p", type=float, default=0.95)
+    p.add_argument("--lr", type=float, default=1e-6)
+    p.add_argument("--kl-beta", type=float, default=0.10)
+    p.add_argument("--block-size", type=int, default=4096)
+    p.add_argument("--save-every", type=int, default=10)
+    p.add_argument("--out-dir", default="/tmp/V24-GRPO-chain-v3")
+    return p.parse_args()
+
+
+def make_prompt(tok, messages):
+    # Use the exact V22 tokenizer chat template instead of reconstructing
+    # headers manually. This keeps the GRPO prompt aligned with SFT.
+    msgs = []
+    for m in messages:
+        role = m.get("role", "")
+        content = str(m.get("content", ""))
+        if role in ("system", "user", "assistant", "ipython", "tool"):
+            if role == "tool":
+                role = "ipython"
+            msgs.append({"role": role, "content": content})
+
+    try:
+        return tok.apply_chat_template(
+            msgs,
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+    except Exception as exc:
+        raise RuntimeError(
+            "Tokenizer chat_template could not render the V22 chain prompt. "
+            f"Roles={[m['role'] for m in msgs]!r}"
+        ) from exc
+
+
+def extract_chain_rows(ds):
+    rows = []
+    for r in ds:
+        msgs = r.get("messages") or []
+        expected = r.get("expected_second_tool")
+        expected_args = r.get("expected_second_args") or {}
+        if not expected or not isinstance(msgs, list):
+            continue
+
+        # Require: assistant first tool_call -> ipython/tool response.
+        first = None
+        response = None
+        for i, m in enumerate(msgs):
+            if m.get("role") == "assistant" and "<tool_call>" in str(m.get("content") or ""):
+                if i + 1 < len(msgs) and msgs[i + 1].get("role") in ("ipython", "tool"):
+                    first = i
+                    response = i + 1
+                    break
+        if first is None:
+            continue
+
+        prefix = msgs[: response + 1]
+        rows.append((prefix, str(expected), expected_args))
+    return rows
+
+
+def extract_tool_call(text):
+    # Primary V22 format: explicit <tool_call>{...}</tool_call>.
+    m = re.search(r"<tool_call>\\s*(\\{.*?\\})\\s*</tool_call>", text, flags=re.S)
+    if m:
+        try:
+            obj = json.loads(m.group(1))
+            if isinstance(obj, dict):
+                return obj
+        except Exception:
+            pass
+
+    # Robust fallback: recover a complete JSON tool object even if the model
+    # polluted the opening tag with extra BPE tokens/whitespace.
+    for m in re.finditer(r'\\{\\s*"name"\\s*:\\s*"([^"]+)"\\s*,\\s*"arguments"\\s*:\\s*\\{', text):
+        start = m.start()
+        depth = 0
+        in_str = False
+        esc = False
+        for k in range(start, len(text)):
+            c = text[k]
+            if in_str:
+                if esc:
+                    esc = False
+                elif c == "\\\\":
+                    esc = True
+                elif c == '"':
+                    in_str = False
+            else:
+                if c == '"':
+                    in_str = True
+                elif c == "{":
+                    depth += 1
+                elif c == "}":
+                    depth -= 1
+                    if depth == 0:
+                        try:
+                            obj = json.loads(text[start:k+1])
+                            if isinstance(obj, dict):
+                                return obj
+                        except Exception:
+                            break
+    return None
+
+
+def reward_chain(text, expected_name, expected_args):
+    call = extract_tool_call(text)
+
+    # Plain answers are neutral: do not punish the model merely for failing
+    # to emit a tool call. This prevents four identical -0.5 rewards.
+    if call is None:
+        return 0.0
+
+    name = call.get("name")
+    args = call.get("arguments")
+
+    # Valid tool structure gives a small positive base.
+    score = 0.25
+
+    if name == expected_name:
+        score += 2.0
+    elif isinstance(name, str) and name:
+        score -= 0.25
+
+    if isinstance(args, dict):
+        if expected_args and args == expected_args:
+            score += 1.0
+        elif expected_args:
+            hits = sum(args.get(k) == v for k, v in expected_args.items())
+            if hits:
+                score += 0.50 * hits / len(expected_args)
+    else:
+        score -= 0.50
+
+    # Header contamination is a mild penalty, not a complete rejection:
+    # valid tool calls should retain most of their structural reward.
+    contamination_tokens = (
+        "<|start_header_id|>",
+        "<|end_header_id|>",
+        "<|assistant|>",
+        "<|user|>",
+        "<|system|>",
+        "<|ipython|>",
+    )
+    if any(x in text for x in contamination_tokens):
+        score -= 0.25
+
+    return float(score)
+
+
+@torch.no_grad()
+def sample_group(model, tok, prompt_ids, group_size, max_new, temperature, top_p):
+    device = next(model.parameters()).device
+    outs = []
+    for _ in range(group_size):
+        ids = prompt_ids.clone().to(device)
+        for _ in range(max_new):
+            x = ids[:, -4096:]
+            logits = model(x)
+            if isinstance(logits, tuple):
+                logits = logits[0]
+            if hasattr(logits, "logits"):
+                logits = logits.logits
+            next_logits = logits[:, -1, :] / max(temperature, 1e-6)
+            probs = F.softmax(next_logits, dim=-1)
+            if top_p < 1.0:
+                sorted_probs, sorted_idx = torch.sort(probs, descending=True)
+                cum = torch.cumsum(sorted_probs, dim=-1)
+                mask = cum - sorted_probs > top_p
+                sorted_probs[mask] = 0
+                sorted_probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True)
+                nxt = sorted_idx.gather(-1, torch.multinomial(sorted_probs, 1))
+            else:
+                nxt = torch.multinomial(probs, 1)
+            ids = torch.cat([ids, nxt], dim=1)
+            if nxt.item() == tok.eos_token_id:
+                break
+        outs.append(ids[:, prompt_ids.shape[1]:].clone())
+    return outs
+
+
+def seq_logprob(model, prompt_ids, response_ids, block_size):
+    device = next(model.parameters()).device
+    full = torch.cat([prompt_ids, response_ids], dim=1)
+    full = full[:, -block_size:]
+    # Keep response-token indexing aligned after any left truncation.
+    plen = max(full.shape[1] - response_ids.shape[1], 0)
+    logits = model(full)
+    if isinstance(logits, tuple):
+        logits = logits[0]
+    if hasattr(logits, "logits"):
+        logits = logits.logits
+    start = max(plen - 1, 0)
+    target = full[:, start + 1:]
+    pred = logits[:, start:-1, :]
+    if target.numel() == 0:
+        return torch.zeros((), device=device)
+    lp = F.log_softmax(pred.float(), dim=-1)
+    tok_lp = lp.gather(-1, target.unsqueeze(-1)).squeeze(-1)
+    return tok_lp.mean()
+
+
+def seq_logprob_and_forward_kl(model, ref, prompt_ids, response_ids, block_size):
+    device = next(model.parameters()).device
+    full = torch.cat([prompt_ids, response_ids], dim=1)
+    if full.shape[1] > block_size:
+        full = full[:, -block_size:]
+    plen = max(full.shape[1] - response_ids.shape[1], 0)
+
+    logits = model(full)
+    if isinstance(logits, tuple):
+        logits = logits[0]
+    if hasattr(logits, "logits"):
+        logits = logits.logits
+
+    with torch.no_grad():
+        ref_logits = ref(full)
+        if isinstance(ref_logits, tuple):
+            ref_logits = ref_logits[0]
+        if hasattr(ref_logits, "logits"):
+            ref_logits = ref_logits.logits
+
+    start = max(plen - 1, 0)
+    target = full[:, start + 1:]
+    pred = logits[:, start:-1, :]
+    ref_pred = ref_logits[:, start:-1, :]
+
+    if target.numel() == 0:
+        z = torch.zeros((), device=device)
+        return z, z
+
+    logp = F.log_softmax(pred.float(), dim=-1)
+    with torch.no_grad():
+        logq = F.log_softmax(ref_pred.float(), dim=-1)
+        q = logq.exp()
+
+    tok_logp = logp.gather(-1, target.unsqueeze(-1)).squeeze(-1)
+    token_kl = (q * (logq - logp)).sum(dim=-1)
+    return tok_logp.mean(), token_kl.mean()
+
+
+def save_state(model, optimizer, step, out_dir):
+    Path(out_dir).mkdir(parents=True, exist_ok=True)
+    state = {
+        "step": step,
+        "model": model.state_dict(),
+        "optimizer": optimizer.state_dict(),
+    }
+    torch.save(state, Path(out_dir) / f"checkpoint_step_{step:03d}.pt")
+    torch.save(state, Path(out_dir) / "checkpoint_last.pt")
+
+
+def main():
+    args = parse_args()
+    random.seed(20260915)
+    torch.manual_seed(20260915)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    dtype = torch.bfloat16 if device.type == "cuda" and torch.cuda.is_bf16_supported() else torch.float32
+    print(f"device={device} dtype={dtype}")
+
+    tok = AutoTokenizer.from_pretrained(args.tokenizer)
+    ds = load_dataset(args.dataset, split="train")
+    rows = extract_chain_rows(ds)
+    print(f"chain_rows={len(rows)} lr={args.lr:g} kl_beta={args.kl_beta:g}")
+    if len(rows) < args.group_size:
+        raise RuntimeError("Too few valid sequential chain rows")
+
+    cfg = preset_config(args.model_size, vocab_size=len(tok), block_size=args.block_size, dropout=0.0, pcs_a=0.8309193524, pcs_b=0.0)
+    model = ScugnizzDecoder(cfg).to(device=device, dtype=dtype)
+    state = load_checkpoint(args.checkpoint, "cpu")
+    if isinstance(state, dict) and "model" in state:
+        state = state["model"]
+    model.load_state_dict(state, strict=True)
+    model.train()
+
+    # Frozen V22 reference for a simple KL-style anchor.
+    ref = ScugnizzDecoder(cfg).to(device=device, dtype=dtype)
+    ref.load_state_dict(state, strict=True)
+    ref.eval()
+    for p in ref.parameters():
+        p.requires_grad_(False)
+
+    opt = torch.optim.AdamW(model.parameters(), lr=args.lr)
+    step = 0
+    for step in range(1, args.max_steps + 1):
+        prefix, expected, expected_args = rows[(step - 1) % len(rows)]
+        prompt = make_prompt(tok, prefix)
+        prompt_ids = tok(prompt, return_tensors="pt", add_special_tokens=False).input_ids.to(device)
+
+        # Sample a group from the state immediately after the first tool response.
+        responses = sample_group(
+            model, tok, prompt_ids, args.group_size, args.max_new,
+            args.temperature, args.top_p
+        )
+        texts = [tok.decode(x[0], skip_special_tokens=False) for x in responses]
+        rewards = torch.tensor(
+            [reward_chain(t, expected, expected_args) for t in texts],
+            device=device, dtype=torch.float32
+        )
+
+        mean_r = rewards.mean()
+        std_r = rewards.std(unbiased=False)
+        adv = (rewards - mean_r) / (std_r + 1e-6)
+
+        # Recompute differentiable sequence log-probs and a true sampled
+        # forward KL against the frozen V22 reference.
+        logps = []
+        kls = []
+        for rids in responses:
+            lp, k = seq_logprob_and_forward_kl(
+                model, ref, prompt_ids, rids.to(device), args.block_size
+            )
+            logps.append(lp)
+            kls.append(k)
+        logps = torch.stack(logps)
+        kls = torch.stack(kls)
+        kl = kls.mean()
+        loss = -(adv.detach() * logps).mean() + args.kl_beta * kl
+
+        opt.zero_grad(set_to_none=True)
+        if loss.requires_grad and torch.isfinite(loss) and torch.isfinite(kl):
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            opt.step()
+
+        print(
+            f"step={step:03d} expected={expected} "
+            f"reward_mean={mean_r.item():+.3f} reward_max={rewards.max().item():+.3f} "
+            f"adv_std={adv.std(unbiased=False).item():.3f} "
+            f"loss={loss.item():+.5f} kl={kl.item():+.5f}"
+        )
+        for i, (t, r) in enumerate(zip(texts, rewards.tolist())):
+            print(f"  sample[{i}] reward={r:+.2f} {t[:500]!r}")
+
+        if step % args.save_every == 0:
+            save_state(model, opt, step, args.out_dir)
+
+    save_state(model, opt, step, args.out_dir)
+    print("DONE")
+
+
+if __name__ == "__main__":
+    main()