Add kaggle/train_grpo.py

kaggle/train_grpo.py

@@ -0,0 +1,603 @@
+"""
+Kaggle GRPO trainer for deepagents-rl coding.
+
+Hard requirements:
+  - W&B metrics every step
+  - BEFORE every checkpoint: held-out eval -> Langfuse scores/trace -> THEN save/push HF
+  - HF adapter revisions tagged grpo-step-{n}
+
+Secrets (Kaggle User Secrets OR env):
+  HF_TOKEN, WANDB_API_KEY, LANGFUSE_PUBLIC_KEY, LANGFUSE_SECRET_KEY
+  optional LANGFUSE_HOST / LANGFUSE_BASE_URL
+"""
+from __future__ import annotations
+
+import json
+import os
+import re
+import subprocess
+import sys
+import traceback
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+# ---------------------------------------------------------------------------
+# Inline reward helpers (also shipped in package; duplicated for notebook isolation)
+# ---------------------------------------------------------------------------
+
+
+def extract_python_code(completion: str) -> str:
+    if "```python" in completion:
+        start = completion.find("```python") + len("```python")
+        end = completion.find("```", start)
+        return completion[start:end if end != -1 else None].strip()
+    if "```" in completion:
+        start = completion.find("```") + 3
+        end = completion.find("```", start)
+        if end != -1:
+            return completion[start:end].strip()
+    return completion.strip()
+
+
+def run_hidden_tests(solution_code: str, hidden_tests: list[str], timeout_s: float = 5.0) -> dict[str, Any]:
+    per_test = []
+    passed = 0
+    for i, assertion in enumerate(hidden_tests):
+        script = f"{solution_code}\n{assertion}\n"
+        try:
+            proc = subprocess.run(
+                [sys.executable, "-c", script],
+                capture_output=True,
+                text=True,
+                timeout=timeout_s,
+            )
+            ok = proc.returncode == 0
+            err = (proc.stderr or "")[:300]
+        except Exception as e:  # noqa: BLE001
+            ok = False
+            err = f"{type(e).__name__}: {e}"
+        per_test.append({"i": i, "ok": ok, "stderr": err})
+        if ok:
+            passed += 1
+    total = len(hidden_tests)
+    return {
+        "pass_frac": (passed / total) if total else 0.0,
+        "passed": passed,
+        "total": total,
+        "per_test": per_test,
+    }
+
+
+def grade_completion(completion: str, hidden_tests: list[str]) -> dict[str, Any]:
+    code = extract_python_code(completion)
+    g = run_hidden_tests(code, hidden_tests)
+    g["extracted_code"] = code
+    return g
+
+
+def coding_reward(completions: list[str], hidden_tests: list[Any], **kwargs: Any) -> list[float]:
+    """TRL GRPO reward: fraction of hidden unit tests passed."""
+    rewards = []
+    for completion, tests in zip(completions, hidden_tests):
+        if isinstance(tests, str):
+            tests = json.loads(tests)
+        g = grade_completion(completion, list(tests))
+        rewards.append(float(g["pass_frac"]))
+    return rewards
+
+
+# ---------------------------------------------------------------------------
+# Secrets
+# ---------------------------------------------------------------------------
+
+
+def load_secrets() -> None:
+    try:
+        from kaggle_secrets import UserSecretsClient
+
+        us = UserSecretsClient()
+        for key in (
+            "HF_TOKEN",
+            "WANDB_API_KEY",
+            "LANGFUSE_PUBLIC_KEY",
+            "LANGFUSE_SECRET_KEY",
+            "LANGFUSE_HOST",
+            "LANGFUSE_BASE_URL",
+        ):
+            try:
+                val = us.get_secret(key)
+                if val:
+                    os.environ[key] = val
+            except Exception:
+                pass
+    except Exception:
+        pass
+
+    missing = [
+        k
+        for k in (
+            "HF_TOKEN",
+            "WANDB_API_KEY",
+            "LANGFUSE_PUBLIC_KEY",
+            "LANGFUSE_SECRET_KEY",
+        )
+        if not os.getenv(k)
+    ]
+    if missing:
+        raise RuntimeError(
+            f"Missing required secrets/env: {missing}. "
+            "Add them in Kaggle Notebook Secrets (Add-ons -> Secrets)."
+        )
+    os.environ.setdefault("LANGFUSE_HOST", "https://cloud.langfuse.com")
+    os.environ.setdefault("LANGFUSE_BASE_URL", os.environ["LANGFUSE_HOST"])
+    os.environ.setdefault("WANDB_PROJECT", "deepagents-rl-coding")
+    os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "0")
+
+
+# ---------------------------------------------------------------------------
+# Model resolve
+# ---------------------------------------------------------------------------
+
+CANDIDATES = [
+    "Qwen/Qwen2.5-Coder-1.5B-Instruct",
+    "Qwen/Qwen2.5-Coder-0.5B-Instruct",
+    "Qwen/Qwen2.5-1.5B-Instruct",
+    "HuggingFaceTB/SmolLM2-1.7B-Instruct",
+    "google/gemma-2-2b-it",
+]
+
+
+def resolve_model() -> str:
+    from huggingface_hub import HfApi
+
+    override = os.getenv("TRAIN_MODEL_ID")
+    api = HfApi(token=os.environ["HF_TOKEN"])
+    tried = []
+    for repo in ([override] if override else []) + CANDIDATES:
+        if not repo or repo in tried:
+            continue
+        tried.append(repo)
+        try:
+            api.model_info(repo)
+            print(f"[resolve] using {repo}")
+            Path("/kaggle/working/resolved_model.json").write_text(
+                json.dumps({"train_model_id": repo, "tried": tried}, indent=2),
+                encoding="utf-8",
+            )
+            return repo
+        except Exception as e:
+            print(f"[resolve] skip {repo}: {e}")
+    raise RuntimeError(f"No model resolved; tried {tried}")
+
+
+# ---------------------------------------------------------------------------
+# Langfuse checkpoint eval
+# ---------------------------------------------------------------------------
+
+
+def langfuse_log_ckpt_eval(
+    *,
+    step: int,
+    hf_revision: str,
+    train_model_id: str,
+    results: list[dict[str, Any]],
+    wandb_run_id: str | None,
+) -> dict[str, Any]:
+    from langfuse import Langfuse, get_client
+
+    Langfuse(
+        public_key=os.environ["LANGFUSE_PUBLIC_KEY"],
+        secret_key=os.environ["LANGFUSE_SECRET_KEY"],
+        host=os.getenv("LANGFUSE_HOST") or os.getenv("LANGFUSE_BASE_URL"),
+    )
+    lf = get_client()
+    if not lf.auth_check():
+        raise RuntimeError("Langfuse auth failed during ckpt eval")
+
+    n = len(results)
+    mean_pass = sum(r["pass_frac"] for r in results) / n if n else 0.0
+    mean_reward = sum(r["reward_total"] for r in results) / n if n else 0.0
+    n_perfect = sum(1 for r in results if r["pass_frac"] >= 1.0)
+
+    with lf.start_as_current_observation(
+        name="kaggle-ckpt-eval",
+        as_type="evaluator",
+        metadata={
+            "step": step,
+            "hf_revision": hf_revision,
+            "train_model_id": train_model_id,
+            "wandb_run_id": wandb_run_id,
+            "phase": "ckpt_eval",
+        },
+        input={"step": step, "hf_revision": hf_revision, "n": n},
+    ) as root:
+        for r in results:
+            with lf.start_as_current_observation(
+                name=f"task:{r['task_id']}",
+                as_type="span",
+                input={"prompt": r["prompt"][:2000]},
+            ) as span:
+                span.update(
+                    output={
+                        "completion": r["completion"][:2000],
+                        "pass_frac": r["pass_frac"],
+                        "reward_total": r["reward_total"],
+                    }
+                )
+        summary = {
+            "mean_pass_frac": mean_pass,
+            "mean_reward_total": mean_reward,
+            "n_perfect": n_perfect,
+            "n_problems": n,
+            "step": step,
+            "hf_revision": hf_revision,
+        }
+        root.update(output=summary)
+        lf.score_current_trace(name="mean_pass_frac", value=mean_pass)
+        lf.score_current_trace(name="mean_reward_total", value=mean_reward)
+        lf.score_current_trace(name="n_perfect", value=float(n_perfect))
+        lf.score_current_trace(name="checkpoint_step", value=float(step))
+        trace_id = lf.get_current_trace_id() or ""
+    lf.flush()
+    print(f"[langfuse] ckpt-eval step={step} rev={hf_revision} mean_pass={mean_pass:.3f} trace={trace_id}")
+    return {"langfuse_trace_id": trace_id, **summary}
+
+
+def run_heldout_eval(model, tokenizer, problems: list[dict], max_new_tokens: int = 512) -> list[dict]:
+    import torch
+
+    rows = []
+    model.eval()
+    for p in problems:
+        messages = [{"role": "user", "content": p["prompt"]}]
+        try:
+            prompt_text = tokenizer.apply_chat_template(
+                messages, tokenize=False, add_generation_prompt=True
+            )
+        except Exception:
+            prompt_text = p["prompt"]
+        inputs = tokenizer(prompt_text, return_tensors="pt").to(model.device)
+        with torch.no_grad():
+            out = model.generate(
+                **inputs,
+                max_new_tokens=max_new_tokens,
+                do_sample=False,
+                pad_token_id=tokenizer.eos_token_id,
+            )
+        gen = out[0][inputs["input_ids"].shape[-1] :]
+        completion = tokenizer.decode(gen, skip_special_tokens=True)
+        g = grade_completion(completion, list(p["hidden_tests"]))
+        rows.append(
+            {
+                "task_id": p["task_id"],
+                "prompt": p["prompt"],
+                "completion": completion,
+                "pass_frac": float(g["pass_frac"]),
+                "reward_total": float(g["pass_frac"]),
+            }
+        )
+    return rows
+
+
+# ---------------------------------------------------------------------------
+# Callback: eval + langfuse THEN save
+# ---------------------------------------------------------------------------
+
+
+def make_eval_before_save_callback(
+    *,
+    eval_problems: list[dict],
+    train_model_id: str,
+    hub_model_id: str,
+    save_steps: int,
+    n_eval: int,
+):
+    from transformers import TrainerCallback, TrainerControl, TrainerState, TrainingArguments
+
+    class EvalBeforeCheckpointCallback(TrainerCallback):
+        def __init__(self) -> None:
+            self._last_eval_step = -1
+
+        def on_save(
+            self,
+            args: TrainingArguments,
+            state: TrainerState,
+            control: TrainerControl,
+            **kwargs,
+        ):
+            # Transformers calls on_save AFTER saving locally. We also hook on_step_end
+            # to eval+langfuse BEFORE hub push. See on_step_end.
+            return control
+
+        def on_step_end(
+            self,
+            args: TrainingArguments,
+            state: TrainerState,
+            control: TrainerControl,
+            **kwargs,
+        ):
+            step = state.global_step
+            if save_steps <= 0 or step <= 0 or step % save_steps != 0:
+                return control
+            if step == self._last_eval_step:
+                return control
+            self._last_eval_step = step
+
+            model = kwargs.get("model")
+            tokenizer = kwargs.get("processing_class") or kwargs.get("tokenizer")
+            # GRPOTrainer may put tokenizer on trainer; fallback via kwargs["train"] not available
+            trainer = kwargs.get("trainer")  # usually absent
+            if tokenizer is None:
+                # recover from module global set in main()
+                tokenizer = globals().get("_TOKENIZER")
+            if model is None or tokenizer is None:
+                print("[ckpt-eval] missing model/tokenizer; skipping")
+                return control
+
+            rev = f"grpo-step-{step}"
+            problems = eval_problems[:n_eval]
+            print(f"\n=== EVAL BEFORE CHECKPOINT step={step} rev={rev} n={len(problems)} ===")
+            try:
+                results = run_heldout_eval(model, tokenizer, problems)
+                import wandb
+
+                wandb_run_id = None
+                if wandb.run is not None:
+                    wandb_run_id = wandb.run.id
+                    mean_pass = sum(r["pass_frac"] for r in results) / max(len(results), 1)
+                    wandb.log(
+                        {
+                            "eval/mean_pass_frac": mean_pass,
+                            "eval/n_perfect": sum(1 for r in results if r["pass_frac"] >= 1.0),
+                            "eval/checkpoint_step": step,
+                        },
+                        step=step,
+                    )
+
+                summary = langfuse_log_ckpt_eval(
+                    step=step,
+                    hf_revision=rev,
+                    train_model_id=train_model_id,
+                    results=results,
+                    wandb_run_id=wandb_run_id,
+                )
+                # Persist summary next to checkpoint
+                out = Path(args.output_dir) / f"ckpt_eval_step_{step}.json"
+                out.write_text(json.dumps(summary | {"results": results}, indent=2), encoding="utf-8")
+
+                # Push adapter AFTER langfuse success
+                try:
+                    model.push_to_hub(
+                        hub_model_id,
+                        token=os.environ["HF_TOKEN"],
+                        revision=rev,
+                        commit_message=f"GRPO adapter {rev}; mean_pass={summary['mean_pass_frac']:.3f}",
+                        private=False,
+                    )
+                    print(f"[hf] pushed {hub_model_id}@{rev}")
+                    if wandb.run is not None:
+                        wandb.log(
+                            {
+                                "hf/revision": rev,
+                                "hf/mean_pass_at_push": summary["mean_pass_frac"],
+                                "langfuse/ckpt_trace_id": summary.get("langfuse_trace_id"),
+                            },
+                            step=step,
+                        )
+                except Exception as e:
+                    print(f"[hf] push failed (local ckpt still ok): {e}")
+                    traceback.print_exc()
+            except Exception:
+                print("[ckpt-eval] FAILED — blocking is policy; re-raising so bad ckpts aren't silent")
+                traceback.print_exc()
+                raise
+            return control
+
+    return EvalBeforeCheckpointCallback()
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class Cfg:
+    dataset_id: str = "vanishingradient/deepagents-rl-mbpp"
+    hub_model_id: str = "vanishingradient/deepagents-rl-coding"
+    max_steps: int = 40
+    save_steps: int = 5
+    num_generations: int = 4
+    learning_rate: float = 1e-6
+    max_completion_length: int = 512
+    max_seq_length: int = 2048
+    n_eval: int = 16
+    lora_r: int = 16
+
+
+def main(cfg: Cfg | None = None) -> None:
+    global _TOKENIZER
+    cfg = cfg or Cfg()
+    load_secrets()
+
+    import torch
+    import wandb
+    from datasets import load_dataset
+    from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+    from trl import GRPOConfig, GRPOTrainer
+
+    os.environ["WANDB_PROJECT"] = os.getenv("WANDB_PROJECT", "deepagents-rl-coding")
+    wandb.login(key=os.environ["WANDB_API_KEY"])
+
+    model_id = resolve_model()
+    ds = load_dataset(cfg.dataset_id, token=os.environ["HF_TOKEN"])
+    train = ds["train"]
+    ev = ds["eval"]
+    # Ensure list columns serializable for reward
+    def _ensure(example):
+        ht = example["hidden_tests"]
+        if not isinstance(ht, list):
+            ht = list(ht)
+        example["hidden_tests"] = ht
+        return example
+
+    train = train.map(_ensure)
+    eval_problems = [
+        {
+            "task_id": row["task_id"],
+            "prompt": row["prompt"],
+            "hidden_tests": list(row["hidden_tests"]),
+        }
+        for row in ev
+    ]
+
+    bnb = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_use_double_quant=True,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(model_id, token=os.environ["HF_TOKEN"], trust_remote_code=True)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    _TOKENIZER = tokenizer
+
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        quantization_config=bnb,
+        device_map="auto",
+        token=os.environ["HF_TOKEN"],
+        trust_remote_code=True,
+    )
+    model = prepare_model_for_kbit_training(model)
+    lora = LoraConfig(
+        r=cfg.lora_r,
+        lora_alpha=cfg.lora_r * 2,
+        lora_dropout=0.0,
+        bias="none",
+        task_type="CAUSAL_LM",
+        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+    )
+    model = get_peft_model(model, lora)
+
+    output_dir = "/kaggle/working/grpo-out"
+    Path(output_dir).mkdir(parents=True, exist_ok=True)
+    globals()["_TOKENIZER"] = tokenizer
+
+    wandb.init(
+        project=os.environ["WANDB_PROJECT"],
+        name="kaggle-grpo-mbpp",
+        job_type="train",
+        config={
+            "train_model_id": model_id,
+            "dataset_id": cfg.dataset_id,
+            "max_steps": cfg.max_steps,
+            "save_steps": cfg.save_steps,
+            "n_eval": cfg.n_eval,
+            "hub_model_id": cfg.hub_model_id,
+            "eval_before_every_checkpoint": True,
+        },
+    )
+
+    # Baseline eval @ step 0 BEFORE any training / checkpoint
+    print("=== BASELINE EVAL (step 0) before any training ===")
+    baseline = run_heldout_eval(model, tokenizer, eval_problems[: cfg.n_eval])
+    base_summary = langfuse_log_ckpt_eval(
+        step=0,
+        hf_revision="grpo-step-0-baseline",
+        train_model_id=model_id,
+        results=baseline,
+        wandb_run_id=wandb.run.id if wandb.run else None,
+    )
+    Path(output_dir, "ckpt_eval_step_0.json").write_text(
+        json.dumps(base_summary | {"results": baseline}, indent=2), encoding="utf-8"
+    )
+    wandb.log(
+        {
+            "eval/mean_pass_frac": base_summary["mean_pass_frac"],
+            "eval/n_perfect": base_summary["n_perfect"],
+            "eval/checkpoint_step": 0,
+            "langfuse/ckpt_trace_id": base_summary.get("langfuse_trace_id"),
+        },
+        step=0,
+    )
+
+    args = GRPOConfig(
+        output_dir=output_dir,
+        run_name="kaggle-grpo-mbpp",
+        learning_rate=cfg.learning_rate,
+        per_device_train_batch_size=1,
+        gradient_accumulation_steps=4,
+        num_generations=cfg.num_generations,
+        max_completion_length=cfg.max_completion_length,
+        max_steps=cfg.max_steps,
+        logging_steps=1,
+        save_steps=cfg.save_steps,
+        save_total_limit=3,
+        report_to=["wandb"],
+        bf16=torch.cuda.is_bf16_supported(),
+        fp16=not torch.cuda.is_bf16_supported(),
+        remove_unused_columns=False,
+        hub_model_id=cfg.hub_model_id,
+        hub_token=os.environ["HF_TOKEN"],
+        push_to_hub=False,  # push ONLY after Langfuse ckpt-eval succeeds
+    )
+
+    cb = make_eval_before_save_callback(
+        eval_problems=eval_problems,
+        train_model_id=model_id,
+        hub_model_id=cfg.hub_model_id,
+        save_steps=cfg.save_steps,
+        n_eval=cfg.n_eval,
+    )
+
+    trainer = GRPOTrainer(
+        model=model,
+        reward_funcs=coding_reward,
+        args=args,
+        train_dataset=train,
+        processing_class=tokenizer,
+        callbacks=[cb],
+    )
+
+    trainer.train()
+    # Final eval + push
+    final_step = int(trainer.state.global_step)
+    rev = f"grpo-step-{final_step}-final"
+    print(f"=== FINAL EVAL step={final_step} ===")
+    final_results = run_heldout_eval(model, tokenizer, eval_problems[: cfg.n_eval])
+    final_summary = langfuse_log_ckpt_eval(
+        step=final_step,
+        hf_revision=rev,
+        train_model_id=model_id,
+        results=final_results,
+        wandb_run_id=wandb.run.id if wandb.run else None,
+    )
+    model.push_to_hub(
+        cfg.hub_model_id,
+        token=os.environ["HF_TOKEN"],
+        revision=rev,
+        commit_message=f"Final GRPO {rev}; mean_pass={final_summary['mean_pass_frac']:.3f}",
+    )
+    # also push main
+    model.push_to_hub(
+        cfg.hub_model_id,
+        token=os.environ["HF_TOKEN"],
+        commit_message=f"Final GRPO main; mean_pass={final_summary['mean_pass_frac']:.3f}",
+    )
+    wandb.log(
+        {
+            "eval/final_mean_pass_frac": final_summary["mean_pass_frac"],
+            "hf/final_revision": rev,
+        }
+    )
+    wandb.finish()
+    print("TRAIN DONE", final_summary)
+
+
+_TOKENIZER = None
+
+if __name__ == "__main__":
+    main()