Fix NaN after backprop: stay in train mode, use no_grad for generation only, lower LR to 1e-5

GRPO_Training_Pipeline.ipynb

@@ -80,7 +80,7 @@
   },
   {
    "cell_type": "code",
-   "source": "###############################################################\n# REAL GRPO TRAINING\n###############################################################\n\nimport torch, requests, json, time, gc, re\nimport torch.nn.functional as F\nimport matplotlib\nmatplotlib.use(\"Agg\")\nimport matplotlib.pyplot as plt\nfrom unsloth import FastLanguageModel\n\nENV = \"https://cobedigger-dataengenv.hf.space\"\n\nG           = 4\nTRAIN_STEPS = 25\nLR          = 2e-5\nMAX_NEW     = 150\nMAX_PROMPT  = 350\nMAX_FULL    = MAX_PROMPT + MAX_NEW  # 500\n\nSYSTEM = (\n    \"You are an ML engineer. Respond with ONLY a JSON object, nothing else.\\n\"\n    \"Format: {\\\"action_type\\\": \\\"<action>\\\", \\\"payload\\\": {}}\\n\"\n    \"Actions: inspect_data, run_script, edit_script, query_actor, submit\\n\"\n    \"For edit_script use: {\\\"action_type\\\": \\\"edit_script\\\", \\\"payload\\\": {\\\"old\\\": \\\"...\\\", \\\"new\\\": \\\"...\\\"}}\\n\\n\"\n    \"Stage 1 bug: column named 'age_years' should be 'age', and missing df.dropna()\\n\"\n    \"Stage 2 bug: MLPClassifier needs StandardScaler before fitting\\n\"\n    \"Stage 3 bug: scaler.fit called before train_test_split (data leakage)\\n\"\n    \"Stage 4 bug: LogisticRegression needs class_weight='balanced'\"\n)\n\nSTAGE_PROMPTS = {\n    1: \"Stage 1 error: KeyError 'age_years'. The column should be named 'age'. There are also NaN rows. Pick one action.\",\n    2: \"Stage 2 error: MLPClassifier loss is NaN. Features need StandardScaler normalization. Pick one action.\",\n    3: \"Stage 3: accuracy is suspiciously 0.98. Scaler is fit before the train/test split causing leakage. Pick one action.\",\n    4: \"Stage 4: deployment blocked, fairness_gap too high. Model needs class_weight='balanced'. Pick one action.\",\n}\n\n# What we want to see in completions for each stage\nSTAGE_HINTS = {\n    1: [\"age_years\", \"age\", \"dropna\", \"rename\", \"KeyError\"],\n    2: [\"StandardScaler\", \"scaler\", \"normalize\", \"MLPClassifier\", \"NaN\"],\n    3: [\"train_test_split\", \"leakage\", \"fit_transform\", \"after split\"],\n    4: [\"class_weight\", \"balanced\", \"fairness\", \"LogisticRegression\"],\n}\n\nVALID_ACTIONS = {\"inspect_data\", \"run_script\", \"edit_script\", \"query_actor\", \"submit\"}\n\ndef parse_action(text):\n    \"\"\"Very lenient parser — tries multiple strategies.\"\"\"\n    text = text.strip()\n\n    # Strategy 1: find first complete JSON object\n    s = text.find(\"{\")\n    if s != -1:\n        depth = 0\n        for i, ch in enumerate(text[s:], s):\n            if ch == \"{\": depth += 1\n            elif ch == \"}\":\n                depth -= 1\n                if depth == 0:\n                    try:\n                        obj = json.loads(text[s:i+1])\n                        return obj\n                    except Exception:\n                        break\n\n    # Strategy 2: look for action_type anywhere in text (case-insensitive)\n    m = re.search(r'\"?action_?type\"?\\s*[=:]\\s*\"?(\\w+)\"?', text, re.IGNORECASE)\n    if m:\n        atype = m.group(1).lower()\n        for va in VALID_ACTIONS:\n            if va.replace(\"_\", \"\") in atype.replace(\"_\", \"\"):\n                return {\"action_type\": va, \"payload\": {}}\n\n    return None\n\ndef score_action(text, stage):\n    \"\"\"\n    Reward layers:\n      0.00 — no JSON / unparseable\n      0.05 — parseable but wrong action name\n      0.15 — valid action_type\n      0.25 — contextually right action for stage\n      0.40 — edit_script with stage-relevant keyword in payload\n      0.60+ — edit_script that actually moves the grader\n    \"\"\"\n    action = parse_action(text)\n    if action is None:\n        return 0.0\n\n    # Fix payload type\n    action.setdefault(\"payload\", {})\n    if isinstance(action[\"payload\"], str):\n        action[\"payload\"] = {\"old\": action[\"payload\"], \"new\": \"\"}\n\n    atype = str(action.get(\"action_type\", \"\")).lower()\n\n    # Normalise action name (model sometimes writes \"Action\" or \"run script\")\n    for va in VALID_ACTIONS:\n        if va.replace(\"_\", \"\") in atype.replace(\"_\", \"\").replace(\" \", \"\"):\n            atype = va\n            break\n    else:\n        return 0.05  # parseable but unknown action\n\n    action[\"action_type\"] = atype\n    reward = 0.15  # valid action\n\n    # Contextually good action\n    good = {1: {\"inspect_data\",\"edit_script\",\"run_script\"},\n            2: {\"inspect_data\",\"edit_script\",\"run_script\"},\n            3: {\"run_script\",\"edit_script\"},\n            4: {\"query_actor\",\"edit_script\",\"run_script\"}}\n    if atype in good.get(stage, set()):\n        reward = 0.25\n\n    # edit_script keyword check\n    if atype == \"edit_script\":\n        payload_str = json.dumps(action[\"payload\"]).lower()\n        full_str    = (text + payload_str).lower()\n        hints       = STAGE_HINTS.get(stage, [])\n        hits        = sum(1 for h in hints if h.lower() in full_str)\n        if hits >= 1: reward = 0.40\n        if hits >= 2: reward = 0.55\n\n        # Try executing against live env for grader signal\n        try:\n            requests.post(f\"{ENV}/reset\", timeout=15)\n            r     = requests.post(f\"{ENV}/step\", json=action, timeout=15).json()\n            score = float(r.get(\"reward\", {}).get(\"score\", 0.0))\n            if score > 0:\n                reward = min(1.0, 0.55 + score * 0.45)\n        except Exception:\n            pass\n\n    return reward\n\ndef make_prompt(stage):\n    msgs = [{\"role\": \"system\", \"content\": SYSTEM},\n            {\"role\": \"user\",   \"content\": STAGE_PROMPTS[stage]}]\n    return tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)\n\ndef reinforce_loss(input_ids, attention_mask, advantages):\n    logits   = model(input_ids=input_ids, attention_mask=attention_mask).logits\n    shift_l  = logits[..., :-1, :].contiguous()\n    shift_lb = input_ids[..., 1:].contiguous()\n    lp       = F.log_softmax(shift_l, dim=-1)\n    tok_lp   = lp.gather(-1, shift_lb.unsqueeze(-1)).squeeze(-1)\n    mask     = attention_mask[..., 1:].float()\n    seq_lp   = (tok_lp * mask).sum(-1) / mask.sum(-1).clamp(min=1)\n    return -(advantages * seq_lp).mean()\n\n# ── Setup ────────────────────────────────────────────────────\nprint(\"Setting up GRPO training...\")\ntorch.cuda.empty_cache(); gc.collect()\nFastLanguageModel.for_training(model)\n\noptimizer = torch.optim.AdamW(\n    [p for p in model.parameters() if p.requires_grad], lr=LR\n)\n\ngrpo_rewards, grpo_losses = [], []\nprint(f\"\\nStarting GRPO: {TRAIN_STEPS} steps, G={G} completions/step\")\nprint(\"=\" * 60)\n\nfor step in range(TRAIN_STEPS):\n    stage  = [1, 2, 3, 4][step % 4]\n    prompt = make_prompt(stage)\n\n    enc = tokenizer(prompt, return_tensors=\"pt\",\n                    truncation=True, max_length=MAX_PROMPT).to(\"cuda\")\n    prompt_len = enc[\"input_ids\"].shape[1]\n\n    # 1. Generate G completions\n    FastLanguageModel.for_inference(model)\n    model.eval()\n    with torch.no_grad():\n        outs = model.generate(\n            **enc,\n            max_new_tokens=MAX_NEW,\n            do_sample=True,\n            temperature=0.9,\n            num_return_sequences=G,\n            pad_token_id=tokenizer.eos_token_id,\n            repetition_penalty=1.1,\n        )\n    completions = [\n        tokenizer.decode(outs[i][prompt_len:], skip_special_tokens=True)\n        for i in range(G)\n    ]\n\n    # 2. Score\n    rewards = [score_action(c, stage) for c in completions]\n    time.sleep(0.2)\n\n    rewards_t = torch.tensor(rewards, dtype=torch.float32)\n    mean_r    = rewards_t.mean()\n    std_r     = rewards_t.std().clamp(min=1e-6)\n\n    print(f\"  Step {step+1:02d}/{TRAIN_STEPS} | Stage {stage} | \"\n          f\"Rewards={[f'{r:.2f}' for r in rewards]} | Avg={mean_r.item():.3f}\", end=\"\")\n\n    if std_r < 1e-4:\n        print(\" | SKIP (no variance)\")\n        grpo_rewards.append(mean_r.item())\n        grpo_losses.append(0.0)\n        torch.cuda.empty_cache()\n        continue\n\n    advantages = ((rewards_t - mean_r) / std_r).to(\"cuda\")\n\n    # 3. Tokenise\n    full_ids_list, full_mask_list = [], []\n    for c in completions:\n        enc_full = tokenizer(\n            prompt + c, return_tensors=\"pt\",\n            truncation=True, padding=\"max_length\", max_length=MAX_FULL,\n        )\n        full_ids_list.append(enc_full[\"input_ids\"])\n        full_mask_list.append(enc_full[\"attention_mask\"])\n    full_ids  = torch.cat(full_ids_list,  dim=0).to(\"cuda\")\n    full_mask = torch.cat(full_mask_list, dim=0).to(\"cuda\")\n\n    # 4. Backprop\n    FastLanguageModel.for_training(model)\n    model.train()\n    loss = reinforce_loss(full_ids, full_mask, advantages)\n    loss.backward()\n    torch.nn.utils.clip_grad_norm_(\n        [p for p in model.parameters() if p.requires_grad], 1.0\n    )\n    optimizer.step()\n    optimizer.zero_grad()\n    torch.cuda.empty_cache()\n\n    grpo_rewards.append(mean_r.item())\n    grpo_losses.append(loss.item())\n    print(f\" | Loss={loss.item():.4f}\")\n\nprint(\"\\nGRPO training complete!\")\nprint(f\"   First 5 avg reward : {sum(grpo_rewards[:5])/5:.3f}\")\nprint(f\"   Last  5 avg reward : {sum(grpo_rewards[-5:])/5:.3f}\")\n\nfig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 4))\nax1.plot(grpo_rewards, color=\"#22c55e\", linewidth=2, marker=\"o\", markersize=4)\nax1.set_title(\"GRPO Avg Reward per Step\", fontweight=\"bold\")\nax1.set_xlabel(\"Training Step\"); ax1.set_ylabel(\"Avg Group Reward\")\nax1.set_ylim(-0.05, 1.1); ax1.grid(True, alpha=0.3)\nax2.plot(grpo_losses, color=\"#f59e0b\", linewidth=2, marker=\"o\", markersize=4)\nax2.set_title(\"GRPO Policy Loss per Step\", fontweight=\"bold\")\nax2.set_xlabel(\"Training Step\"); ax2.set_ylabel(\"Loss\")\nax2.grid(True, alpha=0.3)\nplt.tight_layout()\nplt.savefig(\"grpo_training_curve.png\", dpi=150, bbox_inches=\"tight\")\nplt.show()\n\nwith open(\"grpo_rewards.json\", \"w\") as f:\n    json.dump({\"rewards\": grpo_rewards, \"losses\": grpo_losses}, f)\nprint(\"Saved grpo_training_curve.png and grpo_rewards.json\")\n",
+   "source": "###############################################################\n# REAL GRPO TRAINING\n###############################################################\n\nimport torch, requests, json, time, gc, re\nimport torch.nn.functional as F\nimport matplotlib\nmatplotlib.use(\"Agg\")\nimport matplotlib.pyplot as plt\nfrom unsloth import FastLanguageModel\n\nENV = \"https://cobedigger-dataengenv.hf.space\"\n\nG           = 4\nTRAIN_STEPS = 25\nLR          = 1e-5   # lower LR — 4bit models are sensitive\nMAX_NEW     = 150\nMAX_PROMPT  = 350\nMAX_FULL    = MAX_PROMPT + MAX_NEW\n\nSYSTEM = (\n    \"You are an ML engineer. Respond with ONLY a JSON object, nothing else.\\n\"\n    \"Format: {\\\"action_type\\\": \\\"<action>\\\", \\\"payload\\\": {}}\\n\"\n    \"Actions: inspect_data, run_script, edit_script, query_actor, submit\\n\"\n    \"For edit_script: {\\\"action_type\\\": \\\"edit_script\\\", \\\"payload\\\": {\\\"old\\\": \\\"...\\\", \\\"new\\\": \\\"...\\\"}}\\n\\n\"\n    \"Stage 1 bug: column 'age_years' should be 'age', missing df.dropna()\\n\"\n    \"Stage 2 bug: MLPClassifier needs StandardScaler\\n\"\n    \"Stage 3 bug: scaler.fit called before train_test_split\\n\"\n    \"Stage 4 bug: LogisticRegression needs class_weight='balanced'\"\n)\n\nSTAGE_PROMPTS = {\n    1: \"Stage 1 error: KeyError 'age_years'. Column should be 'age'. NaN rows present. Pick one action.\",\n    2: \"Stage 2 error: MLPClassifier loss is NaN. Needs StandardScaler normalization. Pick one action.\",\n    3: \"Stage 3: accuracy 0.98 — scaler fit before split causes leakage. Pick one action.\",\n    4: \"Stage 4: fairness_gap too high — needs class_weight='balanced'. Pick one action.\",\n}\n\nSTAGE_HINTS = {\n    1: [\"age_years\", \"age\", \"dropna\", \"rename\"],\n    2: [\"StandardScaler\", \"scaler\", \"normalize\"],\n    3: [\"train_test_split\", \"leakage\", \"after\"],\n    4: [\"class_weight\", \"balanced\"],\n}\n\nVALID_ACTIONS = {\"inspect_data\", \"run_script\", \"edit_script\", \"query_actor\", \"submit\"}\n\ndef parse_action(text):\n    s = text.find(\"{\")\n    if s != -1:\n        depth = 0\n        for i, ch in enumerate(text[s:], s):\n            if ch == \"{\": depth += 1\n            elif ch == \"}\":\n                depth -= 1\n                if depth == 0:\n                    try:\n                        return json.loads(text[s:i+1])\n                    except Exception:\n                        break\n    m = re.search(r'\"?action_?type\"?\\s*[=:]\\s*\"?(\\w+)\"?', text, re.IGNORECASE)\n    if m:\n        atype = m.group(1).lower().replace(\" \", \"_\")\n        for va in VALID_ACTIONS:\n            if va.replace(\"_\",\"\") in atype.replace(\"_\",\"\"):\n                return {\"action_type\": va, \"payload\": {}}\n    return None\n\ndef score_action(text, stage):\n    action = parse_action(text)\n    if action is None:\n        return 0.0\n    action.setdefault(\"payload\", {})\n    if isinstance(action[\"payload\"], str):\n        action[\"payload\"] = {\"old\": action[\"payload\"], \"new\": \"\"}\n\n    atype = str(action.get(\"action_type\",\"\")).lower().replace(\" \",\"_\")\n    for va in VALID_ACTIONS:\n        if va.replace(\"_\",\"\") in atype.replace(\"_\",\"\"):\n            atype = va; break\n    else:\n        return 0.05\n    action[\"action_type\"] = atype\n\n    reward = 0.15\n    good = {1:{\"inspect_data\",\"edit_script\",\"run_script\"},\n            2:{\"inspect_data\",\"edit_script\",\"run_script\"},\n            3:{\"run_script\",\"edit_script\"},\n            4:{\"query_actor\",\"edit_script\",\"run_script\"}}\n    if atype in good.get(stage, set()):\n        reward = 0.25\n\n    if atype == \"edit_script\":\n        full = (text + json.dumps(action[\"payload\"])).lower()\n        hits = sum(1 for h in STAGE_HINTS.get(stage,[]) if h.lower() in full)\n        if hits >= 1: reward = 0.40\n        if hits >= 2: reward = 0.55\n        try:\n            requests.post(f\"{ENV}/reset\", timeout=15)\n            r = requests.post(f\"{ENV}/step\", json=action, timeout=15).json()\n            sc = float(r.get(\"reward\",{}).get(\"score\", 0.0))\n            if sc > 0: reward = min(1.0, 0.55 + sc * 0.45)\n        except Exception:\n            pass\n    return reward\n\ndef make_prompt(stage):\n    msgs = [{\"role\":\"system\",\"content\":SYSTEM},\n            {\"role\":\"user\",\"content\":STAGE_PROMPTS[stage]}]\n    return tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)\n\ndef reinforce_loss(input_ids, attention_mask, advantages):\n    logits  = model(input_ids=input_ids, attention_mask=attention_mask).logits\n    shift_l = logits[..., :-1, :].contiguous()\n    shift_b = input_ids[..., 1:].contiguous()\n    lp      = F.log_softmax(shift_l, dim=-1)\n    tok_lp  = lp.gather(-1, shift_b.unsqueeze(-1)).squeeze(-1)\n    mask    = attention_mask[..., 1:].float()\n    seq_lp  = (tok_lp * mask).sum(-1) / mask.sum(-1).clamp(min=1)\n    return -(advantages * seq_lp).mean()\n\n# ── Setup — stay in training mode the whole time ─────────────\nprint(\"Setting up GRPO training...\")\ntorch.cuda.empty_cache(); gc.collect()\nFastLanguageModel.for_training(model)\nmodel.train()\n\noptimizer = torch.optim.AdamW(\n    [p for p in model.parameters() if p.requires_grad], lr=LR\n)\n\ngrpo_rewards, grpo_losses = [], []\nprint(f\"\\nStarting GRPO: {TRAIN_STEPS} steps, G={G} completions/step\")\nprint(\"=\" * 60)\n\nfor step in range(TRAIN_STEPS):\n    stage  = [1, 2, 3, 4][step % 4]\n    prompt = make_prompt(stage)\n\n    enc = tokenizer(prompt, return_tensors=\"pt\",\n                    truncation=True, max_length=MAX_PROMPT).to(\"cuda\")\n    prompt_len = enc[\"input_ids\"].shape[1]\n\n    # 1. Generate with torch.no_grad() — stay in train mode, just no grad tracking\n    with torch.no_grad():\n        outs = model.generate(\n            **enc,\n            max_new_tokens=MAX_NEW,\n            do_sample=True,\n            temperature=0.9,\n            num_return_sequences=G,\n            pad_token_id=tokenizer.eos_token_id,\n            repetition_penalty=1.1,\n        )\n    completions = [\n        tokenizer.decode(outs[i][prompt_len:], skip_special_tokens=True)\n        for i in range(G)\n    ]\n\n    # 2. Score\n    rewards = [score_action(c, stage) for c in completions]\n    time.sleep(0.2)\n\n    rewards_t = torch.tensor(rewards, dtype=torch.float32)\n    mean_r    = rewards_t.mean()\n    std_r     = rewards_t.std().clamp(min=1e-6)\n\n    print(f\"  Step {step+1:02d}/{TRAIN_STEPS} | Stage {stage} | \"\n          f\"Rewards={[f'{r:.2f}' for r in rewards]} | Avg={mean_r.item():.3f}\", end=\"\")\n\n    if std_r < 1e-4:\n        print(\" | SKIP (no variance)\")\n        grpo_rewards.append(mean_r.item())\n        grpo_losses.append(0.0)\n        continue\n\n    advantages = ((rewards_t - mean_r) / std_r).to(\"cuda\")\n\n    # 3. Tokenise\n    full_ids_list, full_mask_list = [], []\n    for c in completions:\n        ef = tokenizer(prompt + c, return_tensors=\"pt\",\n                       truncation=True, padding=\"max_length\", max_length=MAX_FULL)\n        full_ids_list.append(ef[\"input_ids\"])\n        full_mask_list.append(ef[\"attention_mask\"])\n    full_ids  = torch.cat(full_ids_list,  dim=0).to(\"cuda\")\n    full_mask = torch.cat(full_mask_list, dim=0).to(\"cuda\")\n\n    # 4. Backprop — model already in train mode, no mode switch needed\n    loss = reinforce_loss(full_ids, full_mask, advantages)\n    loss.backward()\n    torch.nn.utils.clip_grad_norm_(\n        [p for p in model.parameters() if p.requires_grad], 0.5\n    )\n    optimizer.step()\n    optimizer.zero_grad()\n    torch.cuda.empty_cache()\n\n    grpo_rewards.append(mean_r.item())\n    grpo_losses.append(loss.item())\n    print(f\" | Loss={loss.item():.4f}\")\n\nprint(\"\\nGRPO training complete!\")\nprint(f\"   First 5 avg reward : {sum(grpo_rewards[:5])/5:.3f}\")\nprint(f\"   Last  5 avg reward : {sum(grpo_rewards[-5:])/5:.3f}\")\n\nfig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 4))\nax1.plot(grpo_rewards, color=\"#22c55e\", linewidth=2, marker=\"o\", markersize=4)\nax1.set_title(\"GRPO Avg Reward per Step\", fontweight=\"bold\")\nax1.set_xlabel(\"Training Step\"); ax1.set_ylabel(\"Avg Group Reward\")\nax1.set_ylim(-0.05, 1.1); ax1.grid(True, alpha=0.3)\nax2.plot(grpo_losses, color=\"#f59e0b\", linewidth=2, marker=\"o\", markersize=4)\nax2.set_title(\"GRPO Policy Loss per Step\", fontweight=\"bold\")\nax2.set_xlabel(\"Training Step\"); ax2.set_ylabel(\"Loss\")\nax2.grid(True, alpha=0.3)\nplt.tight_layout()\nplt.savefig(\"grpo_training_curve.png\", dpi=150, bbox_inches=\"tight\")\nplt.show()\n\nwith open(\"grpo_rewards.json\", \"w\") as f:\n    json.dump({\"rewards\": grpo_rewards, \"losses\": grpo_losses}, f)\nprint(\"Saved grpo_training_curve.png and grpo_rewards.json\")\n",
    "metadata": {
     "trusted": true
    },