mirror of https://github.com/razor-ai/soup.git
fix(prm): code-review fixes (PRM head dtype cast, empty-string reject, override notice, prod/render caveats) (v0.71.30)
This commit is contained in:
parent
db2dc7ed65
commit
ec4e78c3ed
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue