diff --git a/src/soup_cli/config/schema.py b/src/soup_cli/config/schema.py index 4045026..175c8da 100644 --- a/src/soup_cli/config/schema.py +++ b/src/soup_cli/config/schema.py @@ -1741,6 +1741,66 @@ class TrainingConfig(BaseModel): "vote. Allowlist: info_rm, rm_ensemble, length_trend, repetition." ), ) + # ---- v0.71.26 Stage 2 — PID-Lagrangian controller + rollback --------- + reward_hack_pid_kp: float = Field( + default=0.5, + ge=0.0, + le=1000.0, + description="v0.71.26 — PID proportional gain (pid_lagrangian mode).", + ) + reward_hack_pid_ki: float = Field( + default=0.1, + ge=0.0, + le=1000.0, + description="v0.71.26 — PID integral gain (pid_lagrangian mode).", + ) + reward_hack_pid_kd: float = Field( + default=0.05, + ge=0.0, + le=1000.0, + description="v0.71.26 — PID derivative gain (pid_lagrangian mode).", + ) + reward_hack_signal_target: float = Field( + default=0.15, + ge=0.0, + lt=1.0, + description=( + "v0.71.26 — target hacking drop_pct the PID controller holds " + "(pid_lagrangian mode)." + ), + ) + reward_hack_rollback: bool = Field( + default=False, + description=( + "v0.71.26 — enable rollback to the last-good RL checkpoint in the " + "escalation ladder. Requires rl_checkpoint_save_every_steps set." + ), + ) + reward_hack_rollback_patience: int = Field( + default=3, + ge=1, + le=100_000, + description=( + "v0.71.26 — consecutive HACK steps before a rollback is triggered." + ), + ) + reward_hack_max_recovery_attempts: int = Field( + default=2, + ge=0, + le=1000, + description=( + "v0.71.26 — max rollbacks before the controller early-stops " + "training (terminal rung of the escalation ladder)." + ), + ) + + @field_validator("reward_hack_rollback", mode="before") + @classmethod + def _validate_reward_hack_rollback(cls, v): + """v0.71.26 — bool guard so YAML ``yes`` / ``1`` cannot silently coerce.""" + if v is None or isinstance(v, bool): + return v + raise TypeError(f"reward_hack_rollback must be bool, got {type(v).__name__}") @field_validator("reward_hack_mitigation", mode="before") @classmethod @@ -3148,6 +3208,18 @@ class EvalConfig(BaseModel): # any to a non-default value while reward_hack_mitigation='off' is a silent # no-op footgun (mirrors the v0.70.0 minillm offenders-list policy). Extended # per stage (Stage 2/3 tunables added with their fields). +# Stage-2 (PID-Lagrangian + rollback) tunables — meaningful only in +# pid_lagrangian mode. Setting one under any other mode is a no-op footgun. +_REWARD_HACK_STAGE2_DEFAULTS: dict = { + "reward_hack_pid_kp": 0.5, + "reward_hack_pid_ki": 0.1, + "reward_hack_pid_kd": 0.05, + "reward_hack_signal_target": 0.15, + "reward_hack_rollback": False, + "reward_hack_rollback_patience": 3, + "reward_hack_max_recovery_attempts": 2, +} + _REWARD_HACK_TUNABLE_DEFAULTS: dict = { "reward_hack_beta_floor": 0.02, "reward_hack_beta_ceil": 1.0, @@ -3157,6 +3229,7 @@ _REWARD_HACK_TUNABLE_DEFAULTS: dict = { "reward_hack_release_patience": 3, "reward_hack_kl_gain": 1.5, "reward_hack_signals": ["info_rm"], + **_REWARD_HACK_STAGE2_DEFAULTS, } @@ -3206,6 +3279,24 @@ def _validate_reward_hack_controller(tcfg) -> None: "exclusive with ref_model_ema_alpha (both drive the KL/ref " "dynamics); pick one" ) + # v0.71.26 Stage 2 — PID / rollback tunables require pid_lagrangian mode. + if tcfg.reward_hack_mitigation != "pid_lagrangian": + stage2_offenders = [ + name + for name, default in _REWARD_HACK_STAGE2_DEFAULTS.items() + if getattr(tcfg, name, default) != default + ] + if stage2_offenders: + raise ValueError( + f"PID/rollback tunables {stage2_offenders} require " + "reward_hack_mitigation='pid_lagrangian'" + ) + # Rollback needs an RL-checkpoint cadence to roll back to. + if tcfg.reward_hack_rollback and tcfg.rl_checkpoint_save_every_steps is None: + raise ValueError( + "reward_hack_rollback=True requires rl_checkpoint_save_every_steps " + "to be set (a cadence to roll back to)" + ) class SoupConfig(BaseModel): diff --git a/src/soup_cli/utils/peft_wiring.py b/src/soup_cli/utils/peft_wiring.py index 0497f17..c9915f9 100644 --- a/src/soup_cli/utils/peft_wiring.py +++ b/src/soup_cli/utils/peft_wiring.py @@ -219,13 +219,16 @@ def _attach_reward_hack( tokenizer: Any, output_dir: str, task: str, + rl_checkpoint_cb: Any = None, ) -> int: """Attach the reward-hack callback: mitigation controller (v0.71.26) when a ``reward_hack_mitigation`` mode is set, else the plain v0.70.0 detector. Returns 1 when a callback was attached, 0 otherwise. The mitigation controller SUBSUMES the plain detector (they share the same signal), so - exactly one of the two is ever attached. + exactly one of the two is ever attached. ``rl_checkpoint_cb`` is the + (already-built) RL-checkpoint callback the pid_lagrangian rollback ladder + restores from. """ import os @@ -235,6 +238,7 @@ def _attach_reward_hack( from soup_cli.utils.reward_hack_control import ( BangBangPolicy, MitigationLogWriter, + PIDLagrangianPolicy, RewardHackMitigationCallback, ) @@ -246,6 +250,7 @@ def _attach_reward_hack( getattr(tcfg, "reward_hack_signals", None) or ("info_rm",) ) bang_bang = None + pid = None if mitigation == "kl_control": bang_bang = BangBangPolicy( beta_floor=tcfg.reward_hack_beta_floor, @@ -256,6 +261,16 @@ def _attach_reward_hack( release_patience=tcfg.reward_hack_release_patience, kl_gain=tcfg.reward_hack_kl_gain, ) + elif mitigation == "pid_lagrangian": + pid = PIDLagrangianPolicy( + kp=tcfg.reward_hack_pid_kp, + ki=tcfg.reward_hack_pid_ki, + kd=tcfg.reward_hack_pid_kd, + signal_target=tcfg.reward_hack_signal_target, + beta_floor=tcfg.reward_hack_beta_floor, + beta_ceil=tcfg.reward_hack_beta_ceil, + integral_clamp=tcfg.reward_hack_beta_ceil, + ) callback = RewardHackMitigationCallback( mode=mitigation, detector=detector, @@ -265,6 +280,15 @@ def _attach_reward_hack( tokenizer=tokenizer, task=task, bang_bang=bang_bang, + pid=pid, + rollback=bool(getattr(tcfg, "reward_hack_rollback", False)), + rollback_patience=int( + getattr(tcfg, "reward_hack_rollback_patience", 3) + ), + max_recovery_attempts=int( + getattr(tcfg, "reward_hack_max_recovery_attempts", 2) + ), + rl_checkpoint_cb=rl_checkpoint_cb, ) trainer.add_callback(callback) callback.attach(trainer) @@ -310,8 +334,21 @@ def attach_rl_callbacks( non-mlx backends, so this helper trusts the caller's config. """ attached = 0 + # Build the RL-checkpoint callback FIRST so the pid_lagrangian rollback + # ladder can be handed a reference to restore from. + ckpt_cb = _build_rl_checkpoint_cb(tcfg, output_dir=output_dir, task=task) + if ckpt_cb is not None: + trainer.add_callback(ckpt_cb) + attached += 1 + attached += _attach_reward_hack( - trainer, tcfg, buffer=buffer, tokenizer=tokenizer, output_dir=output_dir, task=task + trainer, + tcfg, + buffer=buffer, + tokenizer=tokenizer, + output_dir=output_dir, + task=task, + rl_checkpoint_cb=ckpt_cb, ) if bool(getattr(tcfg, "echo_trap_enabled", False)): @@ -333,39 +370,41 @@ def attach_rl_callbacks( except (TypeError, ValueError) as exc: logger.debug("attach echo-trap callback rejected: %s", exc) - save_every = getattr(tcfg, "rl_checkpoint_save_every_steps", None) - if save_every is not None: - from soup_cli.utils.rl_checkpoint import ( - RLCheckpointConfig, - build_rl_checkpoint_callback, - ) - - try: - ckpt_cfg = RLCheckpointConfig( - save_every_steps=int(save_every), - include_optimizer_state=bool( - getattr(tcfg, "rl_checkpoint_include_optimizer", True) - ), - include_ref_model=bool( - getattr(tcfg, "rl_checkpoint_include_ref_model", False) - ), - include_rollout_buffer=bool( - getattr(tcfg, "rl_checkpoint_include_rollout_buffer", False) - ), - keep_last=int(getattr(tcfg, "rl_checkpoint_keep_last", 3)), - ) - trainer.add_callback( - build_rl_checkpoint_callback( - ckpt_cfg, output_dir=output_dir, task=task - ) - ) - attached += 1 - except (TypeError, ValueError) as exc: - logger.debug("attach RL-checkpoint callback rejected: %s", exc) - return attached +def _build_rl_checkpoint_cb(tcfg: Any, *, output_dir: str, task: str) -> Any: + """Build the mid-epoch RL-checkpoint callback (or None if not configured).""" + save_every = getattr(tcfg, "rl_checkpoint_save_every_steps", None) + if save_every is None: + return None + from soup_cli.utils.rl_checkpoint import ( + RLCheckpointConfig, + build_rl_checkpoint_callback, + ) + + try: + ckpt_cfg = RLCheckpointConfig( + save_every_steps=int(save_every), + include_optimizer_state=bool( + getattr(tcfg, "rl_checkpoint_include_optimizer", True) + ), + include_ref_model=bool( + getattr(tcfg, "rl_checkpoint_include_ref_model", False) + ), + include_rollout_buffer=bool( + getattr(tcfg, "rl_checkpoint_include_rollout_buffer", False) + ), + keep_last=int(getattr(tcfg, "rl_checkpoint_keep_last", 3)), + ) + return build_rl_checkpoint_callback( + ckpt_cfg, output_dir=output_dir, task=task + ) + except (TypeError, ValueError) as exc: + logger.debug("build RL-checkpoint callback rejected: %s", exc) + return None + + def attach_plugin_callback(trainer: Any, console: Any = None) -> bool: """Attach :class:`SoupPluginCallback` when any enabled plugin implements a hook. diff --git a/src/soup_cli/utils/reward_hack_control.py b/src/soup_cli/utils/reward_hack_control.py index 691ee46..3bb6ec3 100644 --- a/src/soup_cli/utils/reward_hack_control.py +++ b/src/soup_cli/utils/reward_hack_control.py @@ -358,6 +358,89 @@ def bang_bang_step( return new_state, action +# --- PID-Lagrangian controller (Stage 2) --- + + +@dataclass(frozen=True) +class PIDLagrangianPolicy: + """PID-Lagrangian controller policy (Stooke et al. 2020). + + Treats "hacking signal ≤ ``signal_target``" as a constraint whose Lagrange + multiplier (the β / kl_coef) is updated by a PID law on the constraint + violation ``error = signal - target``. The integral term is clamped + (anti-windup) and the output is clamped to ``[beta_floor, beta_ceil]`` and + never crosses 0. + """ + + kp: float + ki: float + kd: float + signal_target: float + beta_floor: float + beta_ceil: float + integral_clamp: float + + def __post_init__(self) -> None: + _check_finite_float(self.kp, "kp", nonneg=True) + _check_finite_float(self.ki, "ki", nonneg=True) + _check_finite_float(self.kd, "kd", nonneg=True) + _check_finite_float(self.signal_target, "signal_target", nonneg=True) + _check_finite_float(self.beta_floor, "beta_floor", nonneg=True) + _check_finite_float(self.beta_ceil, "beta_ceil", nonneg=True) + _check_finite_float(self.integral_clamp, "integral_clamp", nonneg=True) + if not 0.0 <= self.signal_target < 1.0: + raise ValueError( + f"signal_target must be in [0, 1), got {self.signal_target}" + ) + if self.beta_floor <= 0.0: + raise ValueError(f"beta_floor must be > 0, got {self.beta_floor}") + if self.beta_floor >= self.beta_ceil: + raise ValueError( + f"beta_floor ({self.beta_floor}) must be < " + f"beta_ceil ({self.beta_ceil})" + ) + if self.integral_clamp <= 0.0: + raise ValueError( + f"integral_clamp must be > 0, got {self.integral_clamp}" + ) + + +def pid_step( + policy: PIDLagrangianPolicy, state: ControllerState, *, signal: float +) -> tuple[ControllerState, MitigationAction]: + """Advance the PID-Lagrangian controller one step for the hacking ``signal``. + + ``error = signal - target``; the integral accumulates (clamped ± + ``integral_clamp``); the multiplier β = clamp(floor..ceil, floor + Kp·error + + Ki·∫error + Kd·Δerror). β never crosses 0. + """ + fsignal = float(signal) + error = fsignal - policy.signal_target + integral = state.integral + error + integral = max(-policy.integral_clamp, min(policy.integral_clamp, integral)) + derivative = error - state.prev_error + control = policy.kp * error + policy.ki * integral + policy.kd * derivative + beta = max(policy.beta_floor, min(policy.beta_ceil, policy.beta_floor + control)) + tripped = beta > policy.beta_floor + + new_state = replace( + state, + step=state.step + 1, + beta=beta, + tripped=tripped, + integral=integral, + prev_error=error, + last_signal=min(1.0, max(0.0, fsignal)), + ) + action = MitigationAction( + new_beta=beta, + tripped=tripped, + verdict=_verdict_for(fsignal), + reason=f"pid: error={error:.3f} integral={integral:.3f} beta={beta:.4f}", + ) + return new_state, action + + class MitigationLogWriter: """Thread-safe append-only JSONL log for the reward-hack controller. @@ -483,6 +566,11 @@ class RewardHackMitigationCallback(_TrainerCallbackBase): # type: ignore[misc, tokenizer: Any = None, task: str = "grpo", bang_bang: BangBangPolicy | None = None, + pid: PIDLagrangianPolicy | None = None, + rollback: bool = False, + rollback_patience: int = 3, + max_recovery_attempts: int = 2, + rl_checkpoint_cb: Any = None, ) -> None: if mode not in MITIGATION_MODES: raise ValueError( @@ -508,6 +596,13 @@ class RewardHackMitigationCallback(_TrainerCallbackBase): # type: ignore[misc, self.bang_bang = bang_bang if mode == "kl_control" and bang_bang is None: raise ValueError("kl_control mode requires a BangBangPolicy") + self.pid = pid + if mode == "pid_lagrangian" and pid is None: + raise ValueError("pid_lagrangian mode requires a PIDLagrangianPolicy") + self.rollback = bool(rollback) + self.rollback_patience = int(rollback_patience) + self.max_recovery_attempts = int(max_recovery_attempts) + self.rl_checkpoint_cb = rl_checkpoint_cb # Compose the v0.70.0 detector callback for the info_rm/rm_ensemble # baseline + drop_pct logic (DRY — no re-implementation). self._detector_cb = RewardHackCallback( @@ -516,6 +611,8 @@ class RewardHackMitigationCallback(_TrainerCallbackBase): # type: ignore[misc, self._trainer: Any = None self._state = ControllerState() self._length_baseline: float | None = None + self._hack_streak = 0 + self._last_good_step: int | None = None def attach(self, trainer: Any) -> None: """Store the trainer reference (the β / kl_coef mutation target).""" @@ -640,6 +737,72 @@ class RewardHackMitigationCallback(_TrainerCallbackBase): # type: ignore[misc, telemetry["tripped"] = action.tripped telemetry["action"] = action.reason + def _request_stop(self, control: Any) -> None: + if control is not None: + try: + control.should_training_stop = True + except Exception: # noqa: BLE001 + pass + + def _escalate( + self, model: Any, optimizer: Any, control: Any, telemetry: dict[str, Any] + ) -> Any: + """Escalation ladder rung: rollback to last-good, else early-stop.""" + if self._state.recovery_attempts >= self.max_recovery_attempts: + telemetry["escalation"] = "early_stop" + self._request_stop(control) + return control + target = self._last_good_step + restored = False + if target is not None and self.rl_checkpoint_cb is not None: + try: + restored = bool( + self.rl_checkpoint_cb.restore_checkpoint( + step=target, model=model, optimizer=optimizer + ) + ) + except Exception: # noqa: BLE001 — rollback must never crash the run + restored = False + self._state = replace( + self._state, recovery_attempts=self._state.recovery_attempts + 1 + ) + self._hack_streak = 0 + telemetry["escalation"] = f"rollback to step {target} (restored={restored})" + return control + + def _run_pid( + self, + telemetry: dict[str, Any], + signals: Mapping[str, float], + model: Any, + optimizer: Any, + control: Any, + ) -> Any: + """pid_lagrangian: PID β update + rollback escalation ladder.""" + policy = self.pid + if policy is None: + return control + self._seed_coefficient(policy.beta_floor) + vote = combine_signals(signals, self.signals) + new_state, action = pid_step(policy, self._state, signal=vote) + self._state = new_state + self._apply_coefficient(action.new_beta) + telemetry["vote"] = vote + telemetry["new_beta"] = action.new_beta + telemetry["tripped"] = action.tripped + telemetry["action"] = action.reason + # Escalation ladder: raise (above) → rollback → early-stop. + if action.verdict == "HACK": + self._hack_streak += 1 + else: + self._hack_streak = 0 + saved = getattr(self.rl_checkpoint_cb, "_saved", None) + if saved: + self._last_good_step = max(saved) + if self.rollback and self._hack_streak >= self.rollback_patience: + control = self._escalate(model, optimizer, control, telemetry) + return control + def on_step_end(self, args, state, control, **kwargs): """Per-step hook — read the buffer, compute telemetry, act by mode. @@ -652,10 +815,18 @@ class RewardHackMitigationCallback(_TrainerCallbackBase): # type: ignore[misc, snapshot = self.buffer.snapshot() step = int(getattr(state, "global_step", 0) or 0) telemetry, signals = self._observe(snapshot, step) - # log_only observes; kl_control drives the bang-bang controller. - # pid_lagrangian is wired in Stage 2. + # log_only observes; kl_control drives bang-bang; pid_lagrangian + # drives the PID controller + rollback escalation ladder. if self.mode == "kl_control": self._run_bang_bang(telemetry, signals) + elif self.mode == "pid_lagrangian": + control = self._run_pid( + telemetry, + signals, + kwargs.get("model"), + kwargs.get("optimizer"), + control, + ) self.log_writer.record(step=step, snapshot=telemetry) return control except Exception: # noqa: BLE001 — instrumentation must never crash diff --git a/src/soup_cli/utils/rl_checkpoint.py b/src/soup_cli/utils/rl_checkpoint.py index 42a2e07..4b5b543 100644 --- a/src/soup_cli/utils/rl_checkpoint.py +++ b/src/soup_cli/utils/rl_checkpoint.py @@ -313,6 +313,52 @@ class RLCheckpointCallback(_TrainerCallbackBase): # type: ignore[misc, valid-ty self._prune() return ckpt_dir + def restore_checkpoint(self, *, step: int, model, optimizer) -> bool: + """Reload a saved checkpoint's PEFT adapter + optimizer state (v0.71.26). + + Used by the reward-hack mitigation rollback ladder to hop back to the + last-good checkpoint. Reloads the adapter via + ``set_peft_model_state_dict`` and the optimizer via ``load_state_dict``. + Best-effort — returns ``True`` when at least one artifact was restored, + ``False`` otherwise, and NEVER raises (rollback must not crash the run). + """ + import os + + ckpt_dir = os.path.join(self._ckpt_root(), f"step-{int(step)}") + if not os.path.isdir(ckpt_dir): + return False + restored = False + if model is not None: + try: + from peft import ( + PeftModel, + load_peft_weights, + set_peft_model_state_dict, + ) + + if isinstance(model, PeftModel): + set_peft_model_state_dict(model, load_peft_weights(ckpt_dir)) + restored = True + except Exception: # noqa: BLE001 — best-effort restore, never crash + pass + if optimizer is not None: + opt_path = os.path.join(ckpt_dir, "optimizer.pt") + if os.path.isfile(opt_path): + try: + import torch + + # Trusted file (we wrote it under the cwd-contained run dir); + # weights_only=False loads the full optimizer state_dict. + optimizer.load_state_dict( + torch.load( + opt_path, map_location="cpu", weights_only=False + ) + ) + restored = True + except Exception: # noqa: BLE001 + pass + return restored + def _prune(self) -> None: """Keep only the ``keep_last`` most-recent step checkpoints.""" import os diff --git a/tests/test_v07126.py b/tests/test_v07126.py index d7ef616..ad8e95e 100644 --- a/tests/test_v07126.py +++ b/tests/test_v07126.py @@ -1037,3 +1037,397 @@ class TestRewardHackMitigationCli: src = inspect.getsource(train_mod) assert "--reward-hack-mitigation" in src assert "cfg.training.reward_hack_mitigation" in src + + +# ===================================================================== +# Part C / Stage 2 — schema fields (PID + rollback) (Task C1) +# ===================================================================== + + +class TestStage2Schema: + """PID-Lagrangian + rollback tunables + bounds + cross-validators.""" + + def _cfg(self, extra: str = "", *, mitigation: str = "pid_lagrangian"): + from soup_cli.config.loader import load_config_from_string + + return load_config_from_string(_yaml("grpo", mitigation=mitigation, extra=extra)) + + def test_pid_parses(self): + cfg = self._cfg("reward_hack_pid_kp: 0.8\nreward_hack_signal_target: 0.2") + assert cfg.training.reward_hack_mitigation == "pid_lagrangian" + assert cfg.training.reward_hack_pid_kp == 0.8 + assert cfg.training.reward_hack_signal_target == 0.2 + + def test_defaults(self): + from soup_cli.config.schema import TrainingConfig + + t = TrainingConfig() + assert t.reward_hack_pid_kp == 0.5 and t.reward_hack_pid_ki == 0.1 + assert t.reward_hack_pid_kd == 0.05 and t.reward_hack_signal_target == 0.15 + assert t.reward_hack_rollback is False + assert t.reward_hack_rollback_patience == 3 + assert t.reward_hack_max_recovery_attempts == 2 + + def test_rollback_requires_checkpoint_cadence(self): + with pytest.raises(ValueError, match="rl_checkpoint_save_every_steps"): + self._cfg("reward_hack_rollback: true") + + def test_rollback_with_cadence_ok(self): + cfg = self._cfg("reward_hack_rollback: true\nrl_checkpoint_save_every_steps: 10") + assert cfg.training.reward_hack_rollback is True + + def test_pid_param_requires_pid_mode(self): + # setting a PID param under kl_control is a no-op footgun. + with pytest.raises(ValueError, match="pid_lagrangian"): + self._cfg("reward_hack_pid_kp: 0.9", mitigation="kl_control") + + def test_pid_param_under_off_rejected(self): + from soup_cli.config.loader import load_config_from_string + + with pytest.raises(ValueError, match="reward_hack_mitigation"): + load_config_from_string( + _yaml( + "grpo", + detector=None, + mitigation="off", + extra="reward_hack_pid_kp: 0.9", + ) + ) + + def test_pid_kp_non_negative(self): + with pytest.raises(ValueError): + self._cfg("reward_hack_pid_kp: -0.1") + + def test_signal_target_below_one(self): + with pytest.raises(ValueError): + self._cfg("reward_hack_signal_target: 1.0") + + +# ===================================================================== +# Part C / Stage 2 — PIDLagrangianPolicy + pid_step (Task C2) +# ===================================================================== + + +def _pid_policy(**kw): + from soup_cli.utils.reward_hack_control import PIDLagrangianPolicy + + defaults = dict( + kp=1.0, + ki=0.0, + kd=0.0, + signal_target=0.15, + beta_floor=0.02, + beta_ceil=10.0, + integral_clamp=10.0, + ) + defaults.update(kw) + return PIDLagrangianPolicy(**defaults) + + +def _run_pid(policy, signals, beta0=None): + from soup_cli.utils.reward_hack_control import ControllerState, pid_step + + state = ControllerState(beta=beta0 if beta0 is not None else policy.beta_floor) + states = [] + for signal in signals: + state, _ = pid_step(policy, state, signal=signal) + states.append(state) + return states + + +class TestPidLagrangian: + """PID-Lagrangian controller: P/I/D isolated, anti-windup, output clamp.""" + + def test_policy_gains_non_negative(self): + with pytest.raises(ValueError, match="kp"): + _pid_policy(kp=-1.0) + + def test_policy_floor_lt_ceil(self): + with pytest.raises(ValueError, match="floor"): + _pid_policy(beta_floor=5.0, beta_ceil=1.0) + + def test_policy_integral_clamp_positive(self): + with pytest.raises(ValueError, match="integral_clamp"): + _pid_policy(integral_clamp=0.0) + + def test_proportional_only(self): + # kp=1, error = 0.65 - 0.15 = 0.5 → β = floor + 0.5 + states = _run_pid(_pid_policy(kp=1.0, ki=0.0, kd=0.0), [0.65]) + assert states[0].beta == pytest.approx(0.02 + 0.5) + + def test_integral_accumulates(self): + states = _run_pid(_pid_policy(kp=0.0, ki=1.0, kd=0.0), [0.35, 0.35]) + assert states[1].beta > states[0].beta # integral grows β + + def test_integral_anti_windup(self): + # constant error 0.5 with clamp 0.3 → integral saturates → β flat. + states = _run_pid( + _pid_policy(kp=0.0, ki=1.0, kd=0.0, integral_clamp=0.3, beta_ceil=100.0), + [0.65] * 5, + ) + assert states[0].beta == pytest.approx(0.02 + 0.3) + assert states[4].beta == pytest.approx(states[0].beta) + + def test_derivative_spikes_then_settles(self): + # kd=1: β spikes on the error jump then returns to floor on constant error. + states = _run_pid( + _pid_policy(kp=0.0, ki=0.0, kd=1.0), [0.15, 0.65, 0.65] + ) + assert states[1].beta > 0.02 # jump + assert states[2].beta == pytest.approx(0.02) # deriv back to 0 + + def test_output_clamped_to_ceil(self): + states = _run_pid(_pid_policy(kp=1000.0, beta_ceil=0.5), [0.9]) + assert states[0].beta == pytest.approx(0.5) + + def test_output_clamped_to_floor(self): + # negative error → control negative → β clamped to floor (never < 0). + states = _run_pid(_pid_policy(kp=1.0), [0.0]) + assert states[0].beta == pytest.approx(0.02) + + def test_relaxes_after_raise(self): + # raise on positive error, then relax when the signal drops below target. + pol = _pid_policy(kp=0.0, ki=1.0, kd=0.0, integral_clamp=10.0, beta_ceil=100.0) + states = _run_pid(pol, [0.65, 0.65, 0.0, 0.0, 0.0]) + assert states[-1].beta < states[1].beta + + +# ===================================================================== +# Part C / Stage 2 — RLCheckpointCallback.restore_checkpoint (Task C3) +# ===================================================================== + + +class _FakeSavableModel: + def __init__(self): + self.saved_to = None + + def save_pretrained(self, path): + import os + + os.makedirs(path, exist_ok=True) + self.saved_to = path + + +class TestRestoreCheckpoint: + """restore_checkpoint reloads adapter + optimizer state (RL rollback).""" + + def _cb(self, tmp_path): + from soup_cli.utils.rl_checkpoint import ( + RLCheckpointConfig, + build_rl_checkpoint_callback, + ) + + return build_rl_checkpoint_callback( + RLCheckpointConfig(save_every_steps=1), output_dir="run", task="grpo" + ) + + def test_restore_optimizer_roundtrip(self, tmp_path, monkeypatch): + import copy + + import torch + + monkeypatch.chdir(tmp_path) + cb = self._cb(tmp_path) + param = torch.nn.Parameter(torch.zeros(2)) + opt = torch.optim.SGD([param], lr=0.1, momentum=0.9) + param.grad = torch.ones(2) + opt.step() # creates a momentum buffer + saved = copy.deepcopy(opt.state_dict()["state"][0]["momentum_buffer"]) + cb.save_checkpoint(step=1, model=_FakeSavableModel(), optimizer=opt) + # mutate optimizer state + param.grad = torch.ones(2) * 7 + opt.step() + assert not torch.allclose( + opt.state_dict()["state"][0]["momentum_buffer"], saved + ) + ok = cb.restore_checkpoint(step=1, model=_FakeSavableModel(), optimizer=opt) + assert ok is True + assert torch.allclose( + opt.state_dict()["state"][0]["momentum_buffer"], saved + ) + + def test_restore_missing_step_returns_false(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + cb = self._cb(tmp_path) + assert cb.restore_checkpoint(step=99, model=None, optimizer=None) is False + + def test_restore_adapter_roundtrip(self, tmp_path, monkeypatch): + torch = pytest.importorskip("torch") + peft = pytest.importorskip("peft") + import torch.nn as nn + + monkeypatch.chdir(tmp_path) + + class Tiny(nn.Module): + def __init__(self): + super().__init__() + self.lin = nn.Linear(4, 4, bias=False) + + def forward(self, x): + return self.lin(x) + + model = peft.get_peft_model( + Tiny(), peft.LoraConfig(target_modules=["lin"], r=2) + ) + key = "base_model.model.lin.lora_A.default.weight" + original = model.state_dict()[key].clone() + cb = self._cb(tmp_path) + cb.save_checkpoint(step=1, model=model, optimizer=None) + # mutate the LoRA weight in place + with torch.no_grad(): + dict(model.named_parameters())[ + "base_model.model.lin.lora_A.default.weight" + ].add_(5.0) + assert not torch.allclose(model.state_dict()[key], original) + ok = cb.restore_checkpoint(step=1, model=model, optimizer=None) + assert ok is True + assert torch.allclose(model.state_dict()[key], original) + + +# ===================================================================== +# Part C / Stage 2 — pid_lagrangian callback + escalation ladder (Task C4) +# ===================================================================== + + +class _FakeCkptCb: + def __init__(self, saved): + self._saved = list(saved) + self.restore_calls = [] + + def restore_checkpoint(self, *, step, model, optimizer): + self.restore_calls.append(step) + return True + + +def _pid_callback(tmp_path, buffer, *, rollback=False, rollback_patience=2, + max_recovery_attempts=1, ckpt_cb=None): + from soup_cli.utils.reward_hack_control import ( + MitigationLogWriter, + PIDLagrangianPolicy, + RewardHackMitigationCallback, + ) + + pid = PIDLagrangianPolicy( + kp=1.0, ki=0.0, kd=0.0, signal_target=0.15, + beta_floor=0.02, beta_ceil=10.0, integral_clamp=10.0, + ) + return RewardHackMitigationCallback( + mode="pid_lagrangian", + detector="info_rm", + log_writer=MitigationLogWriter(str(tmp_path / "m.jsonl")), + buffer=buffer, + task="grpo", + pid=pid, + rollback=rollback, + rollback_patience=rollback_patience, + max_recovery_attempts=max_recovery_attempts, + rl_checkpoint_cb=ckpt_cb, + ) + + +class TestPidCallback: + """pid_lagrangian drives β via PID and runs the rollback escalation ladder.""" + + def test_pid_requires_policy(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + from soup_cli.utils.reward_hack_control import ( + MitigationLogWriter, + RewardHackMitigationCallback, + ) + + with pytest.raises(ValueError, match="pid_lagrangian"): + RewardHackMitigationCallback( + mode="pid_lagrangian", + detector="info_rm", + log_writer=MitigationLogWriter("m.jsonl"), + ) + + def test_pid_mutates_beta_on_hack(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + cb = _pid_callback(tmp_path, _SeqBuffer([_HEALTHY, _HACK])) + trainer = _fake_grpo_trainer(beta=0.02) + cb.attach(trainer) + cb.on_step_end(None, types.SimpleNamespace(global_step=1), None) + cb.on_step_end(None, types.SimpleNamespace(global_step=2), None) + assert trainer.beta > 0.5 # PID raised β on the hacking step + assert trainer.args.beta == trainer.beta # dual write + + def test_escalation_rollback_then_stop(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + ckpt = _FakeCkptCb(saved=[10]) + cb = _pid_callback( + tmp_path, + _SeqBuffer([_HEALTHY, _HACK, _HACK, _HACK, _HACK]), + rollback=True, + rollback_patience=2, + max_recovery_attempts=1, + ckpt_cb=ckpt, + ) + cb.attach(_fake_grpo_trainer(beta=0.02)) + control = types.SimpleNamespace(should_training_stop=False) + model, opt = object(), object() + for step in range(1, 6): + control = cb.on_step_end( + None, + types.SimpleNamespace(global_step=step), + control, + model=model, + optimizer=opt, + ) + # rolled back to the last-good saved checkpoint (step 10)… + assert ckpt.restore_calls == [10] + # …then early-stopped after recovery attempts were exhausted. + assert control.should_training_stop is True + + def test_no_rollback_when_disabled(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + ckpt = _FakeCkptCb(saved=[10]) + cb = _pid_callback( + tmp_path, + _SeqBuffer([_HEALTHY, _HACK, _HACK, _HACK]), + rollback=False, + ckpt_cb=ckpt, + ) + cb.attach(_fake_grpo_trainer(beta=0.02)) + control = types.SimpleNamespace(should_training_stop=False) + for step in range(1, 5): + control = cb.on_step_end( + None, types.SimpleNamespace(global_step=step), control, + model=object(), optimizer=object(), + ) + assert ckpt.restore_calls == [] # rollback disabled + assert control.should_training_stop is False + + +class TestAttachPidControl: + """attach_rl_callbacks builds the PID policy + wires the checkpoint ref.""" + + def test_attach_pid_builds_policy_and_ckpt_ref(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + from soup_cli.config.schema import TrainingConfig + from soup_cli.utils.peft_wiring import attach_rl_callbacks + from soup_cli.utils.reward_hack_control import RewardHackMitigationCallback + from soup_cli.utils.rl_checkpoint import RLCheckpointCallback + + tcfg = TrainingConfig( + reward_hack_mitigation="pid_lagrangian", + reward_hack_detector="info_rm", + reward_hack_rollback=True, + rl_checkpoint_save_every_steps=5, + reward_hack_pid_kp=0.7, + ) + added: list = [] + attach_rl_callbacks( + _fake_trainer_recording(added), + tcfg, + buffer=object(), + output_dir=str(tmp_path), + task="grpo", + ) + mit = [c for c in added if isinstance(c, RewardHackMitigationCallback)] + assert len(mit) == 1 and mit[0].mode == "pid_lagrangian" + assert mit[0].pid is not None and mit[0].pid.kp == 0.7 + assert mit[0].rollback is True + ckpts = [c for c in added if isinstance(c, RLCheckpointCallback)] + assert len(ckpts) == 1 + assert mit[0].rl_checkpoint_cb is ckpts[0]