Fix OOM: store ref LoRA weights on CPU instead of deepcopy second model

GRPO_Training_Pipeline.ipynb

@@ -80,7 +80,7 @@
   },
   {
    "cell_type": "code",
-   "source": "\n###############################################################\n# REAL GRPO TRAINING\n# The model GENERATES actions → gets rewards from the env →\n# GRPO computes group-relative advantages → backprop updates weights\n###############################################################\n\nimport torch, requests, json, time, gc, os, copy\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          = 5e-5\nKL_COEF     = 0.01\nMAX_NEW     = 256\nMAX_PROMPT  = 512\nMAX_FULL    = MAX_PROMPT + MAX_NEW   # 768\n\nSYSTEM = (\n    \"You are an ML engineer debugging a broken pipeline. \"\n    \"Respond with ONLY a JSON object — no explanation, no markdown.\\n\"\n    \"Actions: inspect_data | run_script | edit_script | query_actor | submit\\n\"\n    \"edit_script payload must have 'old' and 'new' keys.\\n\\n\"\n    \"Stage 1: rename age_years→age AND add df.dropna()\\n\"\n    \"Stage 2: add StandardScaler fitted on X_train AFTER train_test_split\\n\"\n    \"Stage 3: move scaler.fit to AFTER train_test_split\\n\"\n    \"Stage 4: add class_weight='balanced' to LogisticRegression\"\n)\n\nSTAGE_PROMPTS = {\n    1: \"Stage 1 | KeyError: 'age_years'\\nScript uses column 'age_years' but dataset has column 'age'. Also 12 NaN rows will crash the scaler. What is your action?\",\n    2: \"Stage 2 | MLPClassifier loss diverging (NaN loss after 5 iters)\\nFeatures are not normalised. What is your action?\",\n    3: \"Stage 3 | Suspiciously high accuracy: 0.98. Data leakage suspected.\\nscaler.fit_transform is called on full X before train_test_split. What is your action?\",\n    4: \"Stage 4 | Deployment blocked: fairness_gap=0.18 > threshold 0.05\\nModel predicts poorly for minority class. What is your action?\",\n}\n\ndef score_action(action_str: str, stage: int) -> float:\n    try:\n        s = action_str.find('{')\n        e = action_str.rfind('}')\n        if s == -1 or e == -1:\n            return 0.0\n        action = json.loads(action_str[s:e+1])\n        if \"action_type\" not in action:\n            return 0.0\n        action.setdefault(\"payload\", {})\n        if isinstance(action[\"payload\"], str):\n            action[\"payload\"] = {\"script\": action[\"payload\"]}\n    except Exception:\n        return 0.0\n\n    atype = action[\"action_type\"]\n    if atype == \"inspect_data\":\n        return 0.15\n    if atype == \"query_actor\" and stage == 4:\n        return 0.20\n    if atype == \"run_script\":\n        return 0.10\n    if atype == \"edit_script\":\n        try:\n            requests.post(f\"{ENV}/reset\", timeout=30)\n            r = requests.post(f\"{ENV}/step\", json=action, timeout=30).json()\n            score = float(r.get(\"reward\", {}).get(\"score\", 0.0))\n            return min(1.0, score + 0.1)\n        except Exception:\n            return 0.05\n    if atype == \"submit\":\n        return 0.05\n    return 0.0\n\ndef make_prompt(stage: int) -> str:\n    messages = [\n        {\"role\": \"system\", \"content\": SYSTEM},\n        {\"role\": \"user\",   \"content\": STAGE_PROMPTS[stage]},\n    ]\n    return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n\ndef grpo_loss(model, ref_model, input_ids, attention_mask, advantages):\n    with torch.no_grad():\n        ref_logits = ref_model(input_ids=input_ids, attention_mask=attention_mask).logits\n\n    logits = model(input_ids=input_ids, attention_mask=attention_mask).logits\n\n    shift_logits     = logits[..., :-1, :].contiguous()\n    shift_ref_logits = ref_logits[..., :-1, :].contiguous()\n    shift_labels     = input_ids[..., 1:].contiguous()\n\n    log_probs     = F.log_softmax(shift_logits,     dim=-1)\n    ref_log_probs = F.log_softmax(shift_ref_logits, dim=-1)\n\n    tok_lp     = log_probs.gather(-1,     shift_labels.unsqueeze(-1)).squeeze(-1)\n    ref_tok_lp = ref_log_probs.gather(-1, shift_labels.unsqueeze(-1)).squeeze(-1)\n\n    mask = attention_mask[..., 1:].float()\n    seq_lp     = (tok_lp     * mask).sum(-1) / mask.sum(-1).clamp(min=1)\n    ref_seq_lp = (ref_tok_lp * mask).sum(-1) / mask.sum(-1).clamp(min=1)\n\n    kl          = seq_lp - ref_seq_lp\n    policy_loss = -(advantages * seq_lp).mean()\n    kl_loss     = KL_COEF * kl.mean()\n\n    return policy_loss + kl_loss, policy_loss.item(), kl_loss.item()\n\n# ── Setup ────────────────────────────────────────────────────\nprint(\"Setting up GRPO training...\")\ntorch.cuda.empty_cache(); gc.collect()\n\nFastLanguageModel.for_training(model)\nref_model = copy.deepcopy(model)\nref_model.eval()\nfor p in ref_model.parameters():\n    p.requires_grad_(False)\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\", truncation=True, max_length=MAX_PROMPT).to(\"cuda\")\n    prompt_len = enc[\"input_ids\"].shape[1]\n\n    # 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.8,\n            num_return_sequences=G,\n            pad_token_id=tokenizer.eos_token_id,\n        )\n    completions = [\n        tokenizer.decode(outs[i][prompt_len:], skip_special_tokens=True)\n        for i in range(G)\n    ]\n\n    # Score each completion\n    rewards = [score_action(c, stage) for c in completions]\n    time.sleep(0.5)\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    advantages = ((rewards_t - mean_r) / std_r).to(\"cuda\")\n\n    # Tokenise prompt+completion for each — fixed single max_length arg\n    full_ids_list, full_mask_list = [], []\n    for c in completions:\n        enc_full = tokenizer(\n            prompt + c,\n            return_tensors=\"pt\",\n            truncation=True,\n            padding=\"max_length\",\n            max_length=MAX_FULL,\n        )\n        full_ids_list.append(enc_full[\"input_ids\"])\n        full_mask_list.append(enc_full[\"attention_mask\"])\n\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    # Backprop\n    FastLanguageModel.for_training(model)\n    model.train()\n    loss, pol_loss, kl_val = grpo_loss(model, ref_model, full_ids, full_mask, advantages)\n    loss.backward()\n    torch.nn.utils.clip_grad_norm_([p for p in model.parameters() if p.requires_grad], 1.0)\n    optimizer.step()\n    optimizer.zero_grad()\n\n    avg_r = rewards_t.mean().item()\n    grpo_rewards.append(avg_r)\n    grpo_losses.append(pol_loss)\n\n    print(\n        f\"  Step {step+1:02d}/{TRAIN_STEPS} | Stage {stage} | \"\n        f\"Rewards={[f'{r:.2f}' for r in rewards]} | \"\n        f\"Avg={avg_r:.3f} | PolicyLoss={pol_loss:.4f} | KL={kl_val:.4f}\"\n    )\n    torch.cuda.empty_cache()\n\nprint(\"\\n✅ GRPO 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)\n\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(\"Policy Loss\")\nax2.grid(True, alpha=0.3)\n\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# model generates actions -> env rewards -> advantages -> backprop\n###############################################################\n\nimport torch, requests, json, time, gc\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          = 5e-5\nKL_COEF     = 0.01\nMAX_NEW     = 256\nMAX_PROMPT  = 512\nMAX_FULL    = MAX_PROMPT + MAX_NEW\n\nSYSTEM = (\n    \"You are an ML engineer debugging a broken pipeline. \"\n    \"Respond with ONLY a JSON object. No explanation, no markdown.\\n\"\n    \"Actions: inspect_data | run_script | edit_script | query_actor | submit\\n\"\n    \"edit_script payload must have 'old' and 'new' keys.\\n\\n\"\n    \"Stage 1: rename age_years to age AND add df.dropna()\\n\"\n    \"Stage 2: add StandardScaler fitted on X_train AFTER train_test_split\\n\"\n    \"Stage 3: move scaler.fit to AFTER train_test_split\\n\"\n    \"Stage 4: add class_weight='balanced' to LogisticRegression\"\n)\n\nSTAGE_PROMPTS = {\n    1: \"Stage 1 | KeyError: 'age_years'\\nScript uses 'age_years' but dataset column is 'age'. Also 12 NaN rows. What is your action?\",\n    2: \"Stage 2 | MLPClassifier loss diverging (NaN after 5 iters)\\nFeatures not normalised. What is your action?\",\n    3: \"Stage 3 | Accuracy=0.98 — data leakage suspected.\\nscaler.fit_transform called on full X before split. What is your action?\",\n    4: \"Stage 4 | Deployment blocked: fairness_gap=0.18 > 0.05\\nMinority class poorly predicted. What is your action?\",\n}\n\ndef score_action(action_str, stage):\n    try:\n        s = action_str.find(\"{\"); e = action_str.rfind(\"}\")\n        if s == -1 or e == -1: return 0.0\n        action = json.loads(action_str[s:e+1])\n        if \"action_type\" not in action: return 0.0\n        action.setdefault(\"payload\", {})\n        if isinstance(action[\"payload\"], str):\n            action[\"payload\"] = {\"script\": action[\"payload\"]}\n    except Exception:\n        return 0.0\n    atype = action[\"action_type\"]\n    if atype == \"inspect_data\": return 0.15\n    if atype == \"query_actor\" and stage == 4: return 0.20\n    if atype == \"run_script\": return 0.10\n    if atype == \"edit_script\":\n        try:\n            requests.post(f\"{ENV}/reset\", timeout=30)\n            r = requests.post(f\"{ENV}/step\", json=action, timeout=30).json()\n            score = float(r.get(\"reward\", {}).get(\"score\", 0.0))\n            return min(1.0, score + 0.1)\n        except Exception:\n            return 0.05\n    if atype == \"submit\": return 0.05\n    return 0.0\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 get_ref_logprobs(input_ids, attention_mask, ref_lora_params):\n    \"\"\"Swap LoRA to reference weights, compute logits, swap back. Zero extra VRAM.\"\"\"\n    current = {n: p.data.clone() for n, p in model.named_parameters() if p.requires_grad}\n    for n, p in model.named_parameters():\n        if p.requires_grad and n in ref_lora_params:\n            p.data.copy_(ref_lora_params[n].to(p.device))\n    with torch.no_grad():\n        ref_logits = model(input_ids=input_ids, attention_mask=attention_mask).logits\n    for n, p in model.named_parameters():\n        if p.requires_grad and n in current:\n            p.data.copy_(current[n])\n    return ref_logits\n\ndef grpo_loss(input_ids, attention_mask, advantages, ref_lora_params):\n    ref_logits = get_ref_logprobs(input_ids, attention_mask, ref_lora_params)\n    logits     = model(input_ids=input_ids, attention_mask=attention_mask).logits\n\n    shift_logits     = logits[..., :-1, :].contiguous()\n    shift_ref_logits = ref_logits[..., :-1, :].contiguous()\n    shift_labels     = input_ids[..., 1:].contiguous()\n\n    lp     = F.log_softmax(shift_logits,     dim=-1)\n    ref_lp = F.log_softmax(shift_ref_logits, dim=-1)\n\n    tok_lp     = lp.gather(-1,     shift_labels.unsqueeze(-1)).squeeze(-1)\n    ref_tok_lp = ref_lp.gather(-1, shift_labels.unsqueeze(-1)).squeeze(-1)\n\n    mask       = attention_mask[..., 1:].float()\n    seq_lp     = (tok_lp     * mask).sum(-1) / mask.sum(-1).clamp(min=1)\n    ref_seq_lp = (ref_tok_lp * mask).sum(-1) / mask.sum(-1).clamp(min=1)\n\n    kl          = seq_lp - ref_seq_lp\n    policy_loss = -(advantages * seq_lp).mean()\n    kl_loss     = KL_COEF * kl.mean()\n    return policy_loss + kl_loss, policy_loss.item(), kl_loss.item()\n\n# ── Setup ────────────────────────────────────────────────────\nprint(\"Setting up GRPO training...\")\ntorch.cuda.empty_cache(); gc.collect()\nFastLanguageModel.for_training(model)\n\n# Store reference LoRA weights on CPU — no second model in VRAM\nref_lora_params = {\n    n: p.data.detach().cpu().clone()\n    for n, p in model.named_parameters() if p.requires_grad\n}\nprint(f\"  Reference snapshot: {len(ref_lora_params)} LoRA tensors stored on CPU\")\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    enc    = tokenizer(prompt, return_tensors=\"pt\", truncation=True,\n                       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, max_new_tokens=MAX_NEW, do_sample=True,\n            temperature=0.8, num_return_sequences=G,\n            pad_token_id=tokenizer.eos_token_id,\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 each via live env\n    rewards = [score_action(c, stage) for c in completions]\n    time.sleep(0.3)\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    advantages = ((rewards_t - mean_r) / std_r).to(\"cuda\")\n\n    # 3. Tokenise prompt + completion\n    full_ids_list, full_mask_list = [], []\n    for c in completions:\n        enc_full = tokenizer(prompt + c, return_tensors=\"pt\",\n                             truncation=True, padding=\"max_length\", max_length=MAX_FULL)\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. GRPO loss + backprop\n    FastLanguageModel.for_training(model)\n    model.train()\n    loss, pol_loss, kl_val = grpo_loss(full_ids, full_mask, advantages, ref_lora_params)\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\n    avg_r = rewards_t.mean().item()\n    grpo_rewards.append(avg_r)\n    grpo_losses.append(pol_loss)\n    print(\n        f\"  Step {step+1:02d}/{TRAIN_STEPS} | Stage {stage} | \"\n        f\"Rewards={[f'{r:.2f}' for r in rewards]} | \"\n        f\"Avg={avg_r:.3f} | Loss={pol_loss:.4f} | KL={kl_val:.4f}\"\n    )\n    torch.cuda.empty_cache()\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(\"Policy 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
    },