mirror of https://github.com/razor-ai/soup.git
feat(train): PID-Lagrangian controller + rollback escalation ladder (v0.71.26 Part C)
Stage 2 of closed-loop reward-hacking mitigation: principled control + safety net. - schema: 7 Stage-2 tunables (pid_kp/ki/kd, signal_target, rollback, rollback_patience, max_recovery_attempts) + bool guard; cross-validators — PID/rollback tunables require pid_lagrangian mode; reward_hack_rollback requires rl_checkpoint_save_every_steps (a cadence to roll back to). - reward_hack_control: PIDLagrangianPolicy + pid_step (Stooke et al. — constraint 'signal <= target' whose multiplier β is a PID law; integral anti-windup clamp, output clamp [floor,ceil], never crosses 0). Callback pid_lagrangian path drives β via PID + escalation ladder: raise -> rollback last-good checkpoint (after rollback_patience HACK steps) -> early-stop (after max_recovery_attempts), subsuming halt_on_hack. - rl_checkpoint: RLCheckpointCallback.restore_checkpoint (reload PEFT adapter via set_peft_model_state_dict + optimizer.load_state_dict; best-effort, never raises). - peft_wiring: build PIDLagrangianPolicy from tcfg; build the RL-checkpoint cb FIRST and hand its reference to the mitigation callback for rollback. +26 tests (test_v07126: 91 -> 117; real-peft adapter round-trip included).
This commit is contained in:
Alpamys
parent
925e600702
commit
21557b43c1
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
Loading…
Reference in New Issue