feat(prm): schema prm_reward/prm_aggregate + cross-validators (v0.71.30)

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

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