feat(prm): wire PRM reward into GRPO setup via _select_reward_fn (v0.71.30)

This commit is contained in:
Alpamys 2026-07-05 16:56:20 +05:00
parent ceb54584de
commit 28e306731f
2 changed files with 65 additions and 5 deletions

View File

@ -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

View File

@ -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