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
+100 -2
View File
@@ -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"