fix(train_v3): use PMOpsGRPOTrainer to fix reward kwarg plumbing

PatchFastRL strips custom rollout keys before calling reward_func,
so the 'reward' key never arrived. PMOpsGRPOTrainer overrides
_calculate_rewards to inject pre-computed rewards from the rollout
batch directly, bypassing the broken kwargs path entirely.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

training/train_v3.ipynb

@@ -2,8 +2,15 @@
  "nbformat": 4,
  "nbformat_minor": 5,
  "metadata": {
-  "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},
-  "language_info": {"name": "python", "version": "3.11.0"}
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11.0"
+  }
  },
  "cells": [
   {
@@ -34,7 +41,9 @@
    "cell_type": "markdown",
    "id": "cell-1-md",
    "metadata": {},
-   "source": ["## 0. Install"]
+   "source": [
+    "## 0. Install"
+   ]
   },
   {
    "cell_type": "code",
@@ -58,7 +67,9 @@
    "cell_type": "markdown",
    "id": "cell-2-md",
    "metadata": {},
-   "source": ["## 1. Imports + GPU Config"]
+   "source": [
+    "## 1. Imports + GPU Config"
+   ]
   },
   {
    "cell_type": "code",
@@ -98,7 +109,9 @@
    "cell_type": "markdown",
    "id": "cell-3-md",
    "metadata": {},
-   "source": ["## 2. Clone PM-Ops Repo"]
+   "source": [
+    "## 2. Clone PM-Ops Repo"
+   ]
   },
   {
    "cell_type": "code",
@@ -130,7 +143,9 @@
    "cell_type": "markdown",
    "id": "cell-4-md",
    "metadata": {},
-   "source": ["## 3. HuggingFace Login"]
+   "source": [
+    "## 3. HuggingFace Login"
+   ]
   },
   {
    "cell_type": "code",
@@ -147,7 +162,9 @@
    "cell_type": "markdown",
    "id": "cell-5-md",
    "metadata": {},
-   "source": ["## 4. Start Local PM-Ops Server"]
+   "source": [
+    "## 4. Start Local PM-Ops Server"
+   ]
   },
   {
    "cell_type": "code",
@@ -181,7 +198,9 @@
    "cell_type": "markdown",
    "id": "cell-6-md",
    "metadata": {},
-   "source": ["## 5. Verify Env"]
+   "source": [
+    "## 5. Verify Env"
+   ]
   },
   {
    "cell_type": "code",
@@ -206,7 +225,9 @@
    "cell_type": "markdown",
    "id": "cell-7-md",
    "metadata": {},
-   "source": ["## 6. Load Model — Unsloth 4-bit + LoRA"]
+   "source": [
+    "## 6. Load Model — Unsloth 4-bit + LoRA"
+   ]
   },
   {
    "cell_type": "code",
@@ -264,7 +285,9 @@
    "cell_type": "markdown",
    "id": "cell-9-md",
    "metadata": {},
-   "source": ["## 7. Generate SFT Demonstration Dataset"]
+   "source": [
+    "## 7. Generate SFT Demonstration Dataset"
+   ]
   },
   {
    "cell_type": "code",
@@ -344,7 +367,9 @@
    "cell_type": "markdown",
    "id": "cell-10-md",
    "metadata": {},
-   "source": ["## 8. SFT Training (~15 min on A100)"]
+   "source": [
+    "## 8. SFT Training (~15 min on A100)"
+   ]
   },
   {
    "cell_type": "code",
@@ -392,7 +417,9 @@
    "cell_type": "markdown",
    "id": "cell-11-md",
    "metadata": {},
-   "source": ["## 9. Verify SFT Output — Model Must Output Valid JSON"]
+   "source": [
+    "## 9. Verify SFT Output — Model Must Output Valid JSON"
+   ]
   },
   {
    "cell_type": "code",
@@ -459,7 +486,9 @@
    "cell_type": "markdown",
    "id": "cell-13-md",
    "metadata": {},
-   "source": ["## 10. Generate GRPO Training Dataset"]
+   "source": [
+    "## 10. Generate GRPO Training Dataset"
+   ]
   },
   {
    "cell_type": "code",
@@ -480,7 +509,9 @@
    "cell_type": "markdown",
    "id": "cell-14-md",
    "metadata": {},
-   "source": ["## 11. GRPO Rollout + Reward Functions"]
+   "source": [
+    "## 11. GRPO Rollout + Reward Functions"
+   ]
   },
   {
    "cell_type": "code",
@@ -593,7 +624,9 @@
    "cell_type": "markdown",
    "id": "cell-15-md",
    "metadata": {},
-   "source": ["## 12. GRPO Config + Trainer"]
+   "source": [
+    "## 12. GRPO Config + Trainer"
+   ]
   },
   {
    "cell_type": "code",
@@ -601,48 +634,7 @@
    "metadata": {},
    "outputs": [],
    "execution_count": null,
-   "source": [
-    "from trl import GRPOConfig, GRPOTrainer\n",
-    "\n",
-    "OUTPUT_DIR = 'pm-ops-grpo-Qwen3-1.7B-triage-v3'\n",
-    "HF_REPO_ID = f'Saurav1/{OUTPUT_DIR}'\n",
-    "\n",
-    "grpo_cfg = GRPOConfig(\n",
-    "    # Training\n",
-    "    num_train_epochs              = 2,\n",
-    "    learning_rate                 = 1e-6,   # lower LR: model has SFT init, don't overwrite it\n",
-    "    gradient_accumulation_steps   = GRAD_ACCUM,\n",
-    "    per_device_train_batch_size   = 1,\n",
-    "    warmup_steps                  = 5,\n",
-    "    num_generations               = NUM_GEN,\n",
-    "    # Sequence lengths\n",
-    "    max_completion_length         = MAX_COMP_LEN,\n",
-    "    max_prompt_length             = 4096,\n",
-    "    # Unsloth vLLM for fast generation\n",
-    "    use_vllm                      = True,\n",
-    "    # Output\n",
-    "    output_dir                    = OUTPUT_DIR,\n",
-    "    report_to                     = 'trackio',\n",
-    "    trackio_space_id              = OUTPUT_DIR,\n",
-    "    logging_steps                 = 1,\n",
-    "    save_steps                    = 20,\n",
-    "    gradient_checkpointing        = False,  # Unsloth handles this\n",
-    ")\n",
-    "\n",
-    "eff_batch   = grpo_cfg.per_device_train_batch_size * GRAD_ACCUM\n",
-    "total_steps = len(grpo_dataset) * NUM_GEN * grpo_cfg.num_train_epochs // eff_batch\n",
-    "print(f'GRPO: {len(grpo_dataset)} eps x {NUM_GEN} gen x {grpo_cfg.num_train_epochs} epochs -> ~{total_steps} steps')\n",
-    "\n",
-    "trainer = GRPOTrainer(\n",
-    "    model            = model,\n",
-    "    processing_class = tokenizer,\n",
-    "    reward_funcs     = grpo_reward_func,\n",
-    "    train_dataset    = grpo_dataset,\n",
-    "    args             = grpo_cfg,\n",
-    "    rollout_func     = grpo_rollout_func,\n",
-    ")\n",
-    "print('GRPOTrainer ready')"
-   ]
+   "source": "from trl import GRPOConfig\nfrom training.pm_ops_trainer import PMOpsGRPOTrainer\n\nOUTPUT_DIR = 'pm-ops-grpo-Qwen3-1.7B-triage-v3'\nHF_REPO_ID = f'Saurav1/{OUTPUT_DIR}'\n\ngrpo_cfg = GRPOConfig(\n    # Training\n    num_train_epochs              = 2,\n    learning_rate                 = 1e-6,   # lower LR: model has SFT init, don't overwrite it\n    gradient_accumulation_steps   = GRAD_ACCUM,\n    per_device_train_batch_size   = 1,\n    warmup_steps                  = 5,\n    num_generations               = NUM_GEN,\n    # Sequence lengths\n    max_completion_length         = MAX_COMP_LEN,\n    max_prompt_length             = 4096,\n    # Unsloth vLLM for fast generation\n    use_vllm                      = True,\n    # Output\n    output_dir                    = OUTPUT_DIR,\n    report_to                     = 'trackio',\n    trackio_space_id              = OUTPUT_DIR,\n    logging_steps                 = 1,\n    save_steps                    = 20,\n    gradient_checkpointing        = False,  # Unsloth handles this\n)\n\neff_batch   = grpo_cfg.per_device_train_batch_size * GRAD_ACCUM\ntotal_steps = len(grpo_dataset) * NUM_GEN * grpo_cfg.num_train_epochs // eff_batch\nprint(f'GRPO: {len(grpo_dataset)} eps x {NUM_GEN} gen x {grpo_cfg.num_train_epochs} epochs -> ~{total_steps} steps')\n\n# PMOpsGRPOTrainer overrides _calculate_rewards to read the pre-computed\n# 'reward' key directly from the rollout batch — bypasses the broken\n# PatchFastRL kwargs plumbing that strips custom rollout keys.\ntrainer = PMOpsGRPOTrainer(\n    model            = model,\n    processing_class = tokenizer,\n    reward_funcs     = grpo_reward_func,  # kept as fallback only\n    train_dataset    = grpo_dataset,\n    args             = grpo_cfg,\n    rollout_func     = grpo_rollout_func,\n)\nprint('PMOpsGRPOTrainer ready')"
   },
   {
    "cell_type": "markdown",
@@ -680,7 +672,9 @@
    "cell_type": "markdown",
    "id": "cell-17-md",
    "metadata": {},
-   "source": ["## 14. Save Model"]
+   "source": [
+    "## 14. Save Model"
+   ]
   },
   {
    "cell_type": "code",
@@ -702,7 +696,9 @@
    "cell_type": "markdown",
    "id": "cell-18-md",
    "metadata": {},
-   "source": ["## 15. Evaluate: Baseline vs Trained"]
+   "source": [
+    "## 15. Evaluate: Baseline vs Trained"
+   ]
   },
   {
    "cell_type": "code",
@@ -795,7 +791,9 @@
    "cell_type": "markdown",
    "id": "cell-19-md",
    "metadata": {},
-   "source": ["## 16. Plot Results"]
+   "source": [
+    "## 16. Plot Results"
+   ]
   },
   {
    "cell_type": "code",
@@ -837,7 +835,9 @@
    "cell_type": "markdown",
    "id": "cell-20-md",
    "metadata": {},
-   "source": ["## 17. Teardown"]
+   "source": [
+    "## 17. Teardown"
+   ]
   },
   {
    "cell_type": "code",
@@ -851,4 +851,4 @@
    ]
   }
  ]
-}
+}