feat: add GridMind GRPO training environment and Unsloth training script

scripts/gridmind_grpo_colab.ipynb

@@ -332,24 +332,42 @@
    "outputs": [],
    "source": [
     "import torch\n",
-    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
+    "import gc\n",
+    "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
+    "\n",
+    "# Clear any previous model from memory\n",
+    "for var in ['model', 'trainer']:\n",
+    "    if var in dir():\n",
+    "        del var\n",
+    "gc.collect()\n",
+    "torch.cuda.empty_cache()\n",
     "\n",
     "MODEL_NAME = \"Qwen/Qwen2.5-1.5B-Instruct\"\n",
-    "print(f\"Loading {MODEL_NAME}...\")\n",
+    "print(f\"Loading {MODEL_NAME} with 4-bit quantization for T4 16GB...\")\n",
     "\n",
-    "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n",
+    "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)\n",
     "if tokenizer.pad_token is None:\n",
     "    tokenizer.pad_token = tokenizer.eos_token\n",
+    "tokenizer.padding_side = \"left\"  # required for GRPO\n",
+    "\n",
+    "# 4-bit quantization - fits safely on T4 16GB\n",
+    "bnb_config = BitsAndBytesConfig(\n",
+    "    load_in_4bit=True,\n",
+    "    bnb_4bit_compute_dtype=torch.float16,\n",
+    "    bnb_4bit_quant_type=\"nf4\",\n",
+    "    bnb_4bit_use_double_quant=True,\n",
+    ")\n",
     "\n",
     "model = AutoModelForCausalLM.from_pretrained(\n",
     "    MODEL_NAME,\n",
-    "    torch_dtype=torch.float16,\n",
-    "    device_map=\"cuda\"\n",
+    "    quantization_config=bnb_config,\n",
+    "    device_map=\"auto\",\n",
+    "    trust_remote_code=True,\n",
     ")\n",
     "\n",
-    "total_params = sum(p.numel() for p in model.parameters())\n",
-    "print(f\"Model loaded. Parameters: {total_params/1e6:.0f}M\")\n",
-    "print(f\"Device: {next(model.parameters()).device}\")"
+    "print(f\"Model loaded on: {next(model.parameters()).device}\")\n",
+    "print(f\"Memory allocated: {torch.cuda.memory_allocated()/1e9:.2f} GB / 16 GB\")\n",
+    "print(f\"Memory reserved:  {torch.cuda.memory_reserved()/1e9:.2f} GB / 16 GB\")"
    ]
   },
   {
@@ -368,53 +386,103 @@
    "outputs": [],
    "source": [
     "import json as _json\n",
+    "import requests as _requests\n",
+    "import random as _random\n",
+    "import statistics as _statistics\n",
     "\n",
     "training_rewards = []\n",
-    "\n",
-    "def gridmind_reward_fn(completions, **kwargs):\n",
-    "    \"\"\"Reward function that calls the real environment.\"\"\"\n",
+    "_reward_variance_log = []\n",
+    "_call_count = [0]\n",
+    "\n",
+    "def gridmind_reward_fn(completions, prompts=None, **kwargs):\n",
+    "    \"\"\"\n",
+    "    Reward function compatible with trl 0.23.0.\n",
+    "    Called with positional completions list.\n",
+    "    Must return list of floats same length as completions.\n",
+    "    \"\"\"\n",
     "    rewards = []\n",
-    "    \n",
+    "    batch_raw = []\n",
+    "\n",
     "    for completion in completions:\n",
+    "        _call_count[0] += 1\n",
+    "\n",
     "        try:\n",
-    "            # Extract JSON action from completion\n",
-    "            text = str(completion).strip()\n",
+    "            # Handle both string and list completion formats\n",
+    "            if isinstance(completion, list):\n",
+    "                text = str(completion[0]) if completion else \"\"\n",
+    "            else:\n",
+    "                text = str(completion)\n",
+    "            text = text.strip()\n",
+    "\n",
+    "            # Reset env before each reward call for variance\n",
+    "            task_id = _random.choice([1, 2, 3, 4])\n",
+    "            reset_r = _requests.post(f\"{ENV_URL}/reset\", json={\"task_id\": task_id}, timeout=8)\n",
+    "            if reset_r.status_code != 200:\n",
+    "                rewards.append(-0.5)\n",
+    "                batch_raw.append(-0.5)\n",
+    "                continue\n",
+    "\n",
+    "            # Extract JSON from completion\n",
     "            start = text.rfind('{')\n",
     "            end = text.rfind('}') + 1\n",
     "            if start < 0 or end <= start:\n",
     "                rewards.append(-1.0)\n",
+    "                batch_raw.append(-1.0)\n",
     "                continue\n",
-    "            \n",
-    "            action_str = text[start:end]\n",
-    "            action = _json.loads(action_str)\n",
-    "            \n",
-    "            # Clamp action to valid ranges\n",
-    "            action[\"hvac_power_level\"] = max(0.0, min(1.0, float(action.get(\"hvac_power_level\", 0.5))))\n",
-    "            action[\"thermal_charge_rate\"] = max(-1.0, min(1.0, float(action.get(\"thermal_charge_rate\", 0.0))))\n",
-    "            action[\"batch_job_slot\"] = max(0, min(4, int(action.get(\"batch_job_slot\", 0))))\n",
-    "            action[\"load_shed_fraction\"] = max(0.0, min(0.5, float(action.get(\"load_shed_fraction\", 0.0))))\n",
-    "            action[\"building_id\"] = int(action.get(\"building_id\", 0))\n",
-    "            \n",
-    "            # Call environment\n",
-    "            r = requests.post(f\"{ENV_URL}/step\", json=action, timeout=8)\n",
-    "            if r.status_code != 200:\n",
+    "\n",
+    "            action = _json.loads(text[start:end])\n",
+    "            action = {\n",
+    "                \"hvac_power_level\": max(0.0, min(1.0, float(action.get(\"hvac_power_level\", 0.5)))),\n",
+    "                \"thermal_charge_rate\": max(-1.0, min(1.0, float(action.get(\"thermal_charge_rate\", 0.0)))),\n",
+    "                \"batch_job_slot\": max(0, min(4, int(action.get(\"batch_job_slot\", 0)))),\n",
+    "                \"load_shed_fraction\": max(0.0, min(0.5, float(action.get(\"load_shed_fraction\", 0.0)))),\n",
+    "                \"building_id\": int(action.get(\"building_id\", 0)),\n",
+    "            }\n",
+    "\n",
+    "            step_r = _requests.post(f\"{ENV_URL}/step\", json=action, timeout=8)\n",
+    "            if step_r.status_code != 200:\n",
     "                rewards.append(-0.5)\n",
+    "                batch_raw.append(-0.5)\n",
     "                continue\n",
-    "            \n",
-    "            step_data = r.json()\n",
-    "            if isinstance(step_data, list):\n",
-    "                step_data = step_data[0]\n",
-    "            \n",
-    "            reward = float(step_data.get(\"reward\", 0))\n",
-    "            rewards.append(max(-1.0, min(1.0, reward)))  # Clamp to [-1, 1]\n",
-    "            training_rewards.append(reward)\n",
-    "            \n",
-    "        except Exception as e:\n",
-    "            rewards.append(-1.0)\n",
-    "    \n",
+    "\n",
+    "            data = step_r.json()\n",
+    "            if isinstance(data, list):\n",
+    "                data = data[0]\n",
+    "\n",
+    "            base = float(data.get(\"reward\", 0.0))\n",
+    "            comps = data.get(\"rewards\", {})\n",
+    "            bonus = (\n",
+    "                float(comps.get(\"cost_savings\", 0)) * 0.3 +\n",
+    "                float(comps.get(\"task_satisfaction\", 0)) * 0.2 +\n",
+    "                float(comps.get(\"efficiency_bonus\", 0)) * 0.1 +\n",
+    "                float(comps.get(\"temperature_constraint\", 0)) * 0.15\n",
+    "            )\n",
+    "            final = max(-1.0, min(1.0, base + bonus))\n",
+    "            rewards.append(final)\n",
+    "            batch_raw.append(final)\n",
+    "            training_rewards.append(final)\n",
+    "\n",
+    "        except _json.JSONDecodeError:\n",
+    "            rewards.append(-0.8)\n",
+    "            batch_raw.append(-0.8)\n",
+    "        except Exception:\n",
+    "            rewards.append(-0.5)\n",
+    "            batch_raw.append(-0.5)\n",
+    "\n",
+    "    # Log variance every 10 calls\n",
+    "    if len(batch_raw) > 1 and _call_count[0] % 10 == 0:\n",
+    "        try:\n",
+    "            var = _statistics.variance(batch_raw)\n",
+    "            _reward_variance_log.append(var)\n",
+    "            print(f\"  [Call {_call_count[0]}] Rewards: {[f'{r:.3f}' for r in batch_raw]} | Variance: {var:.4f}\")\n",
+    "            if var < 0.001:\n",
+    "                print(\"    Zero variance - no learning signal!\")\n",
+    "        except Exception:\n",
+    "            pass\n",
+    "\n",
     "    return rewards\n",
     "\n",
-    "print(\"Reward function defined.\")"
+    "print(\"Reward function defined (trl 0.23.0 compatible)\")"
    ]
   },
   {
@@ -433,49 +501,133 @@
    "outputs": [],
    "source": [
     "from trl import GRPOTrainer, GRPOConfig\n",
+    "from peft import LoraConfig, prepare_model_for_kbit_training\n",
     "from datasets import Dataset\n",
+    "import inspect\n",
+    "import os\n",
+    "import requests as _requests\n",
+    "import statistics\n",
+    "import torch, gc\n",
     "\n",
     "# Prepare dataset\n",
     "train_data = [{\"prompt\": d[\"prompt\"]} for d in dataset]\n",
     "train_ds = Dataset.from_list(train_data)\n",
-    "\n",
     "print(f\"Training dataset: {len(train_ds)} prompts\")\n",
-    "print(f\"Sample prompt:\\n{train_data[0]['prompt'][:200]}...\\n\")\n",
     "\n",
-    "# GRPO config for free T4 GPU\n",
-    "config = GRPOConfig(\n",
+    "theme_dist = {}\n",
+    "for d in dataset:\n",
+    "    t = d.get(\"theme\", \"unknown\")\n",
+    "    theme_dist[t] = theme_dist.get(t, 0) + 1\n",
+    "print(f\"Theme distribution: {theme_dist}\")\n",
+    "print(f\"Sample prompt preview:\\n{train_data[0]['prompt'][:200]}...\\n\")\n",
+    "\n",
+    "# Prepare model for QLoRA training\n",
+    "model.config.use_cache = False\n",
+    "model.gradient_checkpointing_enable()\n",
+    "model = prepare_model_for_kbit_training(model)\n",
+    "\n",
+    "peft_config = LoraConfig(\n",
+    "    r=16,\n",
+    "    lora_alpha=32,\n",
+    "    target_modules=[\"q_proj\", \"v_proj\", \"k_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
+    "    lora_dropout=0.05,\n",
+    "    bias=\"none\",\n",
+    "    task_type=\"CAUSAL_LM\",\n",
+    ")\n",
+    "\n",
+    "# GRPOConfig - trl==0.23.0 compatible. Pass this as args=, not config=.\n",
+    "# generation_kwargs is not a GRPOTrainer init parameter in trl 0.23.0.\n",
+    "grpo_config = GRPOConfig(\n",
     "    output_dir=\"./gridmind-grpo-output\",\n",
     "    num_train_epochs=1,\n",
-    "    max_steps=60,  # Complete in ~30-40 min on T4\n",
-    "    per_device_train_batch_size=2,\n",
-    "    gradient_accumulation_steps=2,\n",
-    "    max_prompt_length=512,\n",
-    "    learning_rate=5e-6,\n",
-    "    logging_steps=5,\n",
-    "    save_steps=60,\n",
+    "    max_steps=60,\n",
+    "    per_device_train_batch_size=1,\n",
+    "    gradient_accumulation_steps=4,\n",
+    "    max_prompt_length=400,\n",
+    "    max_completion_length=80,\n",
+    "    num_generations=4,\n",
+    "    learning_rate=5e-5,\n",
     "    fp16=True,\n",
-    "    dataloader_num_workers=0,\n",
+    "    logging_steps=1,\n",
+    "    save_steps=60,\n",
     "    report_to=\"none\",\n",
-    "    num_generations=2,  # 2 generations per prompt for speed\n",
+    "    dataloader_num_workers=0,\n",
+    "    remove_unused_columns=False,\n",
     ")\n",
     "\n",
-    "print(\"\\nStarting GRPO training...\")\n",
-    "print(f\"Estimated time: 30-40 minutes on Colab T4 GPU\")\n",
-    "print(f\"Steps: {config.max_steps}, Batch size: {config.per_device_train_batch_size * config.gradient_accumulation_steps}\\n\")\n",
-    "\n",
-    "# Initialize trainer\n",
+    "print(\"=== PRE-TRAINING DIAGNOSTIC ===\\n\")\n",
+    "import trl\n",
+    "print(f\"TRL version: {trl.__version__}\")\n",
+    "sig = inspect.signature(GRPOTrainer.__init__)\n",
+    "params = list(sig.parameters.keys())\n",
+    "print(f\"GRPOTrainer params: {params[:8]}\")\n",
+    "print(f\"Uses 'args=':   {'args' in params}\")\n",
+    "print(f\"Uses 'config=': {'config' in params}\")\n",
+    "\n",
+    "print(\"\\nTesting reward function...\")\n",
+    "test_completions = [\n",
+    "    '{\"hvac_power_level\": 0.2, \"thermal_charge_rate\": 0.8, \"batch_job_slot\": 2, \"load_shed_fraction\": 0.0, \"building_id\": 0}',\n",
+    "    '{\"hvac_power_level\": 1.0, \"thermal_charge_rate\": -1.0, \"batch_job_slot\": 0, \"load_shed_fraction\": 0.5, \"building_id\": 0}',\n",
+    "    '{\"hvac_power_level\": 0.5, \"thermal_charge_rate\": 0.0, \"batch_job_slot\": 0, \"load_shed_fraction\": 0.0, \"building_id\": 0}',\n",
+    "    'not valid json at all',\n",
+    "]\n",
+    "test_rewards = gridmind_reward_fn(test_completions)\n",
+    "print(f\"Test rewards: {[f'{r:.3f}' for r in test_rewards]}\")\n",
+    "reward_var = statistics.variance(test_rewards) if len(set(test_rewards)) > 1 else 0.0\n",
+    "if reward_var <= 0.001:\n",
+    "    print(\"CRITICAL: Reward variance is too low - fix reward function before training\")\n",
+    "else:\n",
+    "    print(f\"Reward variance: {reward_var:.4f} - sufficient for GRPO\")\n",
+    "\n",
+    "print(f\"\\nGPU memory: {torch.cuda.memory_allocated()/1e9:.2f} GB used / 16 GB total\")\n",
+    "print(f\"Free: {(16 - torch.cuda.memory_allocated()/1e9):.2f} GB\")\n",
+    "print(\"\\n=== READY TO TRAIN ===\" if reward_var > 0.001 else \"\\n=== FIX REWARD FUNCTION FIRST ===\")\n",
+    "\n",
+    "# Reset environment before training\n",
+    "_requests.post(f\"{ENV_URL}/reset\", json={\"task_id\": 1}, timeout=10)\n",
+    "print(\"Environment reset before training.\")\n",
+    "\n",
+    "# Initialize GRPOTrainer - trl 0.23.0 API\n",
     "trainer = GRPOTrainer(\n",
     "    model=model,\n",
+    "    args=grpo_config,\n",
     "    processing_class=tokenizer,\n",
-    "    config=config,\n",
     "    train_dataset=train_ds,\n",
     "    reward_funcs=gridmind_reward_fn,\n",
-    "    generation_kwargs={\"max_new_tokens\": 100},\n",
+    "    peft_config=peft_config,\n",
     ")\n",
     "\n",
-    "# Train\n",
-    "trainer.train()\n",
-    "print(\"\\n\u00e2\u0153\u201c Training complete!\")"
+    "print(\"\\nStarting GRPO training with QLoRA...\")\n",
+    "print(f\"Steps: {grpo_config.max_steps} | Batch: {grpo_config.per_device_train_batch_size} | Generations: {grpo_config.num_generations}\")\n",
+    "print(\"Estimated time: ~25-35 min on T4\\n\")\n",
+    "\n",
+    "train_result = trainer.train()\n",
+    "\n",
+    "print(\"\\nTraining complete!\")\n",
+    "print(f\"  Total steps:    {train_result.global_step}\")\n",
+    "print(f\"  Training loss:  {train_result.training_loss:.6f}\")\n",
+    "\n",
+    "if train_result.training_loss == 0.0:\n",
+    "    print(\"\\nWARNING: Loss is 0.0 - reward function may have zero variance.\")\n",
+    "    print(\"Check reward diagnostic output above. This means the model saw no learning signal.\")\n",
+    "else:\n",
+    "    print(\"\\nNon-zero loss confirmed - model received learning signal.\")\n",
+    "\n",
+    "print(f\"\\nMemory after training: {torch.cuda.memory_allocated()/1e9:.2f} GB\")\n",
+    "\n",
+    "# Save LoRA adapter (much smaller than full model)\n",
+    "adapter_path = \"./gridmind-lora-adapter\"\n",
+    "trainer.model.save_pretrained(adapter_path)\n",
+    "tokenizer.save_pretrained(adapter_path)\n",
+    "print(f\"LoRA adapter saved to {adapter_path}\")\n",
+    "\n",
+    "total_size = sum(\n",
+    "    os.path.getsize(os.path.join(adapter_path, f))\n",
+    "    for f in os.listdir(adapter_path)\n",
+    "    if os.path.isfile(os.path.join(adapter_path, f))\n",
+    ")\n",
+    "print(f\"Adapter size: {total_size/1e6:.1f} MB\")\n",
+    "print(\"Full model would be ~3 GB - adapter is the diff only\")"
    ]
   },
   {
@@ -598,7 +750,7 @@
     "    },\n",
     "    \"improvement_percent\": overall_improvement,\n",
     "    \"model\": MODEL_NAME,\n",
-    "    \"training_steps\": config.max_steps,\n",
+    "    \"training_steps\": grpo_config.max_steps,\n",
     "    \"themes_covered\": [\"multi_agent\", \"instruction_following\", \"world_modeling\", \"curriculum\"],\n",
     "    \"training_rewards_log\": training_rewards[-20:] if training_rewards else [],\n",
     "}\n",
@@ -624,4 +776,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}

scripts/train_unsloth.py

@@ -690,7 +690,7 @@ def main():
     
     trainer = GRPOTrainer(
         model=model,
-        tokenizer=tokenizer,
+        processing_class=tokenizer,
         args=training_args,
         train_dataset=dataset,
         reward_funcs=[
@@ -746,4 +746,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()