Merge remote-tracking branch 'hf/main' into train/v1

training/pm_ops_trainer.py

@@ -0,0 +1,106 @@
+"""PMOpsGRPOTrainer — GRPOTrainer subclass with direct reward injection.
+
+Problem being solved
+--------------------
+TRL 1.2.0's GRPOTrainer._calculate_rewards() calls reward_funcs with
+**reward_kwargs built from per-example keys in the input batch. This works
+correctly when reward_funcs produce rewards from text completions alone.
+
+It breaks for multi-turn env rollouts: our rollout_func computes rewards
+inside the episode loop (where the env state is available), stores them as
+a 'reward' key in the rollout output dict, and expects reward_func to read
+them back via kwargs. TRL silently drops these keys when building
+reward_kwargs from the batch, so reward_func always receives an empty dict
+and returns 0.0 for every episode.
+
+Fix
+---
+Override _calculate_rewards to check for a pre-computed 'reward' key in the
+input batch. If present, return it directly as a [batch_size, 1] tensor —
+no reward_funcs called, no kwargs needed. If absent, fall back to the
+standard TRL behaviour so the class works as a drop-in replacement.
+
+Usage
+-----
+    from training.pm_ops_trainer import PMOpsGRPOTrainer
+
+    trainer = PMOpsGRPOTrainer(
+        model=model,
+        processing_class=tokenizer,
+        reward_funcs=reward_func,   # kept as-is; only used as fallback
+        train_dataset=dataset,
+        args=grpo_config,
+        rollout_func=rollout_func,  # must put 'reward' key in output dict
+    )
+
+rollout_func contract
+---------------------
+rollout_func(prompts, trainer=None) must return a dict containing at minimum:
+    {
+        "prompt_ids":    list[list[int]],   # one per episode
+        "completion_ids": list[list[int]],
+        "logprobs":      list[list[float]],
+        "reward":        list[float],       # one combined float per episode
+    }
+
+The 'reward' value is the weighted combination of all sub-signals and is
+injected directly as the GRPO advantage signal.
+"""
+
+import torch
+from trl import GRPOTrainer
+
+
+class PMOpsGRPOTrainer(GRPOTrainer):
+    """Drop-in GRPOTrainer replacement with direct reward injection.
+
+    Overrides _calculate_rewards to read pre-computed rewards from the
+    rollout batch instead of calling reward_funcs with broken **kwargs.
+    Falls back to standard GRPOTrainer behaviour if 'reward' key is absent.
+    """
+
+    # Key written by rollout_func and read by _calculate_rewards
+    ROLLOUT_REWARD_KEY = "reward"
+
+    def _calculate_rewards(
+        self,
+        inputs,
+        prompts,
+        completions,
+        completion_ids_list,
+    ):
+        """Override: inject pre-computed rewards when available.
+
+        Args:
+            inputs:              list[dict] — one dict per example in the batch.
+                                 rollout_func output keys land here per example.
+            prompts:             list[str]  — prompt texts (unused in this path)
+            completions:         list[str]  — completion texts (unused)
+            completion_ids_list: list[list[int]] — token IDs (unused)
+
+        Returns:
+            Tensor of shape [batch_size, 1] when pre-computed rewards are found.
+            Falls back to super()._calculate_rewards(...) otherwise.
+        """
+        # Fast-path: pre-computed reward key present in every example
+        if inputs and all(
+            self.ROLLOUT_REWARD_KEY in ex for ex in inputs
+        ):
+            device = self.accelerator.device
+            rewards = torch.tensor(
+                [float(ex[self.ROLLOUT_REWARD_KEY]) for ex in inputs],
+                dtype=torch.float32,
+                device=device,
+            ).unsqueeze(1)  # [batch_size, 1]
+
+            # Log the mean reward so it shows up in training curves
+            mean_r = rewards.mean().item()
+            self.log({"reward/injected_mean": mean_r})
+
+            return rewards
+
+        # Fallback: standard TRL reward_funcs path
+        # Triggered when reward key is absent (e.g. non-rollout evaluation)
+        return super()._calculate_rewards(
+            inputs, prompts, completions, completion_ids_list
+        )

training/rollout.py

@@ -16,11 +16,81 @@ import json
 import re
 from typing import Any
 
