fix(grpo): harden rollout reward injection on cloud runtime

Add direct rollout probe fallback in PMOpsGRPOTrainer when cache is empty, and update train_v3 GRPO setup for rollout compatibility (disable fast RL patch/vLLM path, add trainer assertions and preflight probe).

training/pm_ops_trainer.py

@@ -105,6 +105,15 @@ class PMOpsGRPOTrainer(GRPOTrainer):
         completion_ids_list,
     ):
         """Inject cached rewards when available; fall back to reward_funcs otherwise."""
+        def _coerce_rewards(raw):
+            if raw is None:
+                return []
+            if isinstance(raw, torch.Tensor):
+                return [float(r) for r in raw.detach().cpu().flatten().tolist()]
+            if isinstance(raw, (int, float)):
+                return [float(raw)]
+            return [float(r) for r in list(raw)]
+
         cache = self._rollout_reward_cache
         n = len(completions)
 
@@ -120,6 +129,20 @@ class PMOpsGRPOTrainer(GRPOTrainer):
             if input_rewards:
                 cache = input_rewards
 
+        # Cloud fallback: some patched runtimes skip the normal rollout capture path
+        # before calling _calculate_rewards. Actively invoke rollout_func once here
+        # to populate reward cache from the exact prompt batch.
+        if not cache and self.rollout_func is not None and prompts:
+            try:
+                print("[PMOpsGRPOTrainer] cache empty — probing rollout_func for rewards")
+                out = self.rollout_func(prompts, trainer=self)
+                cache = self._rollout_reward_cache
+                if not cache and isinstance(out, dict):
+                    cache = _coerce_rewards(out.get("reward", out.get("rewards", [])))
+                    self._rollout_reward_cache = cache
+            except Exception as exc:
+                print(f"[PMOpsGRPOTrainer] rollout probe failed: {exc!r}")
+
         if cache:
             # Handle num_generations > 1: TRL may call with n > len(cache)
             if len(cache) == n:

training/train_v3.ipynb

