{ "cells": [ { "cell_type": "markdown", "id": "v4-title", "metadata": {}, "source": [ "# PM-Ops GRPO Training v4\n", "\n", "Clean rewrite — all known v3 bugs fixed.\n", "\n", "| Fix | Details |\n", "|---|---|\n", "| Constant reward=0.350 | Replaced env_score formula with runbook-compliance scoring |\n", "| `gen_slot=0` override bug | gen_slot now correctly offsets env seed per GRPO generation |\n", "| Double `env.step` per step | Removed duplicate step call inside try/except |\n", "| Dataset/env task-type mismatch | `dataset.py` pre-filters seeds to triage-only episodes |\n", "| Near-greedy generation | Temperature 1.1 + top_k=50 for diverse rollouts |\n", "| PatchFastRL called twice | Called once, unconditionally, before any trl imports |\n", "\n", "**Reward**: read_runbook (+0.10) + valid_label (±0.20/0.10) + valid_priority (±0.15/0.10) + valid_team (±0.20/0.10) + right_channel (±0.25/0.10) + env_bonus (×0.10). Varies per org config — different valid values per seed. \n", "**Stack**: Unsloth 2026.x · TRL 0.22.2 · Qwen3-1.7B · PMOpsGRPOTrainer \n", "**GPU**: A100-40GB → ~2 hrs (SFT 15 min + GRPO 90 min)" ] }, { "cell_type": "markdown", "id": "v4-s0", "metadata": {}, "source": [ "## 0. Install" ] }, { "cell_type": "code", "execution_count": null, "id": "v4-install", "metadata": {}, "outputs": [], "source": [ "%%capture\n", "import os\n", "!pip install --upgrade -qqq uv\n", "if 'COLAB_' not in ''.join(os.environ.keys()):\n", " !uv pip install unsloth vllm\n", "else:\n", " import subprocess\n", " is_t4 = 'Tesla T4' in str(subprocess.check_output(['nvidia-smi']))\n", " _vllm = 'vllm==0.9.2' if is_t4 else 'vllm==0.15.1'\n", " _triton = 'triton==3.2.0' if is_t4 else 'triton'\n", " !uv pip install -qqq --upgrade {_vllm} torchvision bitsandbytes xformers unsloth\n", " !uv pip install -qqq {_triton}\n", "!uv pip install transformers==4.56.2\n", "!uv pip install --no-deps trl==0.22.2\n", "!pip install 'numpy==1.26.4' --break-system-packages -q\n", "!pip install openenv openenv-core -q" ] }, { "cell_type": "markdown", "id": "v4-s1", "metadata": {}, "source": [ "## 1. GPU Config + Patch" ] }, { "cell_type": "code", "execution_count": null, "id": "v4-config", "metadata": {}, "outputs": [], "source": "import torch\nfrom unsloth import FastLanguageModel, PatchFastRL\n\n# Must patch BEFORE any trl imports\nPatchFastRL('GRPO', FastLanguageModel)\n\nimport trl\n\ngpu = torch.cuda.get_device_properties(0)\nTOTAL_GB = round(gpu.total_memory / 1024**3, 1)\nIS_A100 = TOTAL_GB >= 35\n\nNUM_GEN = 6 if IS_A100 else 2\nGRAD_ACCUM = 32 if IS_A100 else 8\nMAX_COMP_LEN = 192\n# T4 needs more SFT to reach low enough loss before GRPO.\n# 20 eps → loss≈2.0 is too high; 50 eps → loss≈0.8 is a better floor.\nN_SFT_EPISODES = 120 if IS_A100 else 50\nN_GRPO_EPISODES = 150 if IS_A100 else 30\nTRAIN_MAX_STEPS = 12\n\nprint(f'torch={torch.__version__} trl={trl.__version__}')\nprint(f'GPU: {gpu.name} ({TOTAL_GB} GB) IS_A100={IS_A100}')\nprint(f'num_gen={NUM_GEN} grad_accum={GRAD_ACCUM} '\n f'sft_eps={N_SFT_EPISODES} grpo_eps={N_GRPO_EPISODES} '\n f'max_steps={TRAIN_MAX_STEPS} max_comp_len={MAX_COMP_LEN}')" }, { "cell_type": "markdown", "id": "v4-s2", "metadata": {}, "source": [ "## 2. Clone PM-Ops Repo" ] }, { "cell_type": "code", "execution_count": null, "id": "v4-clone", "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()}')\n", "!git -C {REPO_DIR} log --oneline -3" ] }, { "cell_type": "markdown", "id": "v4-s3", "metadata": {}, "source": [ "## 3. HuggingFace Login" ] }, { "cell_type": "code", "execution_count": null, "id": "v4-hf-login", "metadata": {}, "outputs": [], "source": [ "from huggingface_hub import notebook_login\n", "notebook_login()" ] }, { "cell_type": "markdown", "id": "v4-s4", "metadata": {}, "source": [ "## 4. Start PM-Ops Server" ] }, { "cell_type": "code", "execution_count": null, "id": "v4-server", "metadata": {}, "outputs": [], "source": [ "import subprocess, time, requests\n", "\n", "# Kill any leftover server from a previous run\n", "subprocess.run(['pkill', '-f', 'uvicorn'], capture_output=True)\n", "time.sleep(1)\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": "v4-s5", "metadata": {}, "source": [ "## 5. Verify Env" ] }, { "cell_type": "code", "execution_count": null, "id": "v4-verify-env", "metadata": {}, "outputs": [], "source": [ "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(seed=42)\n", " obs = _obs_to_dict(r.observation if hasattr(r, 'observation') else r)\n", " print(f'task_brief : {obs.get(\"task_brief\", \"?\")[:80]}...')\n", " step_r = _env.step({'action_type': 'meta.read_runbook', 'args': {}})\n", " rb_obs = _obs_to_dict(step_r.observation if hasattr(step_r, 'observation') else step_r)\n", " data = (rb_obs.get('last_action_result') or {}).get('data', {})\n", " org = data.get('org_config', {}) if isinstance(data, dict) else {}\n", " print(f'org labels : {list(org.get(\"label_taxonomy\", {}).values())}')\n", " print(f'org priorities: {org.get(\"priority_levels\", [])}')\n", " print(f'org channels : {list(org.get(\"oncall_channels\", {}).values())}')\n", " print('Env verify: OK')" ] }, { "cell_type": "markdown", "id": "v4-s6", "metadata": {}, "source": [ "## 6. Load Model — Unsloth 4-bit + LoRA" ] }, { "cell_type": "code", "execution_count": null, "id": "v4-model", "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,\n", " max_lora_rank = LORA_RANK,\n", " gpu_memory_utilization = 0.55,\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", "print(f'GPU after load: {round(torch.cuda.max_memory_reserved()/1024**3, 2)} GB / {TOTAL_GB} GB')" ] }, { "cell_type": "markdown", "id": "v4-phase1", "metadata": {}, "source": [ "---\n", "## Phase 1 — SFT Warmup\n", "\n", "Teach the model the JSON output format and PM-ops workflow before GRPO. \n", "`baseline_agent` runs N deterministic episodes → ~6 steps each → supervised (prompt, completion) pairs. \n", "2 SFT epochs (~15 min on A100). Without SFT, all rollouts output freeform text → zero GRPO gradient." ] }, { "cell_type": "markdown", "id": "v4-s7", "metadata": {}, "source": [ "## 7. Generate SFT Dataset" ] }, { "cell_type": "code", "execution_count": null, "id": "v4-sft-data", "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", " examples = []\n", " with GenericEnvClient(base_url=env_url).sync() as env:\n", " for i in range(n_episodes):\n", " result = env.reset(seed=seed_start + i)\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, org_config, 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", " payload = {'action_type': action_type, 'args': args}\n", " completion = '```json\\n' + _json.dumps(payload) + '\\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", " examples.append({'text': prompt + completion + tok.eos_token})\n", " turn_history.append({\n", " 'obs_text': obs_text, 'completion': completion,\n", " 'is_runbook': (action_type == 'meta.read_runbook'),\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", " 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", " done = bool(getattr(result, 'done', obs_dict.get('done', False)))\n", " step += 1\n", "\n", " if (i + 1) % 20 == 0:\n", " print(f' {i+1}/{n_episodes} eps — {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'SFT dataset: {len(sft_dataset)} examples (~{len(sft_dataset)//6} eps × 6 steps)')\n", "print(f'Sample (first 300 chars):\\n{sft_raw[0][\"text\"][:300]}')" ] }, { "cell_type": "markdown", "id": "v4-s8", "metadata": {}, "source": [ "## 8. SFT Training" ] }, { "cell_type": "code", "execution_count": null, "id": "v4-sft-train", "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", "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 × {sft_cfg.num_train_epochs} epochs → ~{sft_steps} steps')\n", "\n", "sft_trainer = SFTTrainer(model=model, tokenizer=tokenizer,\n", " train_dataset=sft_dataset, args=sft_cfg)\n", "sft_stats = sft_trainer.train()\n", "loss = sft_stats.metrics.get('train_loss', 0)\n", "runtime = sft_stats.metrics.get('train_runtime', 0)\n", "print(f'SFT done: {round(runtime/60, 1)} min loss={loss:.3f}')" ] }, { "cell_type": "markdown", "id": "v4-s9", "metadata": {}, "source": [ "## 9. Verify SFT — Model Must Output Valid JSON" ] }, { "cell_type": "code", "execution_count": null, "id": "v4-sft-verify", "metadata": {}, "outputs": [], "source": [ "from training.rollout import extract_json_action\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(**inputs, max_new_tokens=128, do_sample=False,\n", " pad_token_id=tokenizer.eos_token_id)\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", "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 again with more episodes or epochs.')\n", "model.train()" ] }, { "cell_type": "markdown", "id": "v4-phase2", "metadata": {}, "source": [ "---\n", "## Phase 2 — GRPO\n", "\n", "Reward components — sum ≈ 0.90 when all correct:\n", "\n", "| Signal | Weight | Description |\n", "|---|---|---|\n", "| `read_runbook` | +0.10 | Did agent read the runbook first? |\n", "| `valid_label` | +0.20 / −0.10 | Ticket label in org's `label_taxonomy`? |\n", "| `valid_priority` | +0.15 / −0.10 | Ticket priority in org's `priority_levels`? |\n", "| `valid_team` | +0.20 / −0.10 | Assigned team in org's `team_map`? |\n", "| `right_channel` | +0.25 / −0.10/ch | Posted to org's `oncall_channels`? |\n", "| `env_bonus` | ×0.10 | Env grader confirmation |\n", "\n", "Reward **varies per org config** (different valid values per seed) → genuine GRPO advantage signal. \n", "Each of the N GRPO generations gets a different env seed via `gen_slot` offset." ] }, { "cell_type": "markdown", "id": "v4-s10", "metadata": {}, "source": [ "## 10. GRPO Training Dataset" ] }, { "cell_type": "code", "execution_count": null, "id": "v4-grpo-data", "metadata": {}, "outputs": [], "source": [ "from training.dataset import generate_triage_dataset\n", "\n", "rows = generate_triage_dataset(n_episodes=N_GRPO_EPISODES, base_seed=42)\n", "grpo_dataset = Dataset.from_list([{'prompt': r['prompt']} for r in rows])\n", "print(f'GRPO dataset: {len(grpo_dataset)} triage episodes')\n", "print(f'Difficulties: {set(r[\"difficulty\"] for r in rows)}')\n", "print(f'Sample: {rows[0][\"prompt\"][:120]}')" ] }, { "cell_type": "markdown", "id": "v4-s11", "metadata": {}, "source": [ "## 11. GRPO Rollout + Reward" ] }, { "cell_type": "code", "execution_count": null, "id": "v4-rollout", "metadata": {}, "outputs": [], "source": [ "import torch.nn.functional as F\n", "from training.rollout import (\n", " _obs_to_dict, _current_obs_text, build_messages,\n", " extract_json_action, step_aware_fallback,\n", ")\n", "from training.dataset import parse_seed_from_prompt\n", "from training.rewards import compute_rollout_reward\n", "\n", "grpo_env = GenericEnvClient(base_url=ENV_URL).sync()\n", "grpo_env.connect()\n", "print('GRPO env connected')\n", "\n", "\n", "def run_grpo_episode(trainer, env, tok, dataset_prompt, max_steps=TRAIN_MAX_STEPS, gen_slot=0):\n", " \"\"\"One full PM-ops episode. gen_slot offsets seed so parallel GRPO generations\n", " explore distinct env episodes even for the same base prompt.\"\"\"\n", " seed = parse_seed_from_prompt(dataset_prompt)\n", " result = env.reset(seed=seed + gen_slot) if seed is not None else 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", "\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", " _model = (trainer.accelerator.unwrap_model(trainer.model)\n", " if hasattr(trainer, 'accelerator') else trainer.model)\n", " _device = (trainer.accelerator.device\n", " if hasattr(trainer, 'accelerator') else next(_model.parameters()).device)\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", " enc = tok(prompt_text, return_tensors='pt', truncation=True, max_length=4096).to(_device)\n", " plen = enc['input_ids'].shape[1]\n", " with torch.no_grad():\n", " out = _model.generate(\n", " **enc,\n", " max_new_tokens = MAX_COMP_LEN,\n", " do_sample = True,\n", " temperature = 1.1,\n", " top_p = 0.95,\n", " top_k = 50,\n", " pad_token_id = tok.pad_token_id or tok.eos_token_id,\n", " output_scores = True,\n", " return_dict_in_generate = True,\n", " )\n", "\n", " cids = out.sequences[0][plen:].tolist()\n", " completion_text = tok.decode(cids, skip_special_tokens=True)\n", "\n", " prompt_ids.extend(enc['input_ids'][0].tolist())\n", " completion_ids.extend(cids)\n", " logprobs.extend([\n", " F.log_softmax(s[0], dim=-1)[t].item()\n", " for s, t in zip(out.scores, cids)\n", " ])\n", "\n", " if step == 0:\n", " print(f' [sample] {repr(completion_text[:180])}')\n", "\n", " parsed = extract_json_action(completion_text)\n", " is_valid = parsed is not None\n", " if not is_valid:\n", " parsed = step_aware_fallback(step, max_steps)\n", " else:\n", " valid_json_count += 1\n", "\n", " action_type = parsed.get('action_type', 'meta.noop')\n", " args = parsed.get('args', {})\n", " print(f' [step {step}] {action_type}')\n", "\n", " # Compliance signal capture\n", " if action_type == 'meta.read_runbook' and is_valid:\n", " read_runbook_done = True\n", " if action_type == 'ticketing.create_ticket' and is_valid and ticket_label is None:\n", " ticket_label = args.get('label')\n", " ticket_priority = args.get('priority')\n", " if action_type == 'ticketing.assign_ticket' and is_valid and assigned_team is None:\n", " assigned_team = args.get('team')\n", " if action_type == 'chat.post_message' and is_valid:\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 is_valid),\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", " # Extract org config from runbook response (one step after the call)\n", " last_result = obs_dict.get('last_action_result') or {}\n", " if action_type == 'meta.read_runbook' and last_result.get('ok'):\n", " data = last_result.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 = '✓' if ticket_label and ticket_label in valid_labels else ('✗' if ticket_label else '-')\n", " pri = '✓' if ticket_priority and ticket_priority in valid_priorities else ('✗' if ticket_priority else '-')\n", " tm = '✓' if assigned_team and assigned_team in valid_teams else ('✗' if assigned_team else '-')\n", " ch = '✓' if any(c in oncall_channels for c in posted_channels) else ('✗' if posted_channels else '-')\n", " print(f' [rollout] steps={step} env={env_score:.3f} '\n", " f'label={lbl} priority={pri} team={tm} channel={ch} → 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", "\n", "def 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(\n", " trainer, grpo_env, tokenizer, prompt, TRAIN_MAX_STEPS, gen_slot=gen_offset\n", " )\n", " for k in out:\n", " out[k].append(ep[k])\n", " return out\n", "\n", "\n", "def grpo_reward_func(completions, **kwargs):\n", " \"\"\"Passthrough — reward pre-computed in grpo_rollout_func.\"\"\"\n", " rewards = kwargs.get('reward', [])\n", " if not rewards:\n", " print(f'[grpo_reward_func] no reward in kwargs — keys: {list(kwargs.keys())}')\n", " return [0.0] * len(completions)\n", " return [float(r) for r in rewards]\n", "\n", "\n", "print(f'GRPO rollout ready max_steps={TRAIN_MAX_STEPS}')" ] }, { "cell_type": "markdown", "id": "v4-s12", "metadata": {}, "source": [ "## 12. GRPO Config + Trainer" ] }, { "cell_type": "code", "execution_count": null, "id": "v4-trainer", "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-v4'\n", "HF_REPO_ID = f'Saurav1/{OUTPUT_DIR}'\n", "\n", "grpo_cfg = GRPOConfig(\n", " num_train_epochs = 2,\n", " learning_rate = 1e-6,\n", " gradient_accumulation_steps = GRAD_ACCUM,\n", " per_device_train_batch_size = 1,\n", " warmup_steps = 5,\n", " num_generations = NUM_GEN,\n", " # max_completion_length must match MAX_COMP_LEN so TRL's internal\n", " # single-turn generation doesn't overflow and zero the policy loss\n", " # (clipped_ratio=1.0 means all gradient is masked → loss=0.000).\n", " max_completion_length = MAX_COMP_LEN,\n", " max_prompt_length = 4096,\n", " use_vllm = False,\n", " output_dir = OUTPUT_DIR,\n", " report_to = 'none',\n", " logging_steps = 1,\n", " save_steps = 20,\n", " gradient_checkpointing = False,\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 × {NUM_GEN} gen × {grpo_cfg.num_train_epochs} epochs → ~{total_steps} steps')\n", "\n", "trainer = PMOpsGRPOTrainer(\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", "assert grpo_cfg.use_vllm is False\n", "print(f'Trainer: {type(trainer).__name__} ready')" ] }, { "cell_type": "markdown", "id": "v4-s13", "metadata": {}, "source": [ "## 13. Preflight Probe" ] }, { "cell_type": "code", "execution_count": null, "id": "v4-probe", "metadata": {}, "outputs": [], "source": [ "# Run 2 episodes and verify reward varies (not stuck at a constant)\n", "probe = grpo_rollout_func(\n", " [grpo_dataset[0]['prompt'], grpo_dataset[1]['prompt']],\n", " trainer=trainer,\n", ")\n", "print(f'probe rewards : {probe[\"reward\"]}')\n", "assert len(probe['reward']) == 2, 'need one reward per prompt'\n", "assert all(isinstance(r, float) for r in probe['reward']), 'rewards must be float'\n", "print('Preflight PASS ✓')" ] }, { "cell_type": "markdown", "id": "v4-s14", "metadata": {}, "source": [ "## 14. Train\n", "\n", "Watch for:\n", "- `[sample] '\\`\\`\\`json ...'` — model should output JSON code blocks\n", "- `label=✓ priority=✓ team=✓ channel=✓` — compliance signals the model is getting right\n", "- `reward/injected_std > 0` in logs — confirms GRPO has a non-zero gradient signal" ] }, { "cell_type": "code", "execution_count": null, "id": "v4-train", "metadata": {}, "outputs": [], "source": [ "trainer_stats = trainer.train()\n", "\n", "train_mins = round(trainer_stats.metrics.get('train_runtime', 0) / 60, 1)\n", "used_gb = round(torch.cuda.max_memory_reserved() / 1024**3, 2)\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)}%)')" ] }, { "cell_type": "markdown", "id": "v4-s15", "metadata": {}, "source": [ "## 15. Save + Push to HF" ] }, { "cell_type": "code", "execution_count": null, "id": "v4-save", "metadata": {}, "outputs": [], "source": [ "grpo_env.close()\n", "\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": "v4-eval-hdr", "metadata": {}, "source": [ "---\n", "## Evaluation — Baseline vs Trained" ] }, { "cell_type": "markdown", "id": "v4-s16", "metadata": {}, "source": [ "## 16. Evaluate" ] }, { "cell_type": "code", "execution_count": null, "id": "v4-eval", "metadata": {}, "outputs": [], "source": "from training.rollout import extract_json_action, step_aware_fallback\nfrom inference import baseline_agent\n\nN_EVAL = 15\nEVAL_MAX_STEPS = 12\nEVAL_SEED_BASE = 9000\n\n# Required after training: switch Unsloth from training hooks to fast inference.\n# Without this, model.generate() is very slow and hangs on T4.\nFastLanguageModel.for_inference(model)\nmodel.eval()\nprint('Model switched to fast inference mode')\n\n\ndef _model_step(obs_dict, task_brief, history, step):\n \"\"\"One greedy decode step. Returns (action_type, args, json_ok).\"\"\"\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(model.device) for k, v in inputs.items()}\n with torch.no_grad():\n out_ids = model.generate(\n **inputs,\n max_new_tokens = 128,\n 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)\n json_ok = parsed is not None\n if not json_ok:\n parsed = step_aware_fallback(step, EVAL_MAX_STEPS)\n history.append({'obs_text': obs_text, 'completion': completion, 'is_runbook': False})\n return parsed.get('action_type', 'meta.noop'), parsed.get('args', {}), json_ok, obs_text\n\n\ndef run_eval(n=N_EVAL, verbose=False):\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 actions = []\n\n while not done and step < EVAL_MAX_STEPS:\n action_type, args, json_ok, _ = _model_step(obs_dict, task_brief, history, step)\n actions.append(f'{action_type}{\"\" if json_ok else \"!\"}')\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 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\n scores.append(score)\n action_summary = ' → '.join(a.split('.')[-1] for a in actions)\n print(f' Trained ep {i+1}/{n}: score={score:.3f} [{action_summary}]')\n return scores\n\n\ndef run_baseline(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 org_config, step, score, done = {}, 0, 0.0, False\n actions = []\n while not done and step < EVAL_MAX_STEPS:\n at, args = baseline_agent(obs_dict, org_config)\n actions.append(at.split('.')[-1])\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 if at == '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 step += 1\n scores.append(score)\n action_summary = ' → '.join(actions)\n print(f' Baseline ep {i+1}/{n}: score={score:.3f} [{action_summary}]')\n return scores\n\n\nprint('--- Baseline ---')\nbaseline_scores = run_baseline()\nprint('\\n--- Trained ---')\ntrained_scores = run_eval()\n\nb_avg = sum(baseline_scores) / N_EVAL\nt_avg = sum(trained_scores) / N_EVAL\nprint(f'\\nBaseline avg : {b_avg:.3f}')\nprint(f'Trained avg : {t_avg:.3f}')\nprint(f'Delta : {t_avg - b_avg:+.3f}')\nprint()\nprint('Action key: ! = fallback (no valid JSON from model)')" }, { "cell_type": "markdown", "id": "v4-s17", "metadata": {}, "source": [ "## 17. Teardown" ] }, { "cell_type": "code", "execution_count": null, "id": "v4-teardown", "metadata": {}, "outputs": [], "source": [ "server_proc.terminate()\n", "print('PM-Ops server stopped')" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.12.0" } }, "nbformat": 4, "nbformat_minor": 5 }