fix(probe): use TRL 0.24.0 log keys — rewards/commerce_reward_fn/mean, grad_norm (not train/ prefix)

notebooks/v4_2_instruct_grpo.ipynb

@@ -631,7 +631,90 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "FastLanguageModel.for_training(model)\n\nprobe_config = GRPOConfig(\n    output_dir=str(CHECKPOINT_DIR / \"probe\"),\n    num_generations=NUM_GENERATIONS,\n    scale_rewards=SCALE_REWARDS,\n    max_completion_length=MAX_COMPLETION_LENGTH,\n    max_steps=10,\n    temperature=TEMPERATURE,\n    beta=BETA,\n    num_train_epochs=1,\n    per_device_train_batch_size=BATCH_SIZE,\n    gradient_accumulation_steps=GRAD_ACCUM,\n    learning_rate=LEARNING_RATE,\n    lr_scheduler_type=LR_SCHEDULER_TYPE,\n    warmup_ratio=WARMUP_RATIO,\n    fp16=False,\n    bf16=True,\n    logging_steps=1,\n    save_steps=999,\n    report_to=\"none\",\n    max_prompt_length=MAX_SEQ_LENGTH // 2,\n    seed=CURRENT_SEED,\n    remove_unused_columns=False,\n)\n\nprobe_trainer = UnslothGRPOTrainer(\n    model=model,\n    reward_funcs=commerce_reward_fn,\n    args=probe_config,\n    train_dataset=train_dataset,\n    processing_class=tokenizer,\n)\n\nt0 = time.time()\nresult = probe_trainer.train()\nelapsed = time.time() - t0\n\n# ── Extract metrics from log history ─────────────────────────────────────────\nrewards = []\ngrad_norms = []\nzero_stds = []\nfor entry in probe_trainer.state.log_history:\n    if \"train/reward\" in entry:\n        rewards.append(entry[\"train/reward\"])\n    if \"train/grad_norm\" in entry:\n        grad_norms.append(entry[\"train/grad_norm\"])\n    if \"train/frac_reward_zero_std\" in entry:\n        zero_stds.append(entry[\"train/frac_reward_zero_std\"])\n\nprint(f\"\\n{'='*60}\")\nprint(f\"PROBE RESULTS ({elapsed:.0f}s, {elapsed/10:.0f}s/step)\")\nprint(f\"  Rewards:    {[f'{r:.3f}' for r in rewards]}\")\nprint(f\"  Grad norms: {[f'{g:.4f}' for g in grad_norms]}\")\nprint(f\"  Zero-std:   {[f'{z:.2f}' for z in zero_stds]}\")\nprint(f\"  Train loss: {result.training_loss:.4f}\")\nprint(f\"{'='*60}\")\n\nif rewards and max(rewards) > 0:\n    print(\"✓ Model generates scoreable output\")\nelse:\n    print(\"✗ WARNING: All rewards are 0. Check reward functions.\")\n\nif grad_norms and max(grad_norms) > 0:\n    print(\"✓ Gradients are flowing\")\nelse:\n    print(\"✗ WARNING: All grad_norms are 0. Check model/LoRA setup.\")\n\nif zero_stds and max(zero_stds) < 0.5:\n    print(\"✓ Batches have reward variance (GRPO has signal)\")\nelse:\n    print(\"⚠️ WARNING: High frac_reward_zero_std. Consider increasing G.\")\n\nprint(\"\\n→ Proceed to full training (Cell 13)\")\n\ndel probe_trainer\ngc.collect(); torch.cuda.empty_cache()"
+   "source": [
+    "FastLanguageModel.for_training(model)\n",
+    "\n",
+    "probe_config = GRPOConfig(\n",
+    "    output_dir=str(CHECKPOINT_DIR / \"probe\"),\n",
+    "    num_generations=NUM_GENERATIONS,\n",
+    "    scale_rewards=SCALE_REWARDS,\n",
+    "    max_completion_length=MAX_COMPLETION_LENGTH,\n",
+    "    max_steps=10,\n",
+    "    temperature=TEMPERATURE,\n",
+    "    beta=BETA,\n",
+    "    num_train_epochs=1,\n",
+    "    per_device_train_batch_size=BATCH_SIZE,\n",
+    "    gradient_accumulation_steps=GRAD_ACCUM,\n",
+    "    learning_rate=LEARNING_RATE,\n",
+    "    lr_scheduler_type=LR_SCHEDULER_TYPE,\n",
+    "    warmup_ratio=WARMUP_RATIO,\n",
+    "    fp16=False,\n",
+    "    bf16=True,\n",
+    "    logging_steps=1,\n",
+    "    save_steps=999,\n",
+    "    report_to=\"none\",\n",
+    "    max_prompt_length=MAX_SEQ_LENGTH // 2,\n",
+    "    seed=CURRENT_SEED,\n",
+    "    remove_unused_columns=False,\n",
+    ")\n",
+    "\n",
+    "probe_trainer = UnslothGRPOTrainer(\n",
+    "    model=model,\n",
+    "    reward_funcs=commerce_reward_fn,\n",
+    "    args=probe_config,\n",
+    "    train_dataset=train_dataset,\n",
+    "    processing_class=tokenizer,\n",
+    ")\n",
+    "\n",
+    "t0 = time.time()\n",
+    "result = probe_trainer.train()\n",
+    "elapsed = time.time() - t0\n",
+    "\n",
+    "# ── Extract metrics from log history ─────────────────────────────────────────\n",
+    "# V4.2.1: TRL 0.24.0 logs under \"reward\" / \"rewards/commerce_reward_fn/mean\"\n",
+    "# and \"grad_norm\" (no \"train/\" prefix in log_history entries).\n",
+    "rewards = []\n",
+    "reward_stds = []\n",
+    "grad_norms = []\n",
+    "for entry in probe_trainer.state.log_history:\n",
+    "    if \"rewards/commerce_reward_fn/mean\" in entry:\n",
+    "        rewards.append(entry[\"rewards/commerce_reward_fn/mean\"])\n",
+    "    if \"rewards/commerce_reward_fn/std\" in entry:\n",
+    "        reward_stds.append(entry[\"rewards/commerce_reward_fn/std\"])\n",
+    "    if \"grad_norm\" in entry:\n",
+    "        grad_norms.append(entry[\"grad_norm\"])\n",
+    "\n",
+    "print(f\"\\n{'='*60}\")\n",
+    "print(f\"PROBE RESULTS ({elapsed:.0f}s, {elapsed/10:.0f}s/step)\")\n",
+    "print(f\"  Rewards:     {[f'{r:.3f}' for r in rewards]}\")\n",
+    "print(f\"  Reward stds: {[f'{s:.3f}' for s in reward_stds]}\")\n",
+    "print(f\"  Grad norms:  {[f'{g:.4f}' for g in grad_norms]}\")\n",
+    "print(f\"  Train loss:  {result.training_loss:.4f}\")\n",
+    "print(f\"{'='*60}\")\n",
+    "\n",
+    "if rewards and max(rewards) > 0:\n",
+    "    print(\"✓ Model generates scoreable output\")\n",
+    "else:\n",
+    "    print(\"✗ WARNING: All rewards are 0. Check reward functions.\")\n",
+    "\n",
+    "if grad_norms and max(grad_norms) > 0:\n",
+    "    print(\"✓ Gradients are flowing\")\n",
+    "else:\n",
+    "    print(\"✗ WARNING: All grad_norms are 0. Check model/LoRA setup.\")\n",
+    "\n",
+    "if reward_stds and min(reward_stds) > 0:\n",
+    "    print(\"✓ Batches have reward variance (GRPO has signal)\")\n",
+    "elif reward_stds:\n",
+    "    n_zero = sum(1 for s in reward_stds if s < 1e-6)\n",
+    "    print(f\"⚠️ WARNING: {n_zero}/{len(reward_stds)} steps had zero reward std. Consider increasing G.\")\n",
+    "else:\n",
+    "    print(\"⚠️ WARNING: No reward_std logged. Check TRL version.\")\n",
+    "\n",
+    "print(\"\\n→ Proceed to full training (Cell 13)\")\n",
+    "\n",
+    "del probe_trainer\n",
+    "gc.collect(); torch.cuda.empty_cache()"
+   ]
   },
   {
    "cell_type": "markdown",