fix(bayesian): surface skipped optimization and clamp warm start to search space

This commit is contained in:
Raja Mukerji
2026-09-07 00:36:39 -07:00
parent 7c332b0e4e
commit 2d31d64857
3 changed files with 196 additions and 14 deletions
+80 -12
View File
@@ -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)