Upload v24_grpo_chain_v6.py with huggingface_hub

v24_grpo_chain_v6.py

@@ -0,0 +1,410 @@
+
+import argparse, copy, json, math, random, re, subprocess, sys
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from transformers import AutoTokenizer
+from datasets import load_dataset
+
+# scugnizz_core.py is a truncated copy of scugnizz-llama.py containing
+# ScugnizzConfig / ScugnizzDecoder / preset_config / load_checkpoint.
+from scugnizz_core import ScugnizzDecoder, preset_config, load_checkpoint
+
+
+TOOL_RE = re.compile(r"<tool_call>\s*(\{.*?\})\s*</tool_call>", re.S)
+HEADER_RE = re.compile(r"<\|start_header_id\|>\s*(?:user|assistant|ipython|system)\s*<\|end_header_id\|>")
+
+def extract_tool_calls(text):
+    text = text or ""
+    calls = []
+    for m in TOOL_RE.finditer(text):
+        raw = m.group(1)
+        try:
+            obj = json.loads(raw)
+            if isinstance(obj, dict) and isinstance(obj.get("name"), str):
+                calls.append((m.start(), m.end(), obj, True))
+        except Exception:
+            pass
+
+    # Fallback for truncated/malformed <tool_call> wrappers: recover a balanced JSON object.
+    for m in re.finditer(r'\{\s*"name"\s*:\s*"([^"\\]+)"\s*,\s*"arguments"\s*:\s*\{', text):
+        start = m.start()
+        if any(a <= start < b for a, b, _, _ in calls):
+            continue
+        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) and isinstance(obj.get("name"), str):
+                                calls.append((start, k+1, obj, False))
+                        except Exception:
+                            pass
+                        break
+    calls.sort(key=lambda x: x[0])
+    return calls
+
+
+def _arg_match(got, gold):
+    if not isinstance(got, dict) or not isinstance(gold, dict):
+        return 0.0
+    if got == gold:
+        return 1.5
+    if not got:
+        return 0.0
+    hits = 0
+    for k, v in gold.items():
+        if k in got and got[k] == v:
+            hits += 1
+    return 0.75 * (hits / max(1, len(gold))) if hits else 0.0
+
+
+def reward_chain(text, expected):
+    """V6 reward: score the FIRST generated tool call, not any later echoed call.
+    expected may be a tool name (legacy) or the gold call dict (preferred).
+    """
+    text = text or ""
+    calls = extract_tool_calls(text)
+    if isinstance(expected, dict):
+        gold_name = expected.get("name", "")
+        gold_args = expected.get("arguments", {})
+    else:
+        gold_name = str(expected or "")
+        gold_args = {}
+
+    if not calls:
+        # Distinguish pure natural-language output from malformed tool syntax.
+        return -1.5 if "tool_call" in text else -1.0
+
+    pos, end, call, wrapped = calls[0]
+    name = call.get("name", "")
+    r = 0.0
+
+    # The target tool must be the FIRST generated action.
+    if name == gold_name:
+        r += 2.0
+        r += _arg_match(call.get("arguments"), gold_args)
+    else:
+        r -= 1.0
+
+    if not wrapped:
+        r -= 0.75
+
+    # Any earlier non-whitespace prose means the model did not enter tool mode cleanly.
+    prefix = text[:pos].strip()
+    if prefix:
+        r -= 0.5
+
+    # A second tool call is not a successful single action: it is usually the
+    # corruption/echo pattern observed in V5.
+    if len(calls) > 1:
+        r -= 0.75 * min(3, len(calls) - 1)
+
+    # Chat-template headers inside a sampled completion are contamination.
+    headers = len(HEADER_RE.findall(text))
+    if headers:
+        r -= 1.25 * min(3, headers)
+
+    # If the model emits a non-empty continuation after the tool call, tolerate
+    # the terminal token but penalize generated tool results / prose.
+    tail = text[end:].strip()
+    if tail:
+        tail = re.sub(r"^<\|eot_id\|>", "", tail).strip()
+    if tail:
+        r -= 0.50
+
+    return float(max(-3.0, min(3.0, r)))
+
+
+def parse_tool_name(text):
+    calls = extract_tool_calls(text)
+    return calls[0][2].get("name") if calls else None
+
+def make_prompt(tok, messages):
+    # messages already contain system/user/assistant/tool history.
+    # sft_data's ROLE_MAP maps tool -> tool and the Llama template renders
+    # that as the ipython header. We reproduce that exact representation.
+    msgs = []
+    for m in messages:
+        role = m["role"]
+        if role == "tool":
+            role = "ipython"
+        msgs.append({"role": role, "content": m["content"]})
+    return tok.apply_chat_template(
+        msgs, tokenize=False, add_generation_prompt=True
+    )
+
+
+def extract_chain_rows(ds):
+    rows = []
+    for row in ds:
+        conv = row.get("conversations") or row.get("messages") or []
+        # Normalize role/value -> role/content.
+        norm = []
+        for m in conv:
+            role = m.get("role") or m.get("from")
+            content = m.get("content")
+            if content is None:
+                content = m.get("value", "")
+            if role in ("human", "user"):
+                role = "user"
+            elif role in ("gpt", "assistant", "reasoning", "tool_call"):
+                role = "assistant"
+            elif role == "tool":
+                role = "tool"
+            elif role == "system":
+                role = "system"
+            else:
+                continue
+            norm.append({"role": role, "content": content})
+
+        # Find every assistant tool call that is preceded by a tool response.
+        for i, m in enumerate(norm):
+            if m["role"] != "assistant" or "<tool_call>" not in m["content"]:
+                continue
+            if i == 0 or norm[i-1]["role"] != "tool":
+                continue
+            gold_calls = extract_tool_calls(m["content"])
+            if not gold_calls:
+                continue
+            expected = gold_calls[0][2]
+
+            # Prompt ends immediately after the previous tool response.
+            prefix = norm[:i]
+            if not prefix or prefix[-1]["role"] != "tool":
+                continue
+
+            # Keep only reasonably sized contexts; left-truncate if needed later.
+            rows.append((prefix, expected))
+
+    return rows
+
+
+@torch.no_grad()
+def sample_group(model, tok, prompt, n, max_new, temperature, top_p, device, block_size):
+    model.eval()
+    ids = tok.encode(prompt, add_special_tokens=False)
+    room = max(32, block_size - max_new)
+    ids = ids[-room:]
+    x = torch.tensor([ids], dtype=torch.long, device=device)
+    out = []
+
+    # Seed variation comes from multinomial sampling.
+    for _ in range(n):
+        past = None
+        cur = x
+        generated = []
+        for _step in range(max_new):
+            logits, _, past = model(cur, past_kvs=past, use_cache=True)
+            logits = logits[:, -1, :] / max(1e-5, temperature)
+            probs = F.softmax(logits, dim=-1)
+
+            if top_p < 1.0:
+                sp, si = torch.sort(probs, descending=True)
+                cs = torch.cumsum(sp, dim=-1)
+                mask = cs > top_p
+                mask[..., 0] = False
+                sp = sp.masked_fill(mask, 0)
+                sp = sp / sp.sum(dim=-1, keepdim=True)
+                nxt = si.gather(-1, torch.multinomial(sp, 1))
+            else:
+                nxt = torch.multinomial(probs, 1)
+
+            tid = int(nxt.item())
+            generated.append(tid)
+
+            stop_ids = set()
+            if tok.eos_token_id is not None:
+                stop_ids.add(int(tok.eos_token_id))
+            for s in ("<|eot_id|>", "<|eom_id|>"):
+                t = tok.convert_tokens_to_ids(s)
+                if isinstance(t, int) and t >= 0:
+                    stop_ids.add(t)
+            if tid in stop_ids:
+                break
+            cur = nxt
+
+        out.append(generated)
+    return out, ids
+
+
+def sequence_logprobs(model, full_ids, prompt_len, response_lens):
+    """Differentiable mean log p(sampled response tokens), excluding padding."""
+    logits, _ = model(full_ids)
+    lp = F.log_softmax(logits[:, :-1, :], dim=-1)
+    targets = full_ids[:, 1:]
+    start = max(0, prompt_len - 1)
+    chosen = lp[:, start:, :].gather(
+        -1, targets[:, start:].unsqueeze(-1)
+    ).squeeze(-1)
+    mask = torch.zeros_like(chosen)
+    for i, n in enumerate(response_lens):
+        mask[i, :int(n)] = 1.0
+    sums = (chosen * mask).sum(dim=-1)
+    means = sums / mask.sum(dim=-1).clamp_min(1.0)
+    return sums, means
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model-size", default="1.7b")
+    ap.add_argument("--tokenizer", required=True)
+    ap.add_argument("--checkpoint", required=True)
+    ap.add_argument("--dataset", required=True)
+    ap.add_argument("--max-steps", type=int, default=30)
+    ap.add_argument("--group-size", type=int, default=4)
+    ap.add_argument("--max-new", type=int, default=64)
+    ap.add_argument("--temperature", type=float, default=0.8)
+    ap.add_argument("--top-p", type=float, default=0.95)
+    ap.add_argument("--lr", type=float, default=5e-6)
+    ap.add_argument("--kl-beta", type=float, default=0.02)
+    ap.add_argument("--block-size", type=int, default=4096)
+    ap.add_argument("--save-every", type=int, default=10)
+    ap.add_argument("--out-dir", default="/tmp/V24-GRPO-chain-smoke")
+    args = ap.parse_args()
+
+    random.seed(42)
+    torch.manual_seed(42)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.bfloat16 if device == "cuda" else torch.float32
+    out = Path(args.out_dir)
+    out.mkdir(parents=True, exist_ok=True)
+
+    print("V24 GRPO CHAIN V6 STRUCTURAL REWARD", flush=True)
+    print(f"device={device} dtype={dtype}", flush=True)
+
+    tok = AutoTokenizer.from_pretrained(args.tokenizer)
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+    if not getattr(tok, "chat_template", None):
+        raise RuntimeError("Tokenizer has no chat_template")
+
+    ds = load_dataset(args.dataset, split="train")
+    rows = extract_chain_rows(ds)
+    if len(rows) < args.group_size:
+        raise RuntimeError(f"Only {len(rows)} chain rows found")
+    random.shuffle(rows)
+    print(f"chain_rows={len(rows)}", flush=True)
+
+    cfg = preset_config(
+        args.model_size, len(tok), args.block_size, 0.0,
+        0.8309193524478643, 0.0, gradient_checkpointing=False
+    )
+    model = ScugnizzDecoder(cfg).to(device=device, dtype=dtype)
+    ck = load_checkpoint(args.checkpoint, "cpu")
+    state = ck["model"] if isinstance(ck, dict) and "model" in ck else ck
+    model.load_state_dict(state, strict=True)
+    del ck
+    if device == "cuda":
+        torch.cuda.empty_cache()
+
+    # Frozen reference = exact V22-best before GRPO.
+    ref = copy.deepcopy(model).eval()
+    for p in ref.parameters():
+        p.requires_grad_(False)
+
+    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.95),
+                            weight_decay=0.0)
+
+    step = 0
+    for step in range(1, args.max_steps + 1):
+        prefix, expected = rows[(step - 1) % len(rows)]
+        prompt = make_prompt(tok, prefix)
+
+        sampled, prompt_ids = sample_group(
+            model, tok, prompt, args.group_size, args.max_new,
+            args.temperature, args.top_p, device, args.block_size
+        )
+        prompt_len = len(prompt_ids)
+
+        texts = [tok.decode(x, skip_special_tokens=True) for x in sampled]
+        rewards = torch.tensor(
+            [reward_chain(t, expected) for t in texts],
+            device=device, dtype=torch.float32
+        )
+
+        # GRPO group-relative advantage.
+        adv = (rewards - rewards.mean()) / (rewards.std(unbiased=False) + 1e-4)
+
+        full = []
+        for gen in sampled:
+            ids = prompt_ids + gen
+            ids = ids[-args.block_size:]
+            full.append(ids)
+
+        # Pad group to a rectangular batch. Padding is outside the scored span.
+        maxlen = max(len(x) for x in full)
+        pad = tok.pad_token_id
+        batch = []
+        response_lens = [len(g) for g in sampled]
+        for ids in full:
+            batch.append(ids + [pad] * (maxlen - len(ids)))
+        x = torch.tensor(batch, dtype=torch.long, device=device)
+
+        # Recompute current-policy log-probs with gradients.
+        cur_sum, cur_mean = sequence_logprobs(
+            model, x, prompt_len, response_lens
+        )
+
+        with torch.no_grad():
+            ref_logits, _ = ref(x)
+            ref_lp = F.log_softmax(ref_logits[:, :-1, :], dim=-1)
+            targets = x[:, 1:]
+            start = max(0, prompt_len - 1)
+            ref_tok = ref_lp[:, start:, :].gather(
+                -1, targets[:, start:].unsqueeze(-1)
+            ).squeeze(-1)
+            ref_mask = torch.zeros_like(ref_tok)
+            for i, n in enumerate(response_lens):
+                ref_mask[i, :int(n)] = 1.0
+            ref_mean = (ref_tok * ref_mask).sum(dim=-1) / ref_mask.sum(dim=-1).clamp_min(1.0)
+            # Sampled-token KL proxy; gradients flow through cur_mean only.
+            kl = cur_mean - ref_mean
+
+        # Normalize by response length for stable small-batch updates.
+        policy_loss = -(adv.detach() * cur_mean).mean()
+        kl_loss = args.kl_beta * kl.mean()
+        loss = policy_loss + kl_loss
+
+        opt.zero_grad(set_to_none=True)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        opt.step()
+
+        print(
+            f"step={step:03d} expected={expected} "
+            f"reward_mean={rewards.mean().item():+.3f} "
+            f"reward_max={rewards.max().item():+.3f} "
+            f"adv_std={adv.std(unbiased=False).item():.3f} "
+            f"loss={loss.item():+.5f} kl={kl.mean().item():+.5f}",
+            flush=True
+        )
+        for j, (t, r) in enumerate(zip(texts, rewards.tolist())):
+            print(f"  sample{j}: r={r:+.2f} {t[:180]!r}", flush=True)
+
+        if step % args.save_every == 0 or step == args.max_steps:
+            torch.save(
+                {"model": model.state_dict(), "step": step},
+                out / f"checkpoint_step_{step:03d}.pt"
+            )
+
+    torch.save({"model": model.state_dict(), "step": step},
+               out / "checkpoint_last.pt")
+    print(f"DONE out={out}", flush=True)
+
+
+if __name__ == "__main__":
+    main()
+
+