mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-09-22 01:10:49 +02:00
fix(bayesian): surface skipped optimization and clamp warm start to search space
This commit is contained in:
@@ -4231,6 +4231,17 @@ class AbliterationPipeline:
|
||||
f"{idx}:{reg:.3f}" for idx, reg in sorted(bayesian_regs.items())
|
||||
)
|
||||
self.log(f" Optimal regs: {regs_str}")
|
||||
else:
|
||||
# Reached whenever the optimizer no-ops for any reason. Without
|
||||
# this the run logs "Running Bayesian optimization (50 trials)",
|
||||
# prints nothing further, and completes as a success — so the
|
||||
# only difference between an optimized and an unoptimized
|
||||
# checkpoint is an absence of output.
|
||||
self.log(
|
||||
"WARNING: Bayesian optimization returned no layer "
|
||||
f"regularizations (reason: {getattr(self, '_bayesian_skipped', 'unknown')}). "
|
||||
"Continuing with method defaults — this checkpoint is NOT optimized."
|
||||
)
|
||||
|
||||
# ── LoRA-based reversible ablation ──────────────────────────────
|
||||
# When enabled, compute LoRA adapters and merge them instead of
|
||||
@@ -7584,6 +7595,11 @@ class AbliterationPipeline:
|
||||
"som_diversity_penalty": self.som_diversity_penalty if self.direction_method == "som" else None,
|
||||
"som_min_signal_to_noise": self.som_min_signal_to_noise if self.direction_method == "som" else None,
|
||||
"layer_selection": self.layer_selection,
|
||||
# None when the optimizer ran (or was never requested); a reason
|
||||
# string when it no-opped. Without it a checkpoint's provenance
|
||||
# records the method that ASKED for an optimization, with nothing
|
||||
# to say whether one happened.
|
||||
"bayesian_optimization_skipped": getattr(self, "_bayesian_skipped", None),
|
||||
"min_layer_fraction": self.min_layer_fraction,
|
||||
"max_layer_fraction": self.max_layer_fraction,
|
||||
"harmless_pc_count": self.harmless_pc_count,
|
||||
|
||||
@@ -41,6 +41,56 @@ if TYPE_CHECKING:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Kernel search space, declared once. `objective` suggests from these bounds and
|
||||
# the warm-start trial is clamped to them; keeping two copies is what let the
|
||||
# warm start drift outside the space it was seeding (see `_clamp_to_space`).
|
||||
KERNEL_SPACE: dict[str, tuple[float, float]] = {
|
||||
"attn_max_weight": (0.5, 1.0),
|
||||
"attn_peak_position": (0.1, 0.9),
|
||||
"attn_min_weight": (0.0, 0.3),
|
||||
"attn_spread": (0.1, 0.6),
|
||||
"mlp_max_weight": (0.3, 1.0),
|
||||
"mlp_peak_position": (0.1, 0.9),
|
||||
"mlp_min_weight": (0.0, 0.3),
|
||||
"mlp_spread": (0.1, 0.6),
|
||||
}
|
||||
|
||||
|
||||
def _skip(pipeline: AbliterationPipeline, reason: str, message: str) -> None:
|
||||
"""Record and surface a no-op so it cannot be mistaken for a completed run.
|
||||
|
||||
Reported three ways because each reaches a different audience: the stdlib
|
||||
logger for anyone capturing logs, the pipeline log for the console panel the
|
||||
operator is watching, and an attribute the caller persists into checkpoint
|
||||
metadata so the provenance of an already-written model still shows it.
|
||||
"""
|
||||
logger.warning(message)
|
||||
log = getattr(pipeline, "log", None)
|
||||
if callable(log):
|
||||
log(f"WARNING: {message}")
|
||||
try:
|
||||
pipeline._bayesian_skipped = reason
|
||||
except AttributeError: # pragma: no cover - pipeline stubs may use __slots__
|
||||
pass
|
||||
|
||||
|
||||
def _clamp_to_space(params: dict[str, float]) -> dict[str, float]:
|
||||
"""Clamp warm-start values into the declared search space.
|
||||
|
||||
Optuna rejects an enqueued value outside a parameter's distribution with a
|
||||
warning and does not use it, so an out-of-range warm start silently costs a
|
||||
trial and the informed starting point it was meant to provide. This is
|
||||
reachable in normal use: `attn_peak_position` is seeded from
|
||||
`peak_layer / (n_layers - 1)`, which for an early peak layer falls below the
|
||||
0.1 floor — a 36-layer model peaking at layer 2 seeds 0.057.
|
||||
"""
|
||||
out = dict(params)
|
||||
for name, (lo, hi) in KERNEL_SPACE.items():
|
||||
if name in out:
|
||||
out[name] = min(max(float(out[name]), lo), hi)
|
||||
return out
|
||||
|
||||
|
||||
def _measure_refusal_rate(
|
||||
pipeline: AbliterationPipeline,
|
||||
n_prompts: int = 10,
|
||||
@@ -265,13 +315,31 @@ def run_bayesian_optimization(
|
||||
import optuna
|
||||
from optuna.samplers import TPESampler
|
||||
except ImportError:
|
||||
logger.warning(
|
||||
"Optuna not installed — skipping Bayesian optimization. "
|
||||
"Install with: pip install optuna"
|
||||
# Optuna is not declared in pyproject.toml, requirements.txt or uv.lock,
|
||||
# so a documented `uv sync --locked --extra dev` install does not have it
|
||||
# and `--method optimized` reaches this branch. `logger.warning` alone is
|
||||
# not enough to make that visible: the `obliterate` command never calls
|
||||
# logging.basicConfig, so this surfaces only through logging's
|
||||
# last-resort stderr handler, while the run panel the user actually
|
||||
# reads is fed by `pipeline.log`. The run then completes, reports
|
||||
# success, and writes metadata recording the method that asked for an
|
||||
# optimization which never ran.
|
||||
_skip(
|
||||
pipeline,
|
||||
"optuna-not-installed",
|
||||
"Optuna not installed — Bayesian optimization SKIPPED. Layer "
|
||||
"weights fall back to method defaults; results are NOT optimized. "
|
||||
"Install with: pip install optuna",
|
||||
)
|
||||
return {}
|
||||
|
||||
if not pipeline.handle or not pipeline._strong_layers:
|
||||
_skip(
|
||||
pipeline,
|
||||
"no-strong-layers",
|
||||
"Bayesian optimization SKIPPED: no model handle or no selected "
|
||||
"layers to optimize over.",
|
||||
)
|
||||
return {}
|
||||
|
||||
model = pipeline.handle.model
|
||||
@@ -401,16 +469,16 @@ def run_bayesian_optimization(
|
||||
_restore_all()
|
||||
|
||||
# Attention kernel: 4 params
|
||||
attn_max = trial.suggest_float("attn_max_weight", 0.5, 1.0)
|
||||
attn_peak = trial.suggest_float("attn_peak_position", 0.1, 0.9)
|
||||
attn_min = trial.suggest_float("attn_min_weight", 0.0, 0.3)
|
||||
attn_spread = trial.suggest_float("attn_spread", 0.1, 0.6)
|
||||
attn_max = trial.suggest_float("attn_max_weight", *KERNEL_SPACE["attn_max_weight"])
|
||||
attn_peak = trial.suggest_float("attn_peak_position", *KERNEL_SPACE["attn_peak_position"])
|
||||
attn_min = trial.suggest_float("attn_min_weight", *KERNEL_SPACE["attn_min_weight"])
|
||||
attn_spread = trial.suggest_float("attn_spread", *KERNEL_SPACE["attn_spread"])
|
||||
|
||||
# MLP kernel: 4 params (separate — can peak at a different layer)
|
||||
mlp_max = trial.suggest_float("mlp_max_weight", 0.3, 1.0)
|
||||
mlp_peak = trial.suggest_float("mlp_peak_position", 0.1, 0.9)
|
||||
mlp_min = trial.suggest_float("mlp_min_weight", 0.0, 0.3)
|
||||
mlp_spread = trial.suggest_float("mlp_spread", 0.1, 0.6)
|
||||
mlp_max = trial.suggest_float("mlp_max_weight", *KERNEL_SPACE["mlp_max_weight"])
|
||||
mlp_peak = trial.suggest_float("mlp_peak_position", *KERNEL_SPACE["mlp_peak_position"])
|
||||
mlp_min = trial.suggest_float("mlp_min_weight", *KERNEL_SPACE["mlp_min_weight"])
|
||||
mlp_spread = trial.suggest_float("mlp_spread", *KERNEL_SPACE["mlp_spread"])
|
||||
|
||||
# Float direction index (cross-layer interpolation, Heretic-style)
|
||||
dir_idx = trial.suggest_float("dir_idx", 0.0, max(n_layers_with_dirs - 1, 0.0))
|
||||
@@ -521,7 +589,7 @@ def run_bayesian_optimization(
|
||||
"mlp_spread": 0.3,
|
||||
"dir_idx": 0.0,
|
||||
}
|
||||
study.enqueue_trial(warm_params)
|
||||
study.enqueue_trial(_clamp_to_space(warm_params))
|
||||
|
||||
pipeline.log(f"Bayesian optimization: running {n_trials} trials (parametric kernel)...")
|
||||
study.optimize(objective, n_trials=n_trials, show_progress_bar=False)
|
||||
|
||||
Reference in New Issue
Block a user