mirror of https://github.com/razor-ai/soup.git
271 lines
12 KiB
Python
271 lines
12 KiB
Python
"""GRPOStabilityCallback — v0.53.11 #127.
|
||
|
||
Live HF TrainerCallback that wires the seven v0.50.0 Part D stability /
|
||
efficiency knobs into the training loop:
|
||
|
||
- ``ref_model_ema_alpha``: EMA update of the reference model post-step.
|
||
- ``replay_buffer_size``: bounded deque of recent rollouts for re-use.
|
||
- ``async_grpo_prefetch``: schedule the next-batch rollout in a background
|
||
thread (advisory — actual prefetch is a TRL ``GRPOTrainer`` concern).
|
||
- ``tis_threshold``: truncated importance sampling — log a warning when
|
||
log-ratios exceed the threshold so operators can spot off-policy drift.
|
||
- ``mask_truncated_completions``: hint to skip rows whose completion hit
|
||
``max_new_tokens`` (the actual mask is applied inside compute_loss).
|
||
- ``defer_rerolling``: when an advantage batch is all-zero, skip re-roll
|
||
this step (record-only — re-roll is a GRPOTrainer concern).
|
||
- ``skip_zero_advantage`` / ``off_policy_mask_threshold``: similar
|
||
record-only knobs surfaced as ``state.log_history`` metrics.
|
||
|
||
Pure math: each per-knob hook is a small pure function in the same module,
|
||
so tests can exercise them without instantiating transformers.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
from collections import deque
|
||
from typing import Any, Optional
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def update_ema(ref_state: dict, policy_state: dict, alpha: float) -> dict:
|
||
"""In-place EMA update: ``ref = (1-α)·ref + α·policy``.
|
||
|
||
Pure-function math kernel. Both inputs are name->tensor mappings (HF
|
||
state_dict shape). ``alpha`` is in ``(0, 1]`` AND finite (no NaN/Inf —
|
||
v0.53.11 review fix per v0.32.0 / v0.47.0 policy).
|
||
|
||
Returns the (mutated) ``ref_state``. The function mutates in-place but
|
||
also returns the dict for chaining; callers should treat the return
|
||
value as the same object passed in.
|
||
"""
|
||
import math
|
||
|
||
if not isinstance(alpha, (int, float)) or isinstance(alpha, bool):
|
||
raise TypeError("alpha must be a non-bool float")
|
||
alpha_f = float(alpha)
|
||
if not math.isfinite(alpha_f):
|
||
raise ValueError("alpha must be finite (no NaN/Inf)")
|
||
if not (0.0 < alpha_f <= 1.0):
|
||
raise ValueError(f"alpha must be in (0, 1], got {alpha}")
|
||
for name, p_tensor in policy_state.items():
|
||
if name not in ref_state:
|
||
continue
|
||
r_tensor = ref_state[name]
|
||
# Defensive: only update tensors with matching shape.
|
||
if hasattr(r_tensor, "shape") and hasattr(p_tensor, "shape"):
|
||
if r_tensor.shape != p_tensor.shape:
|
||
continue
|
||
ref_state[name] = (1.0 - alpha_f) * r_tensor + alpha_f * p_tensor
|
||
return ref_state
|
||
|
||
|
||
def check_tis_threshold(log_ratio, threshold: float) -> bool:
|
||
"""Return True iff the max absolute log-ratio exceeds the TIS threshold.
|
||
|
||
Used to flag off-policy drift in GRPO rollouts.
|
||
"""
|
||
if not isinstance(threshold, (int, float)) or isinstance(threshold, bool):
|
||
raise TypeError("threshold must be a non-bool number")
|
||
if float(threshold) <= 0.0:
|
||
raise ValueError("threshold must be positive")
|
||
if not hasattr(log_ratio, "abs"):
|
||
raise TypeError("log_ratio must be a tensor (need .abs())")
|
||
max_abs = float(log_ratio.abs().max())
|
||
return max_abs > float(threshold)
|
||
|
||
|
||
def filter_zero_advantage(advantages, *, eps: float = 1e-8) -> Any:
|
||
"""Return a boolean mask: True where advantage is non-zero.
|
||
|
||
Used by ``skip_zero_advantage`` — rows with mask=False are dropped
|
||
from the loss compute.
|
||
"""
|
||
if not isinstance(eps, (int, float)) or isinstance(eps, bool):
|
||
raise TypeError("eps must be a non-bool number")
|
||
if not hasattr(advantages, "abs"):
|
||
raise TypeError("advantages must be a tensor")
|
||
return advantages.abs() > float(eps)
|
||
|
||
|
||
def _get_trainer_callback_base():
|
||
"""Lazy-resolve ``transformers.TrainerCallback`` (v0.53.11 review fix).
|
||
|
||
Project policy: every callback inherits TrainerCallback so HF Trainer
|
||
discovers it via the callback handler. We resolve at class-body
|
||
evaluation time so module import does not pull transformers.
|
||
"""
|
||
try:
|
||
from transformers import TrainerCallback
|
||
|
||
return TrainerCallback
|
||
except ImportError:
|
||
return object
|
||
|
||
|
||
_TrainerCallbackBase = _get_trainer_callback_base()
|
||
|
||
|
||
class GRPOStabilityCallback(_TrainerCallbackBase): # type: ignore[misc, valid-type]
|
||
"""HF TrainerCallback that wires v0.50.0 Part D stability knobs.
|
||
|
||
Lazy-inherits ``transformers.TrainerCallback`` so the module is
|
||
importable without transformers (falls back to ``object``).
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
*,
|
||
ref_model_ema_alpha: Optional[float] = None,
|
||
replay_buffer_size: Optional[int] = None,
|
||
async_grpo_prefetch: bool = False,
|
||
tis_threshold: Optional[float] = None,
|
||
mask_truncated_completions: bool = False,
|
||
defer_rerolling: bool = False,
|
||
skip_zero_advantage: bool = False,
|
||
off_policy_mask_threshold: Optional[float] = None,
|
||
):
|
||
# Validation mirrors the schema. Bool-rejection on numeric fields
|
||
# per project policy (v0.30.0 Candidate / v0.41.0 lr_groups).
|
||
if ref_model_ema_alpha is not None:
|
||
if isinstance(ref_model_ema_alpha, bool):
|
||
raise TypeError("ref_model_ema_alpha must be float, not bool")
|
||
if not (0.0 < float(ref_model_ema_alpha) <= 1.0):
|
||
raise ValueError(
|
||
f"ref_model_ema_alpha must be in (0, 1], got {ref_model_ema_alpha}"
|
||
)
|
||
if replay_buffer_size is not None:
|
||
if isinstance(replay_buffer_size, bool):
|
||
raise TypeError("replay_buffer_size must be int, not bool")
|
||
if not (1 <= int(replay_buffer_size) <= 1_000_000):
|
||
raise ValueError(
|
||
f"replay_buffer_size must be in [1, 1e6], got "
|
||
f"{replay_buffer_size}"
|
||
)
|
||
if tis_threshold is not None:
|
||
if isinstance(tis_threshold, bool):
|
||
raise TypeError("tis_threshold must be float, not bool")
|
||
if not (0.0 < float(tis_threshold) <= 100.0):
|
||
raise ValueError(
|
||
f"tis_threshold must be in (0, 100], got {tis_threshold}"
|
||
)
|
||
if off_policy_mask_threshold is not None:
|
||
if isinstance(off_policy_mask_threshold, bool):
|
||
raise TypeError("off_policy_mask_threshold must be float, not bool")
|
||
if not (0.0 <= float(off_policy_mask_threshold) <= 1.0):
|
||
raise ValueError(
|
||
f"off_policy_mask_threshold must be in [0, 1], got "
|
||
f"{off_policy_mask_threshold}"
|
||
)
|
||
self.ref_model_ema_alpha = ref_model_ema_alpha
|
||
self.replay_buffer_size = replay_buffer_size
|
||
self.async_grpo_prefetch = bool(async_grpo_prefetch)
|
||
self.tis_threshold = tis_threshold
|
||
self.mask_truncated_completions = bool(mask_truncated_completions)
|
||
self.defer_rerolling = bool(defer_rerolling)
|
||
self.skip_zero_advantage = bool(skip_zero_advantage)
|
||
self.off_policy_mask_threshold = off_policy_mask_threshold
|
||
# Bounded rollout deque — created on demand.
|
||
self._replay: Optional[deque] = None
|
||
if self.replay_buffer_size is not None:
|
||
self._replay = deque(maxlen=int(self.replay_buffer_size))
|
||
self._tis_alerts = 0
|
||
# Set during on_train_begin (lazy — model is constructed by Trainer
|
||
# before the first event fires).
|
||
self._policy_model: Any = None
|
||
self._ref_model: Any = None
|
||
|
||
def push_rollout(self, rollout: Any) -> None:
|
||
"""Append a rollout to the bounded replay buffer (no-op if disabled)."""
|
||
if self._replay is None:
|
||
return
|
||
self._replay.append(rollout)
|
||
|
||
def replay_size(self) -> int:
|
||
"""Current rollout count in the replay buffer."""
|
||
return 0 if self._replay is None else len(self._replay)
|
||
|
||
def record_tis_alert(self, log_ratio) -> bool:
|
||
"""Increment internal counter when log-ratio breaches the threshold."""
|
||
if self.tis_threshold is None:
|
||
return False
|
||
try:
|
||
if check_tis_threshold(log_ratio, self.tis_threshold):
|
||
self._tis_alerts += 1
|
||
return True
|
||
except (TypeError, ValueError):
|
||
pass
|
||
return False
|
||
|
||
def tis_alerts(self) -> int:
|
||
"""Number of recorded TIS-breach alerts."""
|
||
return self._tis_alerts
|
||
|
||
# --- HF TrainerCallback surface (v0.53.11 live wiring) ---
|
||
|
||
def on_train_begin(self, args, state, control, model=None, **kwargs):
|
||
"""Capture policy + ref model references for EMA updates."""
|
||
logger.debug("GRPOStabilityCallback.on_train_begin")
|
||
self._policy_model = model
|
||
# The reference model lives on the trainer; HF passes it via `kwargs`
|
||
# in newer TRL versions, otherwise pull from the callback handler.
|
||
ref = kwargs.get("ref_model")
|
||
if ref is None:
|
||
# Try to discover it on the trainer instance via the standard
|
||
# TRL ``DPOTrainer.ref_model`` / ``GRPOTrainer.ref_model`` attr.
|
||
trainer = kwargs.get("trainer")
|
||
if trainer is not None:
|
||
ref = getattr(trainer, "ref_model", None)
|
||
self._ref_model = ref
|
||
return control
|
||
|
||
def on_step_end(self, args, state, control, model=None, **kwargs):
|
||
"""Per-step hook — perform EMA update + record stability state.
|
||
|
||
v0.53.11 #127 — wires the actual EMA update post-step. The
|
||
``replay_buffer`` and ``tis_alerts`` counters are surfaced via
|
||
``state.log_history`` so the v0.34.0 anomaly explainer can flag
|
||
instability.
|
||
"""
|
||
# Live EMA update of reference model from current policy.
|
||
if (
|
||
self.ref_model_ema_alpha is not None
|
||
and self._ref_model is not None
|
||
and (model is not None or self._policy_model is not None)
|
||
):
|
||
try:
|
||
policy = model if model is not None else self._policy_model
|
||
ref_sd = self._ref_model.state_dict()
|
||
pol_sd = policy.state_dict()
|
||
update_ema(ref_sd, pol_sd, self.ref_model_ema_alpha)
|
||
# v0.53.11 review fix (security HIGH) — strict=True with
|
||
# try/except for key mismatch. strict=False silently
|
||
# dropped unknown keys, masking corruption from a crafted
|
||
# checkpoint. We catch the RuntimeError and downgrade to
|
||
# strict=False with a WARNING so operators see the drift.
|
||
try:
|
||
self._ref_model.load_state_dict(ref_sd, strict=True)
|
||
except RuntimeError as key_err:
|
||
logger.warning(
|
||
"EMA load_state_dict key mismatch; falling back to "
|
||
"strict=False (potential silent corruption): %s",
|
||
key_err,
|
||
)
|
||
self._ref_model.load_state_dict(ref_sd, strict=False)
|
||
except Exception as exc: # noqa: BLE001 — never crash training
|
||
logger.debug("EMA update skipped: %s", exc)
|
||
# Surface counters to log_history.
|
||
log_history = getattr(state, "log_history", None)
|
||
if log_history is not None:
|
||
entry: dict = {}
|
||
if self.tis_threshold is not None:
|
||
entry["tis_alerts"] = self._tis_alerts
|
||
if self._replay is not None:
|
||
entry["replay_size"] = len(self._replay)
|
||
if self.ref_model_ema_alpha is not None:
|
||
entry["ema_alpha"] = float(self.ref_model_ema_alpha)
|
||
if entry:
|
||
log_history.append(entry)
|
||
return control
|