From ec4e78c3edfec89b1b37b12fdc88fbd792bef255 Mon Sep 17 00:00:00 2001
From: Alpamys <vpn.alpamys@gmail.com>
Date: Sun, 5 Jul 2026 17:30:56 +0500
Subject: [PATCH] fix(prm): code-review fixes (PRM head dtype cast,
 empty-string reject, override notice, prod/render caveats) (v0.71.30)

---
 src/soup_cli/config/schema.py     |  9 ++++++++-
 src/soup_cli/envs/retrieval_qa.py |  2 +-
 src/soup_cli/trainer/prm.py       |  6 +++++-
 src/soup_cli/utils/prm_reward.py  | 20 ++++++++++++++++++--
 tests/test_v07130.py              |  6 ++++++
 5 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/src/soup_cli/config/schema.py b/src/soup_cli/config/schema.py
index 122030f..d54a6e3 100644
--- a/src/soup_cli/config/schema.py
+++ b/src/soup_cli/config/schema.py
@@ -1071,7 +1071,12 @@ class TrainingConfig(BaseModel):
         default="min",
         description=(
             "How PRM per-step scores fold into one reward: min (weakest-link, "
-            "default) | prod | last. Only meaningful when prm_reward is set."
+            "default) | prod | last. Only meaningful when prm_reward is set. "
+            "NOTE: 'prod' assumes per-step scores are bounded in ~[0,1] (a "
+            "probability of step-correctness). Soup's PRM head is trained with "
+            "unconstrained MSE regression, so 'prod' can blow up / flip sign on "
+            "unbounded labels — prefer the default 'min' unless your PRM labels "
+            "are calibrated to [0,1]."
         ),
     )
 
@@ -1085,6 +1090,8 @@ class TrainingConfig(BaseModel):
             raise ValueError(
                 f"prm_reward must be a string path/id, got {type(value).__name__}"
             )
+        if not value:
+            raise ValueError("prm_reward must not be an empty string")
         if "\x00" in value:
             raise ValueError("prm_reward must not contain null bytes")
         if len(value) > 512:
diff --git a/src/soup_cli/envs/retrieval_qa.py b/src/soup_cli/envs/retrieval_qa.py
index ab24167..ce1f684 100644
--- a/src/soup_cli/envs/retrieval_qa.py
+++ b/src/soup_cli/envs/retrieval_qa.py
@@ -38,7 +38,7 @@ def _make_row(rng: random.Random) -> dict[str, str]:
     doc = " ".join(f"{e} {a} {v}." for (e, a, v) in facts)
     prompt = (
         f"Document: {doc}\n"
-        f"Question: What does {entity} {attribute.rstrip()}? "
+        f"Question: What value completes '{entity} {attribute.rstrip()} ___'? "
         "Reply with just the value."
     )
     return {"prompt": prompt, "answer": value}
diff --git a/src/soup_cli/trainer/prm.py b/src/soup_cli/trainer/prm.py
index f00990f..9fc32fa 100644
--- a/src/soup_cli/trainer/prm.py
+++ b/src/soup_cli/trainer/prm.py
@@ -220,7 +220,11 @@ class PRMTrainerWrapper:
             torch_dtype=torch.bfloat16 if self.device == "cuda" else torch.float32,
         )
         hidden_size = base_model.config.hidden_size
-        base_model.reward_head = nn.Linear(hidden_size, 1, bias=True)
+        # Cast the reward head to the base model's dtype. On CUDA the base loads
+        # in bf16 while nn.Linear defaults to fp32, so without this cast the
+        # first compute_loss forward (hidden_states[bf16] @ reward_head[fp32])
+        # raises a dtype-mismatch RuntimeError (v0.71.30 code-review fix).
+        base_model.reward_head = nn.Linear(hidden_size, 1, bias=True).to(base_model.dtype)
         self.model = base_model
         self._dataset = dataset
         console.print(
diff --git a/src/soup_cli/utils/prm_reward.py b/src/soup_cli/utils/prm_reward.py
index f3a2e4f..9b607b9 100644
--- a/src/soup_cli/utils/prm_reward.py
+++ b/src/soup_cli/utils/prm_reward.py
@@ -16,8 +16,16 @@ Design:
   controller observes the PRM reward for free.
 
 Honesty: proof-of-mechanism only — a tiny PRM signal is noisy; this is NOT a
-production reward-model claim (see #286). The step split is a newline
-heuristic (v1).
+production reward-model claim (see #286). Known v1 caveats:
+- The step split is a newline heuristic.
+- ``prm_aggregate='prod'`` assumes per-step scores in ~[0,1]; the PRM head is
+  trained with unconstrained MSE, so 'prod' can blow up on uncalibrated labels
+  — the default 'min' (weakest-link) is the safe choice.
+- The prompt context is rendered by joining message contents (matching the
+  PRM's plaintext training field), NOT via the tokenizer chat template the
+  policy model sees — a distributional gap, acceptable for a proof of mechanism.
+- Completions are scored one forward pass each (no batching) — fine for tiny
+  models; a batched path is a future optimisation.
 
 Security:
 - ``prm_reward`` local paths are containment-checked (realpath + commonpath
@@ -315,6 +323,14 @@ def build_prm_reward_fn(
         prm_path = real
 
     resolved_trust = _resolve_trust(prm_path, trust_remote_code, console)
+    # Announce that the PRM reward is active AND replaces the configured
+    # reward_fn — otherwise a user who also set reward_fn/verifiable_domain has
+    # no signal those are being ignored (code-review MEDIUM/LOW).
+    console.print(
+        f"[dim]Using PRM reward: prm_reward={prm_path!r}, "
+        f"aggregate={tcfg.prm_aggregate!r} "
+        "(this replaces reward_fn/verifiable_domain).[/]"
+    )
     return PRMScorer(
         prm_path=prm_path,
         aggregate=tcfg.prm_aggregate,
diff --git a/tests/test_v07130.py b/tests/test_v07130.py
index d581db9..446c3fd 100644
--- a/tests/test_v07130.py
+++ b/tests/test_v07130.py
@@ -171,6 +171,12 @@ class TestPrmSchema:
         )
         assert cfg.training.prm_reward is None
 
+    def test_rejects_empty_string(self):
+        from soup_cli.config.schema import TrainingConfig
+
+        with pytest.raises(ValueError, match="empty"):
+            TrainingConfig(prm_reward="")
+
     def test_rejects_null_byte(self):
         from soup_cli.config.schema import TrainingConfig