From 40236a150e54ac583ba538830379c4a946e69559 Mon Sep 17 00:00:00 2001 From: Joseph Magly <1159087+jmagly@users.noreply.github.com> Date: Mon, 7 Sep 2026 16:21:29 -0400 Subject: [PATCH] fix(bayesian): reset attempt state and persist every skipped outcome --- obliteratus/abliterate.py | 7 +- obliteratus/bayesian_optimizer.py | 43 ++++-- tests/test_bayesian_optimizer_contracts.py | 146 ++++++++++++++++++++- 3 files changed, 176 insertions(+), 20 deletions(-) diff --git a/obliteratus/abliterate.py b/obliteratus/abliterate.py index e19810f..87648fd 100644 --- a/obliteratus/abliterate.py +++ b/obliteratus/abliterate.py @@ -1242,6 +1242,7 @@ class AbliterationPipeline: # Float layer interpolation: continuous layer weights self._float_layer_weights: dict[int, float] = {} # Bayesian optimizer component-specific scales (set by optimizer) + self._bayesian_skipped: str | None = None self._bayesian_attn_scale: float | None = None self._bayesian_mlp_scale: float | None = None # CoT-aware: identified reasoning-critical directions to preserve @@ -1552,6 +1553,8 @@ class AbliterationPipeline: def run(self) -> Path: """Execute the full abliteration pipeline. Returns path to saved model.""" + from obliteratus.bayesian_optimizer import _reset_bayesian_state + _reset_bayesian_state(self) try: return self._run_pipeline() except PipelineFailure: @@ -4209,11 +4212,13 @@ class AbliterationPipeline: # When enabled, run Optuna TPE to find optimal per-layer regularization # before the standard projection loop. The found values override the # static layer_adaptive_strength weights. + from obliteratus.bayesian_optimizer import _reset_bayesian_state + _reset_bayesian_state(self) bayesian_regs: dict[int, float] = {} bayesian_trials = getattr(self, "_bayesian_trials", 0) or ( METHODS.get(self.method, {}).get("bayesian_trials", 0) ) - if bayesian_trials > 0 and self._strong_layers and self.handle: + if bayesian_trials > 0: self.log(f"Running Bayesian optimization ({bayesian_trials} trials)...") from obliteratus.bayesian_optimizer import run_bayesian_optimization bayesian_regs = run_bayesian_optimization( diff --git a/obliteratus/bayesian_optimizer.py b/obliteratus/bayesian_optimizer.py index de9b441..e2115c9 100644 --- a/obliteratus/bayesian_optimizer.py +++ b/obliteratus/bayesian_optimizer.py @@ -56,6 +56,13 @@ KERNEL_SPACE: dict[str, tuple[float, float]] = { } +def _reset_bayesian_state(pipeline: AbliterationPipeline) -> None: + """Clear status and component overrides from an earlier attempt.""" + pipeline._bayesian_skipped = None + pipeline._bayesian_attn_scale = None + pipeline._bayesian_mlp_scale = None + + def _skip(pipeline: AbliterationPipeline, reason: str, message: str) -> None: """Record and surface a no-op so it cannot be mistaken for a completed run. @@ -77,10 +84,10 @@ def _skip(pipeline: AbliterationPipeline, reason: str, message: str) -> None: 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 + Optuna warns about an out-of-range enqueued value but still evaluates it. + Clamping keeps the informed starting point within the same bounds as later + sampled trials. 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. """ @@ -311,19 +318,17 @@ def run_bayesian_optimization( Returns: Dict mapping layer_idx -> optimal regularization value. """ + _reset_bayesian_state(pipeline) + if n_trials <= 0: + _skip(pipeline, "no-trials", "Bayesian optimization SKIPPED: no trials requested.") + return {} + try: import optuna from optuna.samplers import TPESampler except ImportError: - # 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. + # Keep the operator log and saved provenance explicit even when a + # partial or older installation lacks the optimizer dependency. _skip( pipeline, "optuna-not-installed", @@ -375,7 +380,11 @@ def run_bayesian_optimization( pipeline._free_gpu_memory() if not reference_logits: - pipeline.log(" Failed to collect reference logits — skipping optimization") + _skip( + pipeline, + "reference-logits-unavailable", + "Bayesian optimization SKIPPED: Failed to collect reference logits.", + ) return {} from obliteratus.strategies.utils import ( @@ -653,4 +662,10 @@ def run_bayesian_optimization( del original_params pipeline._free_gpu_memory() + if not best_result: + _skip( + pipeline, + "no-successful-trials", + "Bayesian optimization SKIPPED: no usable trial results.", + ) return best_result diff --git a/tests/test_bayesian_optimizer_contracts.py b/tests/test_bayesian_optimizer_contracts.py index 7b7bdef..543375e 100644 --- a/tests/test_bayesian_optimizer_contracts.py +++ b/tests/test_bayesian_optimizer_contracts.py @@ -253,6 +253,7 @@ def test_run_bayesian_optimization_returns_empty_when_reference_logits_fail(monk assert bo.run_bayesian_optimization(pipeline, n_kl_prompts=2) == {} assert pipeline.freed == 1 + assert pipeline._bayesian_skipped == "reference-logits-unavailable" assert "Failed to collect reference logits" in pipeline.logs[-1] @@ -445,9 +446,8 @@ def test_run_bayesian_optimization_no_pareto_uses_objective_best_and_restores_af # 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. + # warned but still evaluated a value outside the search space. Clamping + # keeps the initial trial within the same bounds as sampled trials. assert study.enqueued == [{ "attn_max_weight": 0.9, "attn_peak_position": 0.1, @@ -511,8 +511,8 @@ def test_clamp_to_space_pulls_out_of_range_warm_start_into_the_distribution(): `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. + Optuna warns but evaluates an enqueued value outside its distribution; + the warm-start trial should obey the same bounds as sampled trials. """ clamped = bo._clamp_to_space( { @@ -558,3 +558,139 @@ def test_objective_bounds_and_warm_start_share_one_declaration(monkeypatch): for name, bounds in bo.KERNEL_SPACE.items(): assert recorded.get(name) == bounds, f"{name} suggested outside KERNEL_SPACE" + + +def test_retry_clears_skip_reason_and_component_overrides(monkeypatch): + study = _FakeStudy(best_trials=[_FakeTrial()]) + _install_fake_optuna(monkeypatch, study) + monkeypatch.setattr(bo, "_measure_refusal_rate", lambda *_a, **_k: 0.25) + monkeypatch.setattr(bo, "_measure_kl_divergence", lambda *_a, **_k: 0.1) + pipeline = _optimization_pipeline([_Layer()]) + tokenizer = pipeline.handle.tokenizer + pipeline.handle.tokenizer = lambda *_a, **_k: (_ for _ in ()).throw(RuntimeError("tokenization")) + assert bo.run_bayesian_optimization(pipeline, n_trials=1) == {} + assert pipeline._bayesian_skipped == "reference-logits-unavailable" + pipeline.handle.tokenizer = tokenizer + assert bo.run_bayesian_optimization(pipeline, n_trials=1) + assert pipeline._bayesian_skipped is None + assert pipeline._bayesian_attn_scale == 0.8 + pipeline._strong_layers = [] + assert bo.run_bayesian_optimization(pipeline, n_trials=1) == {} + assert pipeline._bayesian_skipped == "no-strong-layers" + assert pipeline._bayesian_attn_scale is None + assert pipeline._bayesian_mlp_scale is None + + +@pytest.mark.parametrize("n_trials", [0, -1]) +def test_nonpositive_trial_budget_records_skip(n_trials): + pipeline = _optimization_pipeline([_Layer()]) + assert bo.run_bayesian_optimization(pipeline, n_trials=n_trials) == {} + assert pipeline._bayesian_skipped == "no-trials" + + +def test_empty_study_records_skip(monkeypatch): + study = _FakeStudy(best_trials=[]) + _install_fake_optuna(monkeypatch, study) + monkeypatch.setattr(study, "optimize", lambda *_a, **_k: None) + pipeline = _optimization_pipeline([_Layer()]) + assert bo.run_bayesian_optimization(pipeline, n_trials=1) == {} + assert pipeline._bayesian_skipped == "no-successful-trials" + + +@pytest.mark.parametrize("outcome,reason", [ + ("missing", "optuna-not-installed"), + ("reference-failure", "reference-logits-unavailable"), + ("no-layers", "no-strong-layers"), + ("success", None), + ("not-requested", None), +]) +def test_caller_warning_and_saved_checkpoint_status(monkeypatch, tmp_path, outcome, reason): + """Exercise EXCISE and the actual on-disk checkpoint provenance together.""" + import json + from pathlib import Path + from unittest.mock import MagicMock + + from transformers import LlamaConfig, LlamaForCausalLM + + from obliteratus.abliterate import AbliterationPipeline + from obliteratus.models.loader import ModelHandle + + config = LlamaConfig( + vocab_size=16, hidden_size=8, intermediate_size=16, num_hidden_layers=1, + num_attention_heads=1, num_key_value_heads=1, max_position_embeddings=16, + ) + model = LlamaForCausalLM(config).eval() + tokenizer = MagicMock() + tokenizer.return_value = {"input_ids": torch.tensor([[1, 2]])} + tokenizer.save_pretrained.side_effect = lambda path: ( + Path(path) / "tokenizer_config.json" + ).write_text("{}", encoding="utf-8") + pipeline = AbliterationPipeline( + model_name="test-model", output_dir=str(tmp_path / "checkpoint"), method="basic", + ) + pipeline.handle = ModelHandle( + model=model, tokenizer=tokenizer, config=config, model_name="test-model", task="causal_lm", + ) + logs = [] + pipeline._on_log = logs.append + pipeline._on_stage = lambda _record: None + pipeline._strong_layers = [] if outcome == "no-layers" else [0] + pipeline.refusal_directions = {0: torch.ones(8) / 8**0.5} + pipeline.refusal_subspaces = {0: pipeline.refusal_directions[0].unsqueeze(0)} + pipeline._bayesian_trials = 0 if outcome == "not-requested" else 1 + # Both a previous skip and previous successful component overrides must be reset. + pipeline._bayesian_skipped = "previous-attempt" + pipeline._bayesian_attn_scale = 0.2 + pipeline._bayesian_mlp_scale = 0.3 + + study = _FakeStudy(best_trials=[_FakeTrial(params={**_FakeTrial.params, "dir_idx": 0.0})]) + _install_fake_optuna(monkeypatch, study) + monkeypatch.setattr(bo, "_measure_refusal_rate", lambda *_a, **_k: 0.25) + monkeypatch.setattr(bo, "_measure_kl_divergence", lambda *_a, **_k: 0.1) + if outcome == "missing": + real_import = builtins.__import__ + + def import_without_optuna(name, *args, **kwargs): + if name == "optuna" or name.startswith("optuna."): + raise ImportError("missing optuna") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", import_without_optuna) + elif outcome == "reference-failure": + tokenizer.side_effect = RuntimeError("tokenization failed") + + pipeline._excise() + result = pipeline._rebirth() + metadata = json.loads((result / "abliteration_metadata.json").read_text()) + assert metadata["method_config"]["bayesian_optimization_skipped"] == reason + fallback = [message for message in logs if "this checkpoint is NOT optimized" in message] + if reason: + assert len(fallback) == 1 + assert reason in fallback[0] + else: + assert not fallback + if outcome != "success": + assert pipeline._bayesian_attn_scale is None + assert pipeline._bayesian_mlp_scale is None + if outcome == "success": + assert any("Bayesian optimization complete" in message for message in logs) + if outcome == "not-requested": + assert not study.enqueued + + +def test_run_resets_bayesian_status_before_any_stage(monkeypatch, tmp_path): + from obliteratus.abliterate import AbliterationPipeline + + pipeline = AbliterationPipeline(model_name="test-model", output_dir=str(tmp_path), method="basic") + pipeline._bayesian_skipped = "previous-attempt" + pipeline._bayesian_attn_scale = 0.2 + pipeline._bayesian_mlp_scale = 0.3 + + def stages(): + assert pipeline._bayesian_skipped is None + assert pipeline._bayesian_attn_scale is None + assert pipeline._bayesian_mlp_scale is None + return tmp_path + + monkeypatch.setattr(pipeline, "_run_pipeline", stages) + assert pipeline.run() == tmp_path