diff --git a/src/soup_cli/config/schema.py b/src/soup_cli/config/schema.py index 52d1bad..07c77e9 100644 --- a/src/soup_cli/config/schema.py +++ b/src/soup_cli/config/schema.py @@ -1055,6 +1055,42 @@ class TrainingConfig(BaseModel): "Required when reward_fn='verifiable'." ), ) + # v0.71.30 — PRM-guided GRPO: use a trained Soup PRM as the per-step + # reward inside GRPO. ``prm_reward`` names the PRM directory (or HF id); + # ``prm_aggregate`` folds the per-step scalars into one reward. + prm_reward: Optional[str] = Field( + default=None, + description=( + "Path (or HF id) to a Soup-trained PRM (task='prm') used as the " + "GRPO per-step reward (v0.71.30). When set, the PRM replaces " + "reward_fn. Requires task='grpo', backend='transformers', " + "modality='text'." + ), + ) + prm_aggregate: Literal["min", "prod", "last"] = Field( + 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." + ), + ) + + @field_validator("prm_reward", mode="before") + @classmethod + def _validate_prm_reward_field(cls, value): + """v0.71.30 — shape-only validation (containment enforced at load).""" + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, str): + raise ValueError( + f"prm_reward must be a string path/id, got {type(value).__name__}" + ) + if "\x00" in value: + raise ValueError("prm_reward must not contain null bytes") + if len(value) > 512: + raise ValueError("prm_reward must be <= 512 chars") + return value + # v0.50.0 Part A — GRPO objective variants (unsloth + axolotl parity). # Schema-only in v0.50.0; live loss kernels wired in v0.50.1. grpo_variant: Optional[Literal[ @@ -4213,6 +4249,38 @@ class SoupConfig(BaseModel): ) return self + @model_validator(mode="after") + def _validate_prm_reward(self) -> "SoupConfig": + """v0.71.30 — PRM-guided GRPO gate. + + ``prm_reward`` runs a PRM (base CausalLM + reward head) forward as the + GRPO reward, so it requires ``task='grpo'`` on ``backend='transformers'`` + with ``modality='text'``. A non-default ``prm_aggregate`` while + ``prm_reward`` is unset silently no-ops → reject as a footgun. + """ + if self.training.prm_reward is None: + if self.training.prm_aggregate != "min": + raise ValueError( + "prm_aggregate is only meaningful with prm_reward set; " + f"got prm_aggregate={self.training.prm_aggregate!r} and " + "prm_reward=None" + ) + return self + if self.task != "grpo": + raise ValueError( + f"prm_reward requires task='grpo'; got task={self.task!r}" + ) + if self.backend != "transformers": + raise ValueError( + "prm_reward requires backend='transformers' (the PRM reward " + f"runs a transformers forward); got backend={self.backend!r}" + ) + if self.modality != "text": + raise ValueError( + f"prm_reward requires modality='text'; got modality={self.modality!r}" + ) + return self + @model_validator(mode="after") def _validate_vllm_sleep_mode(self) -> "SoupConfig": """v0.50.0 Part B — ``vllm_sleep_mode`` requires task='grpo' and a diff --git a/tests/test_v07130.py b/tests/test_v07130.py index b3de5b7..1984448 100644 --- a/tests/test_v07130.py +++ b/tests/test_v07130.py @@ -87,6 +87,103 @@ class TestAggregate: assert set(AGGREGATE_MODES) == {"min", "prod", "last"} +# --------------------------------------------------------------------------- +# Task 2 — schema fields + cross-validators +# --------------------------------------------------------------------------- +def _prm_yaml( + *, + task: str = "grpo", + backend: str = "transformers", + modality: str = "text", + prm_reward: str | None = "./prm", + prm_aggregate: str | None = None, +) -> str: + lines = [ + "base: HuggingFaceTB/SmolLM2-135M", + f"task: {task}", + f"backend: {backend}", + f"modality: {modality}", + "data:", + " train: ./data/train.jsonl", + " format: chatml", + "training:", + ] + if prm_reward is not None: + lines.append(f" prm_reward: {prm_reward}") + if prm_aggregate is not None: + lines.append(f" prm_aggregate: {prm_aggregate}") + return "\n".join(lines) + "\n" + + +class TestPrmSchema: + def test_default_fields(self): + from soup_cli.config.schema import TrainingConfig + + tc = TrainingConfig() + assert tc.prm_reward is None + assert tc.prm_aggregate == "min" + + def test_happy_grpo_parses(self): + from soup_cli.config.loader import load_config_from_string + + cfg = load_config_from_string(_prm_yaml(prm_aggregate="prod")) + assert cfg.training.prm_reward == "./prm" + assert cfg.training.prm_aggregate == "prod" + + def test_rejects_non_grpo_task(self): + from soup_cli.config.loader import load_config_from_string + + with pytest.raises(ValueError, match="task='grpo'"): + load_config_from_string(_prm_yaml(task="sft")) + + def test_rejects_mlx_backend(self): + from soup_cli.config.loader import load_config_from_string + + with pytest.raises(ValueError, match="transformers"): + load_config_from_string(_prm_yaml(backend="mlx")) + + def test_rejects_unsloth_backend(self): + from soup_cli.config.loader import load_config_from_string + + with pytest.raises(ValueError, match="transformers"): + load_config_from_string(_prm_yaml(backend="unsloth")) + + def test_rejects_non_text_modality(self): + from soup_cli.config.loader import load_config_from_string + + with pytest.raises(ValueError, match="modality='text'"): + load_config_from_string(_prm_yaml(modality="vision")) + + def test_aggregate_without_prm_reward_is_footgun(self): + from soup_cli.config.loader import load_config_from_string + + with pytest.raises(ValueError, match="prm_reward"): + load_config_from_string( + _prm_yaml(prm_reward=None, prm_aggregate="prod") + ) + + def test_default_aggregate_without_prm_reward_ok(self): + from soup_cli.config.loader import load_config_from_string + + # prm_aggregate at its default is fine even without prm_reward. + cfg = load_config_from_string( + _prm_yaml(task="sft", prm_reward=None, prm_aggregate="min") + ) + assert cfg.training.prm_reward is None + + def test_rejects_null_byte(self): + from soup_cli.config.schema import TrainingConfig + + with pytest.raises(ValueError, match="null"): + TrainingConfig(prm_reward="./prm\x00evil") + + def test_rejects_oversize(self): + from soup_cli.config.schema import TrainingConfig + + with pytest.raises(ValueError): + TrainingConfig(prm_reward="x" * 5000) + + class TestNoTopLevelTorch: def test_prm_reward_has_no_top_level_torch(self): import soup_cli.utils.prm_reward as mod