{ "cells": [ { "cell_type": "markdown", "id": "cell-0", "metadata": {}, "source": [ "# PM-Ops GRPO Training v3 — SFT Warmup + GRPO\n", "\n", "## What changed from v2\n", "| Problem in v2 | Fix in v3 |\n", "|---|---|\n", "| `json=0.00` — model never outputs JSON, GRPO gradient = 0 | **SFT warmup on 60 baseline traces first** |\n", "| Reward = 0.300 every episode (no variance) | **env_score varies 0.25–1.0 after SFT** |\n", "| Slow — no vLLM, eager attention | **Unsloth + `use_vllm=True`** |\n", "| 5 reward signals, all from fallback | **2 signals: env_score × 0.85 + json_ratio × 0.15** |\n", "| `meta.finish` never called in 15-step cap | **Fixed fallback respects `max_steps`** |\n", "\n", "## Why SFT before GRPO?\n", "GRPO learns by comparing rewards across a *group* of generations. If all generations get the\n", "same reward (because the model outputs garbage on every step), the advantage is 0 and weights\n", "don't move. SFT warmup costs ~15 min and unlocks the full GRPO gradient.\n", "\n", "**Stack**: Unsloth + TRL 1.2.0 + OpenEnv · **GPU**: A100 → ~75 min total" ] }, { "cell_type": "markdown", "id": "cell-1-md", "metadata": {}, "source": [ "## 0. Install" ] }, { "cell_type": "code", "execution_count": null, "id": "cell-1", "metadata": {}, "outputs": [], "source": [ "# Unsloth + vLLM first — let Unsloth resolve torch compat\n", "!pip install -q unsloth vllm\n", "# TRL training stack\n", "!pip install -q \"trl==1.2.0\" accelerate datasets\n", "# PM-Ops server runtime\n", "!pip install -q \"openenv-core>=0.2.2\" \"fastapi>=0.110.0\" \"uvicorn[standard]>=0.29.0\" \"pydantic>=2.0.0\"\n", "# Experiment tracking\n", "!pip install -q trackio\n", "print('Done - restart kernel, then run all cells from top.')" ] }, { "cell_type": "markdown", "id": "cell-2-md", "metadata": {}, "source": [ "## 1. Imports + GPU Config" ] }, { "cell_type": "code", "execution_count": null, "id": "cell-2", "metadata": {}, "outputs": [], "source": [ "import torch\n", "\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", "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", "print(f'TRL : {trl.__version__}')\n", "\n", "gpu = torch.cuda.get_device_properties(0)\n", "TOTAL_GB = round(gpu.total_memory / 1024**3, 1)\n", "IS_A100 = TOTAL_GB >= 35\n", "print(f'GPU : {gpu.name} ({TOTAL_GB} GB)')\n", "\n", "# Adaptive config — T4 uses minimal settings for smoke-testing\n", "NUM_GEN = 6 if IS_A100 else 2\n", "GRAD_ACCUM = 32 if IS_A100 else 8\n", "MAX_COMP_LEN = 384\n", "N_SFT_EPISODES = 60 if IS_A100 else 15\n", "N_GRPO_EPISODES = 150 if IS_A100 else 30\n", "TRAIN_MAX_STEPS = 12 # triage solvable in 5; 12 gives exploration room\n", "\n", "print(f'num_gen={NUM_GEN} grad_accum={GRAD_ACCUM} '\n", " f'sft_eps={N_SFT_EPISODES} grpo_eps={N_GRPO_EPISODES}')" ] }, { "cell_type": "markdown", "id": "cell-3-md", "metadata": {}, "source": [ "## 2. Clone PM-Ops Repo" ] }, { "cell_type": "code", "execution_count": null, "id": "cell-3", "metadata": {}, "outputs": [], "source": [ "import os, sys\n", "\n", "REPO_URL = 'https://huggingface.co/spaces/TheCrustaceans/Pm-ops'\n", "REPO_DIR = '/content/Pm_ops'\n", "\n", "if not os.path.exists(REPO_DIR):\n", " !git clone --depth=1 -q {REPO_URL} {REPO_DIR}\n", " print(f'Cloned -> {REPO_DIR}')\n", "else:\n", " !git -C {REPO_DIR} pull -q origin main\n", " print(f'Pulled -> {REPO_DIR}')\n", "\n", "for p in [REPO_DIR, os.path.join(REPO_DIR, 'training')]:\n", " if p not in sys.path:\n", " sys.path.insert(0, p)\n", "os.chdir(REPO_DIR)\n", "print(f'CWD: {os.getcwd()}')" ] }, { "cell_type": "markdown", "id": "cell-4-md", "metadata": {}, "source": [ "## 3. HuggingFace Login" ] }, { "cell_type": "code", "execution_count": null, "id": "cell-4", "metadata": {}, "outputs": [], "source": [ "from huggingface_hub import notebook_login\n", "notebook_login()" ] }, { "cell_type": "markdown", "id": "cell-5-md", "metadata": {}, "source": [ "## 4. Start Local PM-Ops Server" ] }, { "cell_type": "code", "execution_count": null, "id": "cell-5", "metadata": {}, "outputs": [], "source": [ "import subprocess, time, requests\n", "\n", "server_proc = subprocess.Popen(\n", " [sys.executable, '-m', 'uvicorn', 'server.app:app',\n", " '--host', '0.0.0.0', '--port', '8000'],\n", " cwd=REPO_DIR, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,\n", ")\n", "ENV_URL = 'http://localhost:8000'\n", "\n", "for _ in range(30):\n", " try:\n", " if requests.get(f'{ENV_URL}/', timeout=2).status_code == 200:\n", " print(f'PM-Ops server ready pid={server_proc.pid}')\n", " break\n", " except Exception:\n", " pass\n", " time.sleep(1)\n", "else:\n", " raise RuntimeError('Server did not start in 30 s')" ] }, { "cell_type": "markdown", "id": "cell-6-md", "metadata": {}, "source": [ "## 5. Verify Env" ] }, { "cell_type": "code", "execution_count": null, "id": "cell-6", "metadata": {}, "outputs": [], "source": [ "import trl.experimental.openenv # must be importable\n", "from openenv.core import GenericEnvClient\n", "from training.rollout import _obs_to_dict\n", "\n", "with GenericEnvClient(base_url=ENV_URL).sync() as _env:\n", " r = _env.reset()\n", " obs = _obs_to_dict(r.observation if hasattr(r, 'observation') else r)\n", " print(f'task_brief : {obs.get(\"task_brief\", \"?\")[:80]}...')\n", " _env.step({'action_type': 'meta.read_runbook', 'args': {}})\n", " print('Env step : OK')" ] }, { "cell_type": "markdown", "id": "cell-7-md", "metadata": {}, "source": [ "## 6. Load Model — Unsloth 4-bit + LoRA" ] }, { "cell_type": "code", "execution_count": null, "id": "cell-7", "metadata": {}, "outputs": [], "source": [ "MODEL_NAME = 'Qwen/Qwen3-1.7B'\n", "LORA_RANK = 16\n", "\n", "model, tokenizer = FastLanguageModel.from_pretrained(\n", " model_name = MODEL_NAME,\n", " max_seq_length = 4096 + MAX_COMP_LEN,\n", " load_in_4bit = True,\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", "model = FastLanguageModel.get_peft_model(\n", " model,\n", " r = LORA_RANK,\n", " target_modules = ['q_proj','k_proj','v_proj','o_proj',\n", " 'gate_proj','up_proj','down_proj'],\n", " lora_alpha = LORA_RANK,\n", " use_gradient_checkpointing = 'unsloth',\n", " random_state = 42,\n", " )\n", "tokenizer.pad_token = tokenizer.eos_token\n", "tokenizer.padding_side = 'left'\n", "model.print_trainable_parameters()\n", "reserved = round(torch.cuda.max_memory_reserved() / 1024**3, 2)\n", "print(f'GPU after load: {reserved} GB / {TOTAL_GB} GB')" ] }, { "cell_type": "markdown", "id": "cell-8-md", "metadata": {}, "source": [ "---\n", "## Phase 1 — SFT Warmup\n", "\n", "**Goal**: teach the model the JSON output format and PM-ops workflow before GRPO.\n", "\n", "We run `baseline_agent` (the deterministic heuristic) for 60 episodes and record every\n", "(observation, action) pair as a supervised example. Each episode produces ~6 steps:\n", "`read_runbook` → `create_ticket` → `assign_ticket` → `list_channels` → `post_message` → `finish`.\n", "\n", "After 2 SFT epochs (~15 min), the model reliably outputs `\\`\\`\\`json ... \\`\\`\\`` blocks.\n", "Without this, GRPO reward variance ≈ 0 and nothing is learned." ] }, { "cell_type": "markdown", "id": "cell-9-md", "metadata": {}, "source": [ "## 7. Generate SFT Demonstration Dataset" ] }, { "cell_type": "code", "execution_count": null, "id": "cell-9", "metadata": {}, "outputs": [], "source": [ "import json as _json\n", "from datasets import Dataset\n", "from inference import baseline_agent\n", "from training.rollout import _obs_to_dict, _current_obs_text, build_messages\n", "\n", "\n", "def generate_sft_dataset(env_url, tok, n_episodes, seed_start=2000):\n", " \"\"\"Run baseline_agent for each episode; record (prompt, completion) pairs.\"\"\"\n", " examples = []\n", " with GenericEnvClient(base_url=env_url).sync() as env:\n", " for i in range(n_episodes):\n", " seed = seed_start + i\n", " result = env.reset(seed=seed)\n", " obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n", " task_brief = obs_dict.get('task_brief', '')\n", " turn_history = []\n", " org_config = {}\n", " step, done = 0, False\n", "\n", " while not done and step < 8:\n", " obs_text = _current_obs_text(obs_dict, step, task_brief)\n", " action_type, args = baseline_agent(obs_dict, org_config)\n", "\n", " # Target completion: JSON code block (what we want the model to learn)\n", " payload = {'action_type': action_type, 'args': args}\n", " completion = '```json\\n' + _json.dumps(payload) + '\\n```'\n", "\n", " msgs = build_messages(turn_history, obs_text)\n", " prompt = tok.apply_chat_template(\n", " msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False\n", " )\n", " # Full SFT text = prompt + target completion + eos\n", " examples.append({'text': prompt + completion + tok.eos_token})\n", "\n", " turn_history.append({\n", " 'obs_text' : obs_text,\n", " 'completion': completion,\n", " 'is_runbook': (action_type == 'meta.read_runbook'),\n", " })\n", "\n", " result = env.step({'action_type': action_type, 'args': args})\n", " obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n", "\n", " # Sync org_config from runbook response\n", " if action_type == 'meta.read_runbook':\n", " last = obs_dict.get('last_action_result') or {}\n", " if last.get('ok'):\n", " data = last.get('data') or {}\n", " if isinstance(data, dict) and 'org_config' in data:\n", " org_config.update(data['org_config'])\n", "\n", " done = bool(getattr(result, 'done', obs_dict.get('done', False)))\n", " step += 1\n", "\n", " if (i + 1) % 10 == 0:\n", " print(f' {i+1}/{n_episodes} episodes — {len(examples)} examples')\n", "\n", " return examples\n", "\n", "\n", "print(f'Generating {N_SFT_EPISODES} SFT demonstration episodes...')\n", "sft_raw = generate_sft_dataset(ENV_URL, tokenizer, n_episodes=N_SFT_EPISODES)\n", "sft_dataset = Dataset.from_list(sft_raw)\n", "print(f'\\nSFT dataset : {len(sft_dataset)} examples (~{len(sft_dataset)//6} eps x 6 steps)')\n", "print(f'Sample (first 300 chars):\\n{sft_raw[0][\"text\"][:300]}')" ] }, { "cell_type": "markdown", "id": "cell-10-md", "metadata": {}, "source": [ "## 8. SFT Training (~15 min on A100)" ] }, { "cell_type": "code", "execution_count": null, "id": "cell-10", "metadata": {}, "outputs": [], "source": [ "from trl import SFTTrainer, SFTConfig\n", "\n", "sft_cfg = SFTConfig(\n", " dataset_text_field = 'text',\n", " max_seq_length = 2048,\n", " num_train_epochs = 2,\n", " per_device_train_batch_size = 4,\n", " gradient_accumulation_steps = 4,\n", " learning_rate = 2e-4,\n", " warmup_steps = 10,\n", " output_dir = 'pm-ops-sft-warmup',\n", " report_to = 'none',\n", " logging_steps = 5,\n", " save_strategy = 'no',\n", " dataloader_num_workers = 0,\n", ")\n", "\n", "sft_steps = (\n", " len(sft_dataset)\n", " // (sft_cfg.per_device_train_batch_size * sft_cfg.gradient_accumulation_steps)\n", " * sft_cfg.num_train_epochs\n", ")\n", "print(f'SFT: {len(sft_dataset)} examples x {sft_cfg.num_train_epochs} epochs -> ~{sft_steps} steps')\n", "\n", "sft_trainer = SFTTrainer(\n", " model=model, tokenizer=tokenizer,\n", " train_dataset=sft_dataset, args=sft_cfg,\n", ")\n", "sft_stats = sft_trainer.train()\n", "\n", "runtime = sft_stats.metrics.get('train_runtime', 0)\n", "loss = sft_stats.metrics.get('train_loss', 0)\n", "print(f'SFT done: {round(runtime/60, 1)} min, loss={loss:.3f}')" ] }, { "cell_type": "markdown", "id": "cell-11-md", "metadata": {}, "source": [ "## 9. Verify SFT Output — Model Must Output Valid JSON" ] }, { "cell_type": "code", "execution_count": null, "id": "cell-11", "metadata": {}, "outputs": [], "source": [ "from training.rollout import extract_json_action, _obs_to_dict, _current_obs_text, build_messages\n", "\n", "model.eval()\n", "with GenericEnvClient(base_url=ENV_URL).sync() as _env:\n", " r = _env.reset(seed=99001)\n", " obs_dict = _obs_to_dict(r.observation if hasattr(r, 'observation') else r)\n", " obs_text = _current_obs_text(obs_dict, 0, obs_dict.get('task_brief', ''))\n", " msgs = build_messages([], obs_text)\n", " prompt = tokenizer.apply_chat_template(\n", " msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False\n", " )\n", "\n", "inputs = tokenizer([prompt], return_tensors='pt').to(model.device)\n", "with torch.no_grad():\n", " out = model.generate(\n", " **inputs, max_new_tokens=128, do_sample=False,\n", " pad_token_id=tokenizer.eos_token_id\n", " )\n", "completion = tokenizer.decode(out[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)\n", "parsed = extract_json_action(completion)\n", "\n", "print(f'Output : {completion[:400]}')\n", "print(f'Parsed : {parsed}')\n", "\n", "if parsed is not None:\n", " print('PASS: model outputs valid JSON after SFT')\n", "else:\n", " print('FAIL: still no valid JSON — run SFT cell again with more epochs or more data')\n", "\n", "model.train()" ] }, { "cell_type": "markdown", "id": "cell-12-md", "metadata": {}, "source": [ "---\n", "## Phase 2 — GRPO\n", "\n", "Now that the model outputs valid JSON, GRPO can optimize for *correctness*.\n", "\n", "**Reward** (2 components, sum = 1.0):\n", "\n", "| Component | Weight | Signal |\n", "|---|---|---|\n", "| `env_score` | 0.85 | Env grader: 0.25 (ticket) + 0.20 (label) + 0.20 (priority) + 0.20 (team) + 0.15 (channel) |\n", "| `json_ratio` | 0.15 | Fraction of steps with parseable JSON — maintains format quality |\n", "\n", "`env_score` naturally varies 0.25–1.0 per episode (the model may get the ticket right\n", "but pick the wrong label, or get the channel wrong). This is the learning signal.\n", "Anti-hacking: org_config values differ every episode (seeded), so the model cannot memorize answers." ] }, { "cell_type": "markdown", "id": "cell-13-md", "metadata": {}, "source": [ "## 10. Generate GRPO Training Dataset" ] }, { "cell_type": "code", "execution_count": null, "id": "cell-13", "metadata": {}, "outputs": [], "source": "from training.dataset import generate_triage_dataset\n\n# generate_triage_dataset now pre-simulates the env's RNG to only include seeds\n# where env.reset(seed) will actually run a TRIAGE episode — previously the env\n# silently ran release_notes/dep_update for the same seed, guaranteeing env_score=0.\nrows = generate_triage_dataset(n_episodes=N_GRPO_EPISODES, base_seed=42)\ngrpo_dataset = Dataset.from_list([{'prompt': r['prompt']} for r in rows])\nprint(f'GRPO dataset: {len(grpo_dataset)} triage episodes')\nprint(f'Difficulties: {set(r[\"difficulty\"] for r in rows)}')\nprint(f'Sample prompt: {rows[0][\"prompt\"][:120]}')" }, { "cell_type": "markdown", "id": "cell-14-md", "metadata": {}, "source": [ "## 11. GRPO Rollout + Reward Functions" ] }, { "cell_type": "code", "execution_count": null, "id": "cell-14", "metadata": {}, "outputs": [], "source": "from training.rollout import (\n _obs_to_dict, _current_obs_text, build_messages,\n extract_json_action, step_aware_fallback,\n _generate_no_vllm,\n)\nfrom training.dataset import parse_seed_from_prompt\nfrom training.rewards import compute_rollout_reward\n\ngrpo_env = GenericEnvClient(base_url=ENV_URL).sync()\ngrpo_env.connect()\nprint('GRPO training env connected')\n\nROLLOUT_TEMPERATURE = 1.1 # must be > 1.0 for rollout diversity\n\n\ndef run_grpo_episode(trainer, env, tok, dataset_prompt, max_steps=TRAIN_MAX_STEPS,\n gen_offset=0):\n \"\"\"Run one PM-ops triage episode and return trajectory + reward.\n\n Reward is runbook-compliance based (see training/rewards.py):\n - Did the model read the runbook?\n - Did it use a valid label/priority/team from the runbook?\n - Did it notify a correct oncall channel?\n\n This creates genuine reward variance across rollouts because org configs\n vary by seed — the same hardcoded label/team/channel is correct for some\n orgs and wrong for others, giving GRPO a real gradient signal.\n \"\"\"\n seed = parse_seed_from_prompt(dataset_prompt)\n if seed is not None:\n result = env.reset(seed=seed + gen_offset)\n else:\n result = env.reset()\n\n obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n task_brief = obs_dict.get('task_brief') or dataset_prompt\n\n prompt_ids, completion_ids, logprobs = [], [], []\n turn_history = []\n valid_json_count = 0\n env_score = 0.0\n step, done = 0, False\n _sample_logged = False\n\n # Runbook-compliance tracking\n read_runbook_done = False\n valid_labels: set = set()\n valid_priorities: set = set()\n valid_teams: set = set()\n oncall_channels: set = set()\n ticket_label: str | None = None\n ticket_priority: str | None = None\n assigned_team: str | None = None\n posted_channels: list = []\n\n while not done and step < max_steps:\n obs_text = _current_obs_text(obs_dict, step, task_brief)\n msgs = build_messages(turn_history, obs_text)\n prompt_text = tok.apply_chat_template(\n msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False\n )\n\n rollout_out = _generate_no_vllm(\n trainer, prompt_text, tok,\n max_new_tokens=MAX_COMP_LEN,\n temperature=ROLLOUT_TEMPERATURE,\n )\n prompt_ids.extend(rollout_out['prompt_ids'])\n completion_ids.extend(rollout_out['completion_ids'])\n logprobs.extend(rollout_out['logprobs'])\n completion_text = rollout_out['text']\n\n if not _sample_logged:\n print(f' [sample] {repr(completion_text[:180])}')\n _sample_logged = True\n\n parsed = extract_json_action(completion_text)\n if parsed is not None:\n valid_json_count += 1\n else:\n parsed = step_aware_fallback(step, max_steps)\n\n action_type = parsed.get('action_type', 'meta.noop')\n args = parsed.get('args', {})\n\n if action_type == 'meta.read_runbook' and parsed is not None:\n read_runbook_done = True\n\n if action_type == 'ticketing.create_ticket' and parsed is not None and ticket_label is None:\n ticket_label = args.get('label')\n ticket_priority = args.get('priority')\n\n if action_type == 'ticketing.assign_ticket' and parsed is not None and assigned_team is None:\n assigned_team = args.get('team')\n\n if action_type == 'chat.post_message' and parsed is not None:\n ch = args.get('channel', '')\n if ch:\n posted_channels.append(ch)\n\n turn_history.append({\n 'obs_text' : obs_text,\n 'completion': completion_text,\n 'is_runbook': (action_type == 'meta.read_runbook' and parsed is not None),\n })\n\n try:\n result = env.step({'action_type': action_type, 'args': args})\n except RuntimeError as exc:\n if 'VALIDATION_ERROR' in str(exc):\n result = env.step({'action_type': 'meta.noop', 'args': {}})\n else:\n raise\n\n obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n\n # Extract org_config from runbook response (available the step AFTER read_runbook)\n last_res = obs_dict.get('last_action_result') or {}\n if action_type == 'meta.read_runbook' and last_res.get('ok'):\n data = last_res.get('data') or {}\n if isinstance(data, dict):\n org = data.get('org_config') or {}\n valid_labels = set(org.get('label_taxonomy', {}).values())\n valid_priorities = set(org.get('priority_levels', []))\n valid_teams = set(org.get('team_map', {}).values())\n oncall_channels = set(org.get('oncall_channels', {}).values())\n\n done = bool(getattr(result, 'done', obs_dict.get('done', False)))\n env_score = float(getattr(result, 'reward', obs_dict.get('reward', 0.0)))\n step += 1\n\n reward = compute_rollout_reward(\n read_runbook_done = read_runbook_done,\n valid_labels = valid_labels,\n valid_priorities = valid_priorities,\n valid_teams = valid_teams,\n oncall_channels = oncall_channels,\n ticket_label = ticket_label,\n ticket_priority = ticket_priority,\n assigned_team = assigned_team,\n posted_channels = posted_channels,\n env_score = env_score,\n valid_json_count = valid_json_count,\n )\n\n lbl_ok = '✓' if ticket_label and ticket_label in valid_labels else ('✗' if ticket_label else '-')\n pri_ok = '✓' if ticket_priority and ticket_priority in valid_priorities else ('✗' if ticket_priority else '-')\n tm_ok = '✓' if assigned_team and assigned_team in valid_teams else ('✗' if assigned_team else '-')\n ch_ok = '✓' if any(ch in oncall_channels for ch in posted_channels) else ('✗' if posted_channels else '-')\n print(f' [rollout] steps={step} env={env_score:.3f} '\n f'label={lbl_ok} pri={pri_ok} team={tm_ok} ch={ch_ok} '\n f'offset={gen_offset} -> reward={reward:.3f}')\n\n return {\n 'prompt_ids' : prompt_ids,\n 'completion_ids': completion_ids,\n 'logprobs' : logprobs,\n 'reward' : reward,\n }\n\n\ndef grpo_rollout_func(prompts, trainer=None):\n out = {'prompt_ids': [], 'completion_ids': [], 'logprobs': [], 'reward': []}\n prompt_seen: dict = {}\n for prompt in prompts:\n gen_offset = prompt_seen.get(prompt, 0)\n prompt_seen[prompt] = gen_offset + 1\n ep = run_grpo_episode(trainer, grpo_env, tokenizer, prompt,\n TRAIN_MAX_STEPS, gen_offset=gen_offset)\n for k in out:\n out[k].append(ep[k])\n return out\n\n\n_warned_missing_reward = False\ndef grpo_reward_func(completions, **kwargs):\n \"\"\"Passthrough — reward is pre-computed in grpo_rollout_func.\"\"\"\n global _warned_missing_reward\n rewards = kwargs.get('reward', [])\n if not rewards:\n if not _warned_missing_reward:\n print(f\"[WARN] reward_func fallback, no reward key. kwargs: {list(kwargs.keys())}\")\n _warned_missing_reward = True\n return [0.0] * len(completions)\n return [float(r) for r in rewards]\n\n\nprint(f'GRPO rollout ready max_steps={TRAIN_MAX_STEPS} temperature={ROLLOUT_TEMPERATURE}')" }, { "cell_type": "markdown", "id": "cell-15-md", "metadata": {}, "source": [ "## 12. GRPO Config + Trainer" ] }, { "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, "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", "id": "cell-16-md", "metadata": {}, "source": [ "## 13. Train\n", "\n", "Watch stdout for:\n", "- `[sample] '```json...'` — should look like valid JSON code blocks\n", "- `[rollout] env=X.XXX` — **should trend upward over steps** (this is the signal)\n", "- `[rollout] json=X.XX` — should stay > 0.7 (SFT maintains format quality)\n", "\n", "Watch trackio for the reward curve." ] }, { "cell_type": "code", "execution_count": null, "id": "cell-16", "metadata": {}, "outputs": [], "source": [ "trainer_stats = trainer.train()\n", "\n", "used_gb = round(torch.cuda.max_memory_reserved() / 1024**3, 2)\n", "train_mins = round(trainer_stats.metrics.get('train_runtime', 0) / 60, 1)\n", "print(f'Training time : {train_mins} min')\n", "print(f'Peak GPU : {used_gb} GB / {TOTAL_GB} GB ({round(used_gb/TOTAL_GB*100, 1)}%)')\n", "final_reward = trainer_stats.metrics.get('train/reward', trainer_stats.metrics.get('train_loss', '?'))\n", "print(f'Final reward : {final_reward}')" ] }, { "cell_type": "markdown", "id": "cell-17-md", "metadata": {}, "source": [ "## 14. Save Model" ] }, { "cell_type": "code", "execution_count": null, "id": "cell-17", "metadata": {}, "outputs": [], "source": [ "grpo_env.close()\n", "\n", "# Unsloth merged save — dequantises first, then merges LoRA cleanly into bf16.\n", "# Do NOT use trainer.save_model() directly on a 4-bit + LoRA model.\n", "model.save_pretrained_merged(OUTPUT_DIR, tokenizer, save_method='merged_16bit')\n", "model.push_to_hub_merged(HF_REPO_ID, tokenizer, save_method='merged_16bit')\n", "print(f'Pushed -> https://huggingface.co/{HF_REPO_ID}')" ] }, { "cell_type": "markdown", "id": "cell-18-md", "metadata": {}, "source": [ "## 15. Evaluate: Baseline vs Trained" ] }, { "cell_type": "code", "execution_count": null, "id": "cell-18", "metadata": {}, "outputs": [], "source": [ "from transformers import AutoModelForCausalLM\n", "\n", "N_EVAL = 15\n", "EVAL_MAX_STEPS = 12\n", "EVAL_SEED_BASE = 9000\n", "\n", "eval_model = AutoModelForCausalLM.from_pretrained(\n", " OUTPUT_DIR, torch_dtype=torch.bfloat16, device_map='auto'\n", ")\n", "eval_model.eval()\n", "\n", "\n", "def eval_trained(n=N_EVAL):\n", " scores = []\n", " with GenericEnvClient(base_url=ENV_URL).sync() as env:\n", " for i in range(n):\n", " result = env.reset(seed=EVAL_SEED_BASE + i)\n", " obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n", " task_brief = obs_dict.get('task_brief', '')\n", " history, step, score, done = [], 0, 0.0, False\n", "\n", " while not done and step < EVAL_MAX_STEPS:\n", " obs_text = _current_obs_text(obs_dict, step, task_brief)\n", " msgs = build_messages(history, obs_text)\n", " prompt = tokenizer.apply_chat_template(\n", " msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False\n", " )\n", " inputs = tokenizer([prompt], return_tensors='pt', truncation=True, max_length=4096)\n", " inputs = {k: v.to(eval_model.device) for k, v in inputs.items()}\n", " with torch.no_grad():\n", " out_ids = eval_model.generate(\n", " **inputs, max_new_tokens=256, do_sample=False,\n", " pad_token_id=tokenizer.eos_token_id\n", " )\n", " completion = tokenizer.decode(\n", " out_ids[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True\n", " )\n", " parsed = extract_json_action(completion) or step_aware_fallback(step, EVAL_MAX_STEPS)\n", " result = env.step({'action_type': parsed['action_type'], 'args': parsed.get('args', {})})\n", " obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n", " done = bool(getattr(result, 'done', obs_dict.get('done', False)))\n", " score = float(getattr(result, 'reward', obs_dict.get('reward', 0.0)))\n", " history.append({'obs_text': obs_text, 'completion': completion, 'is_runbook': False})\n", " step += 1\n", "\n", " scores.append(score)\n", " print(f' Trained ep {i+1}/{n}: score={score:.3f}')\n", " return scores\n", "\n", "\n", "def eval_baseline(n=N_EVAL):\n", " from inference import baseline_agent\n", " scores = []\n", " with GenericEnvClient(base_url=ENV_URL).sync() as env:\n", " for i in range(n):\n", " result = env.reset(seed=EVAL_SEED_BASE + i)\n", " obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n", " org_config, step, score, done = {}, 0, 0.0, False\n", " while not done and step < EVAL_MAX_STEPS:\n", " at, args = baseline_agent(obs_dict, org_config)\n", " result = env.step({'action_type': at, 'args': args})\n", " obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n", " done = bool(getattr(result, 'done', obs_dict.get('done', False)))\n", " score = float(getattr(result, 'reward', obs_dict.get('reward', 0.0)))\n", " step += 1\n", " scores.append(score)\n", " print(f' Baseline ep {i+1}/{n}: score={score:.3f}')\n", " return scores\n", "\n", "\n", "print('--- Baseline ---')\n", "baseline_scores = eval_baseline()\n", "print('\\n--- Trained ---')\n", "trained_scores = eval_trained()\n", "\n", "print(f'\\nBaseline avg : {sum(baseline_scores)/N_EVAL:.3f}')\n", "print(f'Trained avg : {sum(trained_scores)/N_EVAL:.3f}')\n", "print(f'Delta : {(sum(trained_scores)-sum(baseline_scores))/N_EVAL:+.3f}')" ] }, { "cell_type": "markdown", "id": "cell-19-md", "metadata": {}, "source": [ "## 16. Plot Results" ] }, { "cell_type": "code", "execution_count": null, "id": "cell-19", "metadata": {}, "outputs": [], "source": [ "import matplotlib.pyplot as plt\n", "import numpy as np\n", "\n", "fig, axes = plt.subplots(1, 2, figsize=(12, 4))\n", "\n", "ax, x, w = axes[0], np.arange(N_EVAL), 0.35\n", "ax.bar(x - w/2, baseline_scores, w, label='Baseline', color='steelblue', alpha=0.8)\n", "ax.bar(x + w/2, trained_scores, w, label='GRPO v3', color='coral', alpha=0.8)\n", "ax.axhline(sum(baseline_scores)/N_EVAL, color='steelblue', ls='--', alpha=0.5, lw=1.5)\n", "ax.axhline(sum(trained_scores)/N_EVAL, color='coral', ls='--', alpha=0.5, lw=1.5)\n", "ax.set(xlabel='Eval episode', ylabel='Reward (0-1)',\n", " title='Per-episode reward: Baseline vs GRPO v3', xticks=x, ylim=(0, 1.05))\n", "ax.legend()\n", "\n", "ax2 = axes[1]\n", "avgs = [sum(baseline_scores)/N_EVAL, sum(trained_scores)/N_EVAL]\n", "bars = ax2.bar(['Baseline', 'GRPO v3'], avgs, color=['steelblue', 'coral'], alpha=0.85, width=0.5)\n", "for bar, val in zip(bars, avgs):\n", " ax2.text(bar.get_x() + bar.get_width()/2, val + 0.01, f'{val:.3f}',\n", " ha='center', fontsize=13, fontweight='bold')\n", "ax2.set(ylabel='Average reward (0-1)', title=f'Average over {N_EVAL} triage episodes',\n", " ylim=(0, 1.05))\n", "\n", "plt.tight_layout()\n", "plt.savefig('eval_results_v3.png', dpi=150, bbox_inches='tight')\n", "plt.show()\n", "print('Saved: eval_results_v3.png')" ] }, { "cell_type": "markdown", "id": "cell-20-md", "metadata": {}, "source": [ "## 17. Teardown" ] }, { "cell_type": "code", "execution_count": null, "id": "cell-20", "metadata": {}, "outputs": [], "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 }