From 2d31d64857701644ea7be08cfa5926fc0d8dbeff Mon Sep 17 00:00:00 2001 From: Raja Mukerji Date: Mon, 7 Sep 2026 00:36:39 -0700 Subject: [PATCH] fix(bayesian): surface skipped optimization and clamp warm start to search space --- obliteratus/abliterate.py | 16 ++++ obliteratus/bayesian_optimizer.py | 92 ++++++++++++++++--- tests/test_bayesian_optimizer_contracts.py | 102 ++++++++++++++++++++- 3 files changed, 196 insertions(+), 14 deletions(-) diff --git a/obliteratus/abliterate.py b/obliteratus/abliterate.py index 6daaff9..e19810f 100644 --- a/obliteratus/abliterate.py +++ b/obliteratus/abliterate.py @@ -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, diff --git a/obliteratus/bayesian_optimizer.py b/obliteratus/bayesian_optimizer.py index 8d17816..de9b441 100644 --- a/obliteratus/bayesian_optimizer.py +++ b/obliteratus/bayesian_optimizer.py @@ -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) diff --git a/tests/test_bayesian_optimizer_contracts.py b/tests/test_bayesian_optimizer_contracts.py index dfddb16..7b7bdef 100644 --- a/tests/test_bayesian_optimizer_contracts.py +++ b/tests/test_bayesian_optimizer_contracts.py @@ -443,13 +443,18 @@ def test_run_bayesian_optimization_no_pareto_uses_objective_best_and_restores_af n_kl_prompts=1, ) + # peak_position was 0.0 here before the warm start was clamped. 0.0 is + # outside the distribution `objective` declares for it (0.1..0.9), so Optuna + # rejected the enqueued trial with a warning and the informed warm start was + # silently discarded. This assertion previously encoded that defect as the + # expected behaviour, which is why it went unnoticed; 0.1 is the floor. assert study.enqueued == [{ "attn_max_weight": 0.9, - "attn_peak_position": 0.0, + "attn_peak_position": 0.1, "attn_min_weight": 0.05, "attn_spread": 0.3, "mlp_max_weight": 0.6, - "mlp_peak_position": 0.0, + "mlp_peak_position": 0.1, "mlp_min_weight": 0.05, "mlp_spread": 0.3, "dir_idx": 0.0, @@ -460,3 +465,96 @@ def test_run_bayesian_optimization_no_pareto_uses_objective_best_and_restores_af assert pipeline.moe_calls == [] assert torch.allclose(layers[0].self_attn.o_proj.weight, original_attn) assert torch.allclose(layers[0].mlp.down_proj.weight, original_mlp) + + +def test_missing_optuna_is_surfaced_on_the_pipeline_log_and_recorded(monkeypatch): + """A skipped optimization must not be indistinguishable from a completed one. + + `logger.warning` reaches stderr only through logging's last-resort handler — + the `obliterate` command never configures logging — while the console panel + the operator watches is fed by `pipeline.log`. Before this contract the + pipeline log received nothing at all, so a run whose headline feature never + executed still read as a clean success. + """ + real_import = builtins.__import__ + + def fake_import(name, *args, **kwargs): + if name == "optuna" or name.startswith("optuna."): + raise ImportError("no optuna in this test") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", fake_import) + + pipeline = _Pipeline() + pipeline.handle = object() + pipeline._strong_layers = [0] + + assert bo.run_bayesian_optimization(pipeline) == {} + assert any("SKIPPED" in message for message in pipeline.logs) + assert any("optuna" in message.lower() for message in pipeline.logs) + assert pipeline._bayesian_skipped == "optuna-not-installed" + + +def test_skip_without_handle_or_layers_is_also_recorded(monkeypatch): + _install_fake_optuna(monkeypatch) + + pipeline = _Pipeline() + pipeline.handle = None + pipeline._strong_layers = [0] + + assert bo.run_bayesian_optimization(pipeline) == {} + assert pipeline._bayesian_skipped == "no-strong-layers" + + +def test_clamp_to_space_pulls_out_of_range_warm_start_into_the_distribution(): + """Reachable in normal use, not a synthetic edge case. + + `attn_peak_position` is seeded from ``peak_layer / (n_layers - 1)``; a + 36-layer model whose peak layer is 2 seeds 0.057, below the 0.1 floor. + Optuna rejects an enqueued value outside its distribution, so the warm-start + trial silently loses the informed starting point it exists to provide. + """ + clamped = bo._clamp_to_space( + { + "attn_peak_position": 2 / 35, # early peak layer on a 36-layer model + "mlp_peak_position": 34 / 35, # late peak layer + "mlp_max_weight": 0.25, # max_weight 0.5 x mlp_scale 0.5 + "dir_idx": 3.0, # bounds are dynamic; must pass through + } + ) + + for name, value in clamped.items(): + if name in bo.KERNEL_SPACE: + low, high = bo.KERNEL_SPACE[name] + assert low <= value <= high, f"{name}={value} outside {low}..{high}" + + assert clamped["attn_peak_position"] == pytest.approx(0.1) + assert clamped["mlp_peak_position"] == pytest.approx(0.9) + assert clamped["mlp_max_weight"] == pytest.approx(0.3) + assert clamped["dir_idx"] == 3.0 + + +def test_objective_bounds_and_warm_start_share_one_declaration(monkeypatch): + """The two copies of the search space are why the warm start could drift.""" + study = _FakeStudy(best_trials=[_FakeTrial()]) + _install_fake_optuna(monkeypatch, study=study) + + recorded: dict[str, tuple[float, float]] = {} + + class _RecordingTrial(_FakeTrial): + def suggest_float(self, name, low, high): + recorded[name] = (low, high) + return self.params[name] + + monkeypatch.setattr( + study, + "optimize", + lambda objective, n_trials, show_progress_bar: objective(_RecordingTrial()), + ) + monkeypatch.setattr(bo, "_measure_refusal_rate", lambda *_a, **_k: 0.25) + monkeypatch.setattr(bo, "_measure_kl_divergence", lambda *_a, **_k: 0.1) + + bo.run_bayesian_optimization(_optimization_pipeline([_Layer()]), n_trials=1) + + for name, bounds in bo.KERNEL_SPACE.items(): + assert recorded.get(name) == bounds, f"{name} suggested outside KERNEL_SPACE"