-from trl.experimental.openenv import generate_rollout_completions
+import torch
+import torch.nn.functional as F
 
 from training.dataset import parse_seed_from_prompt
 from training.prompts import SYSTEM_PROMPT, format_observation
 
+
+# ---------------------------------------------------------------------------
+# HF model.generate() — replaces generate_rollout_completions (vLLM-only)
+# ---------------------------------------------------------------------------
+
+def _get_model_for_generation(trainer):
+    """Unwrap the model safely regardless of accelerate/PEFT/DDP wrapping.
+
+    Priority:
+      1. accelerator.unwrap_model — handles DDP + PEFT + DeepSpeed
+      2. trainer.model.module     — plain DDP wrapping
+      3. trainer.model            — unwrapped (local or single-GPU)
+    """
+    if hasattr(trainer, "accelerator"):
+        return trainer.accelerator.unwrap_model(trainer.model)
+    if hasattr(trainer.model, "module"):
+        return trainer.model.module
+    return trainer.model
+
+
+def _generate_no_vllm(trainer, prompt_text: str, tokenizer, max_new_tokens: int = 512) -> dict:
+    """Generate one completion using HF model.generate() without vLLM.
+
+    Returns the same dict shape as generate_rollout_completions so the rest
+    of rollout_once is unchanged:
+        prompt_ids:     list[int]
+        completion_ids: list[int]
+        logprobs:       list[float]  (per-token log-prob under current policy)
+        text:           str
+    """
+    model = _get_model_for_generation(trainer)
+
+    # Device: prefer accelerator.device, fall back to first param device
+    if hasattr(trainer, "accelerator"):
+        device = trainer.accelerator.device
+    else:
+        device = next(model.parameters()).device
+
+    enc = tokenizer(prompt_text, return_tensors="pt").to(device)
+    prompt_len = enc["input_ids"].shape[1]
+
+    with torch.no_grad():
+        out = model.generate(
+            **enc,
+            max_new_tokens=max_new_tokens,
+            do_sample=True,
+            temperature=0.7,
+            pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
+            output_scores=True,
+            return_dict_in_generate=True,
+        )
+
+    prompt_ids     = enc["input_ids"][0].tolist()
+    completion_ids = out.sequences[0][prompt_len:].tolist()
+
+    # Per-token log-probs from output.scores (one score tensor per new token)
+    logprobs = [
+        F.log_softmax(score[0], dim=-1)[tok_id].item()
+        for score, tok_id in zip(out.scores, completion_ids)
+    ]
+
+    text = tokenizer.decode(completion_ids, skip_special_tokens=True)
+    return {
+        "prompt_ids":     prompt_ids,
+        "completion_ids": completion_ids,
+        "logprobs":       logprobs,
+        "text":           text,
+    }
+
 MAX_STEPS = 40
 # ~3000 tokens at 4 chars/token; leaves room for completion tokens
 MAX_PROMPT_CHARS = 12_000
@@ -209,14 +279,12 @@ def rollout_once(
             enable_thinking=False,
         )
 
-        rollout_out = generate_rollout_completions(trainer, [prompt_text])[0]
+        rollout_out = _generate_no_vllm(trainer, prompt_text, tokenizer)
         prompt_ids.extend(rollout_out["prompt_ids"])
         completion_ids.extend(rollout_out["completion_ids"])
         logprobs.extend(rollout_out["logprobs"])
 
-        completion_text = rollout_out.get("text") or tokenizer.decode(
-            rollout_out["completion_ids"], skip_special_tokens=True
-        )
+        completion_text = rollout_out["text"]
 
         # Parse action; fall back gracefully on parse failure
         parsed = extract_json_action(completion_text)
