mirror of https://github.com/razor-ai/soup.git
feat(prm): wire PRM reward into GRPO setup via _select_reward_fn (v0.71.30)
This commit is contained in:
parent
ceb54584de
commit
28e306731f
|
|
@ -134,6 +134,24 @@ def _read_attr(obj: Any, name: str) -> Any:
|
|||
return obj.get(name)
|
||||
return getattr(obj, name, None)
|
||||
|
||||
|
||||
def _select_reward_fn(tcfg: Any, device: str, trust_remote_code: bool) -> Any:
|
||||
"""Choose the GRPO reward function (v0.71.30).
|
||||
|
||||
When ``tcfg.prm_reward`` is set, a trained Soup PRM scores each completion's
|
||||
steps and REPLACES the configured ``reward_fn`` (process-supervision). The
|
||||
returned callable rides the existing shaping + ``wrap_reward_funcs`` seam in
|
||||
:meth:`GRPOTrainerWrapper.setup` unchanged, so the v0.71.26 reward-hack
|
||||
mitigation controller observes the PRM reward for free.
|
||||
"""
|
||||
if tcfg.prm_reward is not None:
|
||||
from soup_cli.utils.prm_reward import build_prm_reward_fn
|
||||
|
||||
return build_prm_reward_fn(tcfg, device, trust_remote_code)
|
||||
from soup_cli.trainer.rewards import load_reward_fn
|
||||
|
||||
return load_reward_fn(tcfg.reward_fn, verifiable_domain=tcfg.verifiable_domain)
|
||||
|
||||
class GRPOTrainerWrapper:
|
||||
"""High-level wrapper for GRPO training from SoupConfig.
|
||||
|
||||
|
|
@ -222,11 +240,9 @@ class GRPOTrainerWrapper:
|
|||
use_unsloth = cfg.backend == "unsloth"
|
||||
|
||||
# --- Load reward function ---
|
||||
from soup_cli.trainer.rewards import load_reward_fn
|
||||
|
||||
reward_fn = load_reward_fn(
|
||||
tcfg.reward_fn, verifiable_domain=tcfg.verifiable_domain
|
||||
)
|
||||
# v0.71.30 — when tcfg.prm_reward is set, a trained Soup PRM replaces
|
||||
# the configured reward (process-supervision); otherwise load reward_fn.
|
||||
reward_fn = _select_reward_fn(tcfg, self.device, self._trust_remote_code)
|
||||
|
||||
# v0.71.11 #235/#240 — when the reward-hack or echo-trap detector is
|
||||
# enabled, wrap the reward function(s) with a capture shim so the
|
||||
|
|
|
|||
|
|
@ -349,6 +349,50 @@ class TestBuildPrmRewardFn:
|
|||
mod.build_prm_reward_fn(_T(), device="cpu", trust_remote_code=False)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task 4 — GRPO wiring
|
||||
# ---------------------------------------------------------------------------
|
||||
class TestGrpoPrmWiring:
|
||||
def test_prm_reward_selected(self, monkeypatch):
|
||||
import soup_cli.trainer.grpo as grpo
|
||||
|
||||
sentinel = object()
|
||||
monkeypatch.setattr(
|
||||
"soup_cli.utils.prm_reward.build_prm_reward_fn",
|
||||
lambda tcfg, device, trust_remote_code: sentinel,
|
||||
)
|
||||
|
||||
class _T:
|
||||
prm_reward = "./prm"
|
||||
prm_aggregate = "min"
|
||||
reward_fn = "accuracy"
|
||||
verifiable_domain = None
|
||||
|
||||
out = grpo._select_reward_fn(_T(), "cpu", False)
|
||||
assert out is sentinel
|
||||
|
||||
def test_standard_reward_selected(self, monkeypatch):
|
||||
import soup_cli.trainer.grpo as grpo
|
||||
|
||||
captured = {}
|
||||
|
||||
def _fake_load(spec, verifiable_domain=None):
|
||||
captured["spec"] = spec
|
||||
return "REWARD_FN"
|
||||
|
||||
monkeypatch.setattr("soup_cli.trainer.rewards.load_reward_fn", _fake_load)
|
||||
|
||||
class _T:
|
||||
prm_reward = None
|
||||
prm_aggregate = "min"
|
||||
reward_fn = "accuracy"
|
||||
verifiable_domain = None
|
||||
|
||||
out = grpo._select_reward_fn(_T(), "cpu", False)
|
||||
assert out == "REWARD_FN"
|
||||
assert captured["spec"] == "accuracy"
|
||||
|
||||
|
||||
class TestNoTopLevelTorch:
|
||||
def test_prm_reward_has_no_top_level_torch(self):
|
||||
import soup_cli.utils.prm_reward as mod
|
||||
|
|
|
|||
Loading…
Reference in New Issue