@@ -1,17 +1,4 @@
 {
- "nbformat": 4,
- "nbformat_minor": 5,
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "name": "python",
-   "version": "3.11.0"
-  }
- },
  "cells": [
   {
    "cell_type": "markdown",
@@ -47,10 +34,10 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "cell-1",
    "metadata": {},
    "outputs": [],
-   "execution_count": null,
    "source": [
     "# Unsloth + vLLM first — let Unsloth resolve torch compat\n",
     "!pip install -q unsloth vllm\n",
@@ -73,16 +60,22 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "cell-2",
    "metadata": {},
    "outputs": [],
-   "execution_count": null,
    "source": [
     "import torch\n",
     "\n",
-    "# Patch GRPO BEFORE importing GRPOTrainer\n",
+    "# PatchFastRL can bypass custom rollout_func on some cloud runtimes.\n",
+    "# Keep it disabled for PMOpsGRPOTrainer rollout compatibility.\n",
     "from unsloth import FastLanguageModel, PatchFastRL\n",
-    "PatchFastRL('GRPO', FastLanguageModel)\n",
+    "ENABLE_FAST_RL_PATCH = False\n",
+    "if ENABLE_FAST_RL_PATCH:\n",
+    "    PatchFastRL('GRPO', FastLanguageModel)\n",
+    "    print('PatchFastRL enabled')\n",
+    "else:\n",
+    "    print('PatchFastRL disabled for rollout_func compatibility')\n",
     "\n",
     "import trl\n",
     "print(f'torch : {torch.__version__}')\n",
@@ -115,10 +108,10 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "cell-3",
    "metadata": {},
    "outputs": [],
-   "execution_count": null,
    "source": [
     "import os, sys\n",
     "\n",
@@ -149,10 +142,10 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "cell-4",
    "metadata": {},
    "outputs": [],
-   "execution_count": null,
    "source": [
     "from huggingface_hub import notebook_login\n",
     "notebook_login()"
@@ -168,10 +161,10 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "cell-5",
    "metadata": {},
    "outputs": [],
-   "execution_count": null,
    "source": [
     "import subprocess, time, requests\n",
     "\n",
@@ -204,10 +197,10 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "cell-6",
    "metadata": {},
    "outputs": [],
-   "execution_count": null,
    "source": [
     "import trl.experimental.openenv  # must be importable\n",
     "from openenv.core import GenericEnvClient\n",
@@ -231,10 +224,10 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "cell-7",
    "metadata": {},
    "outputs": [],
-   "execution_count": null,
    "source": [
     "MODEL_NAME = 'Qwen/Qwen3-1.7B'\n",
     "LORA_RANK  = 16\n",
@@ -243,7 +236,7 @@
     "    model_name           = MODEL_NAME,\n",
     "    max_seq_length       = 4096 + MAX_COMP_LEN,\n",
     "    load_in_4bit         = True,\n",
-    "    fast_inference       = True,   # enables vLLM path for GRPO rollouts\n",
+    "    fast_inference       = False,  # disable fast RL path to preserve rollout_func behaviour\n",
     "    max_lora_rank        = LORA_RANK,\n",
     "    gpu_memory_utilization = 0.50, # leave headroom for SFT activations\n",
     ")\n",
@@ -255,7 +248,7 @@
     "    lora_alpha               = LORA_RANK,\n",
     "    use_gradient_checkpointing = 'unsloth',\n",
     "    random_state             = 42,\n",
-    ")\n",
+    " )\n",
     "tokenizer.pad_token    = tokenizer.eos_token\n",
     "tokenizer.padding_side = 'left'\n",
     "model.print_trainable_parameters()\n",
@@ -291,10 +284,10 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "cell-9",
    "metadata": {},
    "outputs": [],
-   "execution_count": null,
    "source": [
     "import json as _json\n",
     "from datasets import Dataset\n",
@@ -373,10 +366,10 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "cell-10",
    "metadata": {},
    "outputs": [],
-   "execution_count": null,
    "source": [
     "from trl import SFTTrainer, SFTConfig\n",
     "\n",
@@ -423,10 +416,10 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "cell-11",
    "metadata": {},
    "outputs": [],
-   "execution_count": null,
    "source": [
     "from training.rollout import extract_json_action, _obs_to_dict, _current_obs_text, build_messages\n",
     "\n",
@@ -492,10 +485,10 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "cell-13",
    "metadata": {},
    "outputs": [],
-   "execution_count": null,
    "source": [
     "from training.dataset import generate_triage_dataset\n",
     "\n",
@@ -515,10 +508,10 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "cell-14",
    "metadata": {},
    "outputs": [],
-   "execution_count": null,
    "source": [
     "from trl.experimental.openenv import generate_rollout_completions\n",
     "from training.rollout import (\n",
@@ -630,11 +623,71 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "cell-15",
    "metadata": {},
    "outputs": [],
+   "source": [
+    "from trl import GRPOConfig\n",
+    "from training.pm_ops_trainer import PMOpsGRPOTrainer\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",
+    "    # Keep disabled for rollout_func compatibility on cloud runtimes\n",
+    "    use_vllm                      = False,\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",
+    "# PMOpsGRPOTrainer overrides _calculate_rewards to read the pre-computed\n",
+    "# 'reward' key directly from the rollout batch — bypasses broken kwargs plumbing.\n",
+    "trainer = 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",
+    ")\n",
+    "print(f'PMOpsGRPOTrainer ready: {type(trainer).__name__}')\n",
+    "assert isinstance(trainer, PMOpsGRPOTrainer), 'trainer must be PMOpsGRPOTrainer'\n",
+    "assert grpo_cfg.use_vllm is False, 'use_vllm must be False for rollout compatibility'"
+   ]
+  },
+  {
+   "cell_type": "code",
    "execution_count": null,
-   "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')"
+   "id": "9a7c99b4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Preflight: ensure rollout returns reward and trainer received rollout_func\n",
+    "probe = grpo_rollout_func([grpo_dataset[0]['prompt']], trainer=trainer)\n",
+    "print('probe keys:', list(probe.keys()))\n",
+    "print('probe reward sample:', probe['reward'][:1])\n",
+    "assert len(probe['reward']) == 1, 'rollout probe did not return reward values'"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -653,10 +706,10 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "cell-16",
    "metadata": {},
    "outputs": [],
-   "execution_count": null,
    "source": [
     "trainer_stats = trainer.train()\n",
     "\n",
@@ -678,10 +731,10 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "cell-17",
    "metadata": {},
    "outputs": [],
-   "execution_count": null,
    "source": [
     "grpo_env.close()\n",
     "\n",
@@ -702,10 +755,10 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "cell-18",
    "metadata": {},
    "outputs": [],
-   "execution_count": null,
    "source": [
     "from transformers import AutoModelForCausalLM\n",
     "\n",
@@ -797,10 +850,10 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "cell-19",
    "metadata": {},
    "outputs": [],
-   "execution_count": null,
    "source": [
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
@@ -841,14 +894,27 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "id": "cell-20",
    "metadata": {},
    "outputs": [],
-   "execution_count": null,
    "source": [
     "server_proc.terminate()\n",
     "print('Local PM-Ops server stopped')"
    ]
   }
- ]
-}
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}