fix(v4): reduce MAX_COMP_LEN 384→192, fix eval hanging with for_inference

MAX_COMP_LEN=384 caused clipped_ratio=1.000 every step → all completion
gradients masked → Training Loss=0.000. JSON actions are 80-120 tokens;
192 gives headroom without filling TRL's buffer.

Eval cell hung because Unsloth's training hooks were still active after
trainer.train(). FastLanguageModel.for_inference(model) is required before
model.generate() post-training. Also reduced eval max_new_tokens 256→128.

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

training/train_v4.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.12.0"}
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.0"
+  }
  },
  "cells": [
   {
@@ -33,7 +40,9 @@
    "cell_type": "markdown",
    "id": "v4-s0",
    "metadata": {},
-   "source": ["## 0. Install"]
+   "source": [
+    "## 0. Install"
+   ]
   },
   {
    "cell_type": "code",
@@ -64,7 +73,9 @@
    "cell_type": "markdown",
    "id": "v4-s1",
    "metadata": {},
-   "source": ["## 1. GPU Config + Patch"]
+   "source": [
+    "## 1. GPU Config + Patch"
+   ]
   },
   {
    "cell_type": "code",
@@ -72,37 +83,15 @@
    "metadata": {},
    "execution_count": null,
    "outputs": [],
-   "source": [
-    "import torch\n",
-    "from unsloth import FastLanguageModel, PatchFastRL\n",
-    "\n",
-    "# Must patch BEFORE any trl imports\n",
-    "PatchFastRL('GRPO', FastLanguageModel)\n",
-    "\n",
-    "import trl\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",
-    "\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  = 120 if IS_A100 else 20\n",
-    "N_GRPO_EPISODES = 150 if IS_A100 else 30\n",
-    "TRAIN_MAX_STEPS = 12\n",
-    "\n",
-    "print(f'torch={torch.__version__}  trl={trl.__version__}')\n",
-    "print(f'GPU: {gpu.name} ({TOTAL_GB} GB)  IS_A100={IS_A100}')\n",
-    "print(f'num_gen={NUM_GEN}  grad_accum={GRAD_ACCUM}  '\n",
-    "      f'sft_eps={N_SFT_EPISODES}  grpo_eps={N_GRPO_EPISODES}  max_steps={TRAIN_MAX_STEPS}')"
-   ]
+   "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\n# Per-step max tokens: JSON actions are ~80-120 tokens; 192 gives headroom\n# without filling TRL's completion buffer and zeroing the policy loss.\nMAX_COMP_LEN    = 192\nN_SFT_EPISODES  = 120 if IS_A100 else 20\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"]
+   "source": [
+    "## 2. Clone PM-Ops Repo"
+   ]
   },
   {
    "cell_type": "code",
@@ -135,7 +124,9 @@
    "cell_type": "markdown",
    "id": "v4-s3",
    "metadata": {},
-   "source": ["## 3. HuggingFace Login"]
+   "source": [
+    "## 3. HuggingFace Login"
+   ]
   },
   {
    "cell_type": "code",
@@ -152,7 +143,9 @@
    "cell_type": "markdown",
    "id": "v4-s4",
    "metadata": {},
-   "source": ["## 4. Start PM-Ops Server"]
+   "source": [
+    "## 4. Start PM-Ops Server"
+   ]
   },
   {
    "cell_type": "code",
@@ -190,7 +183,9 @@
    "cell_type": "markdown",
    "id": "v4-s5",
    "metadata": {},
-   "source": ["## 5. Verify Env"]
+   "source": [
+    "## 5. Verify Env"
+   ]
   },
   {
    "cell_type": "code",
@@ -220,7 +215,9 @@
    "cell_type": "markdown",
    "id": "v4-s6",
    "metadata": {},
-   "source": ["## 6. Load Model — Unsloth 4-bit + LoRA"]
+   "source": [
+    "## 6. Load Model — Unsloth 4-bit + LoRA"
+   ]
   },
   {
    "cell_type": "code",
@@ -272,7 +269,9 @@
    "cell_type": "markdown",
    "id": "v4-s7",
    "metadata": {},
-   "source": ["## 7. Generate SFT Dataset"]
+   "source": [
+    "## 7. Generate SFT Dataset"
+   ]
   },
   {
    "cell_type": "code",
@@ -338,7 +337,9 @@
    "cell_type": "markdown",
    "id": "v4-s8",
    "metadata": {},
-   "source": ["## 8. SFT Training"]
+   "source": [
+    "## 8. SFT Training"
+   ]
   },
   {
    "cell_type": "code",
@@ -382,7 +383,9 @@
    "cell_type": "markdown",
    "id": "v4-s9",
    "metadata": {},
-   "source": ["## 9. Verify SFT — Model Must Output Valid JSON"]
+   "source": [
+    "## 9. Verify SFT — Model Must Output Valid JSON"
+   ]
   },
   {
    "cell_type": "code",
@@ -446,7 +449,9 @@
    "cell_type": "markdown",
    "id": "v4-s10",
    "metadata": {},
-   "source": ["## 10. GRPO Training Dataset"]
+   "source": [
+    "## 10. GRPO Training Dataset"
+   ]
   },
   {
    "cell_type": "code",
@@ -468,7 +473,9 @@
    "cell_type": "markdown",
    "id": "v4-s11",
    "metadata": {},
-   "source": ["## 11. GRPO Rollout + Reward"]
+   "source": [
+    "## 11. GRPO Rollout + Reward"
+   ]
   },
   {
    "cell_type": "code",
@@ -663,7 +670,9 @@
    "cell_type": "markdown",
    "id": "v4-s12",
    "metadata": {},
-   "source": ["## 12. GRPO Config + Trainer"]
+   "source": [
+    "## 12. GRPO Config + Trainer"
+   ]
   },
   {
    "cell_type": "code",
@@ -671,51 +680,15 @@
    "metadata": {},
    "execution_count": null,
    "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       = 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')"
-   ]
+   "source": "from trl import GRPOConfig\nfrom training.pm_ops_trainer import PMOpsGRPOTrainer\n\nOUTPUT_DIR = 'pm-ops-grpo-Qwen3-1.7B-triage-v4'\nHF_REPO_ID = f'Saurav1/{OUTPUT_DIR}'\n\ngrpo_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\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 × {NUM_GEN} gen × {grpo_cfg.num_train_epochs} epochs → ~{total_steps} steps')\n\ntrainer = 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)\nassert grpo_cfg.use_vllm is False\nprint(f'Trainer: {type(trainer).__name__}  ready')"
   },
   {
    "cell_type": "markdown",
    "id": "v4-s13",
    "metadata": {},
-   "source": ["## 13. Preflight Probe"]
+   "source": [
+    "## 13. Preflight Probe"
+   ]
   },
   {
    "cell_type": "code",
@@ -767,7 +740,9 @@
    "cell_type": "markdown",
    "id": "v4-s15",
    "metadata": {},
-   "source": ["## 15. Save + Push to HF"]
+   "source": [
+    "## 15. Save + Push to HF"
+   ]
   },
   {
    "cell_type": "code",
@@ -796,7 +771,9 @@
    "cell_type": "markdown",
    "id": "v4-s16",
    "metadata": {},
-   "source": ["## 16. Evaluate"]
+   "source": [
+    "## 16. Evaluate"
+   ]
   },
   {
    "cell_type": "code",
@@ -804,95 +781,15 @@
    "metadata": {},
    "execution_count": null,
    "outputs": [],
-   "source": [
-    "from training.rollout import extract_json_action, step_aware_fallback\n",
-    "from inference import baseline_agent\n",
-    "\n",
-    "N_EVAL         = 15\n",
-    "EVAL_MAX_STEPS = 12\n",
-    "EVAL_SEED_BASE = 9000\n",
-    "\n",
-    "\n",
-    "def run_eval(n=N_EVAL):\n",
-    "    scores = []\n",
-    "    model.eval()\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(model.device) for k, v in inputs.items()}\n",
-    "                with torch.no_grad():\n",
-    "                    out_ids = 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",
-    "        model.train()\n",
-    "    return scores\n",
-    "\n",
-    "\n",
-    "def 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",
-    "            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",
-    "                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",
-    "    return scores\n",
-    "\n",
-    "\n",
-    "print('--- Baseline ---')\n",
-    "baseline_scores = run_baseline()\n",
-    "print('--- Trained ---')\n",
-    "trained_scores  = run_eval()\n",
-    "\n",
-    "b_avg = sum(baseline_scores) / N_EVAL\n",
-    "t_avg = sum(trained_scores)  / N_EVAL\n",
-    "print(f'\\nBaseline avg : {b_avg:.3f}')\n",
-    "print(f'Trained  avg : {t_avg:.3f}')\n",
-    "print(f'Delta        : {t_avg - b_avg:+.3f}')"
-   ]
+   "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# Switch model to Unsloth fast inference mode.\n# Required after training — without this, model.generate() is very slow\n# and may hang on T4 because Unsloth's training hooks are still active.\nFastLanguageModel.for_inference(model)\nmodel.eval()\nprint('Model switched to fast inference mode')\n\n\ndef run_eval(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(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) 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\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            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                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            print(f'  Baseline ep {i+1}/{n}: score={score:.3f}')\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}')"
   },
   {
    "cell_type": "markdown",
    "id": "v4-s17",
    "metadata": {},
-   "source": ["## 17. Teardown"]
+   "source": [
+    "## 17. Teardown"
+   ]
   },
   {
    "cell_type": "code",
@@ -906,4 +803,4 @@
    ]
   }
  ]
-}
+}