@@ -287,18 +355,34 @@ def rollout_once(
         no_wrong_channels = 1.0
 
     read_runbook_reward = 1.0 if read_runbook_done else 0.0
-    combined = (
-        final_score       * 0.45
-        + no_wrong_channels * 0.15
-        + valid_json_ratio  * 0.15
-        + read_runbook_reward * 0.15
-        + efficiency        * 0.10
-    )
+
+    # --- Reward gating ---
+    # If model never output valid JSON, it never actually tried anything.
+    # Strip all process rewards and apply a harsh penalty.
+    # Process rewards (runbook, no_wrong, efficiency) only matter if model acted.
+    if valid_action_count == 0:
+        combined = -1.0
+    elif final_score == 0.0:
+        # Model tried (valid JSON) but task failed — small process credit, no final bonus
+        combined = (
+            valid_json_ratio    * 0.15
+            + read_runbook_reward * 0.10
+            - 0.30                        # hard penalty for zero task completion
+        )
+    else:
+        combined = (
+            final_score           * 0.45
+            + no_wrong_channels   * 0.15
+            + valid_json_ratio    * 0.15
+            + read_runbook_reward * 0.15
+            + efficiency          * 0.10
+        )
+
     print(
         f"[rollout] steps={step} final={final_score:.3f} "
         f"json={valid_json_ratio:.2f} runbook={read_runbook_reward:.0f} "
         f"no_wrong={no_wrong_channels:.2f} eff={efficiency:.2f} "
-        f"→ combined={combined:.3f}"
+        f"valid_acts={valid_action_count} → combined={combined:.3f}"
     )
     return {
         "prompt_ids": prompt_ids,

training/train_v2.ipynb

@@ -21,7 +21,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 0. Install Dependencies"]
+   "source": [
+    "## 0. Install Dependencies"
+   ]
   },
   {
    "cell_type": "code",
@@ -47,7 +49,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 1. Version Check + GPU Detect"]
+   "source": [
+    "## 1. Version Check + GPU Detect"
+   ]
   },
   {
    "cell_type": "code",
@@ -81,7 +85,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 2. Clone PM-Ops Repo"]
+   "source": [
+    "## 2. Clone PM-Ops Repo"
+   ]
   },
   {
    "cell_type": "code",
@@ -111,7 +117,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 3. HuggingFace Login"]
+   "source": [
+    "## 3. HuggingFace Login"
+   ]
   },
   {
    "cell_type": "code",
@@ -161,7 +169,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 5. Verify Environment"]
+   "source": [
+    "## 5. Verify Environment"
+   ]
   },
   {
    "cell_type": "code",
@@ -239,7 +249,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 7. Generate Training Dataset"]
+   "source": [
+    "## 7. Generate Training Dataset"
+   ]
   },
   {
    "cell_type": "code",
@@ -375,7 +387,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 10. Configure GRPO Training"]
+   "source": [
+    "## 10. Configure GRPO Training"
+   ]
   },
   {
    "cell_type": "code",
@@ -423,7 +437,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 11. Create Trainer"]
+   "source": [
+    "## 11. Create Trainer"
+   ]
   },
   {
    "cell_type": "code",
@@ -431,17 +447,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from trl import GRPOTrainer\n",
+    "from training.pm_ops_trainer import PMOpsGRPOTrainer\n",
     "\n",
-    "trainer = GRPOTrainer(\n",
+    "trainer = PMOpsGRPOTrainer(\n",
     "    model=model,\n",
     "    processing_class=tokenizer,\n",
-    "    reward_funcs=reward_func,\n",
+    "    reward_funcs=reward_func,   # fallback only — injected rewards take priority\n",
     "    train_dataset=dataset,\n",
     "    args=grpo_config,\n",
     "    rollout_func=rollout_func,\n",
     ")\n",
-    "print('GRPOTrainer ready')"
+    "print(\"PMOpsGRPOTrainer ready — direct reward injection active\")"
    ]
   },
   {
@@ -485,7 +501,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 13. Save + Push"]
+   "source": [
+    "## 13. Save + Push"
+   ]
   },
   {
    "cell_type": "code",
@@ -509,7 +527,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 14. Merge LoRA → bf16 (Optional — for full-weight inference)"]
+   "source": [
+    "## 14. Merge LoRA → bf16 (Optional — for full-weight inference)"
+   ]
   },
   {
    "cell_type": "code",
@@ -533,7 +553,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 15. Evaluate: Baseline vs Trained"]
+   "source": [
+    "## 15. Evaluate: Baseline vs Trained"
+   ]
   },
   {
    "cell_type": "code",
@@ -621,7 +643,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 16. Plot"]
+   "source": [
+    "## 16. Plot"
+   ]
   },
   {
    "cell_type": "code",
@@ -660,7 +684,9 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": ["## 17. Teardown"]
+   "source": [
+    "## 17. Teardown"
+   ]
   },
   {
    "cell_type": "code",
@@ -674,9 +700,16 @@
   }
  ],
  "metadata": {
-  "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},
-  "language_info": {"name": "python", "version": "3.12.0"}
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.0"
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}