mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-17 16:37:30 +02:00
test: enforce gate 2 coverage and vertical contracts
This commit is contained in:
Joseph Magly
@@ -0,0 +1,462 @@
|
||||
"""CPU-safe contract tests for Bayesian optimization helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import builtins
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from obliteratus import bayesian_optimizer as bo
|
||||
|
||||
|
||||
class _Pipeline:
|
||||
def __init__(self):
|
||||
self.refusal_directions = {}
|
||||
self.handle = None
|
||||
self._strong_layers = []
|
||||
self.harmful_prompts = ["harmful one", "harmful two"]
|
||||
self.use_chat_template = False
|
||||
self.freed = 0
|
||||
self.logs = []
|
||||
|
||||
def _get_model_device(self, _model):
|
||||
return torch.device("cpu")
|
||||
|
||||
def _maybe_apply_chat_template(self, prompts):
|
||||
return [f"<chat>{prompt}</chat>" for prompt in prompts]
|
||||
|
||||
def _free_gpu_memory(self):
|
||||
self.freed += 1
|
||||
|
||||
def log(self, message):
|
||||
self.logs.append(message)
|
||||
|
||||
|
||||
def _install_fake_optuna(monkeypatch, study=None):
|
||||
optuna = types.ModuleType("optuna")
|
||||
optuna.logging = types.SimpleNamespace(WARNING=30, set_verbosity=lambda _level: None)
|
||||
optuna.create_study = lambda **_kwargs: study
|
||||
|
||||
samplers = types.ModuleType("optuna.samplers")
|
||||
|
||||
class TPESampler:
|
||||
def __init__(self, **_kwargs):
|
||||
pass
|
||||
|
||||
samplers.TPESampler = TPESampler
|
||||
monkeypatch.setitem(sys.modules, "optuna", optuna)
|
||||
monkeypatch.setitem(sys.modules, "optuna.samplers", samplers)
|
||||
|
||||
|
||||
class _FakeTrial:
|
||||
number = 0
|
||||
|
||||
params = {
|
||||
"attn_max_weight": 0.8,
|
||||
"attn_peak_position": 0.0,
|
||||
"attn_min_weight": 0.1,
|
||||
"attn_spread": 0.6,
|
||||
"mlp_max_weight": 0.6,
|
||||
"mlp_peak_position": 1.0,
|
||||
"mlp_min_weight": 0.2,
|
||||
"mlp_spread": 0.6,
|
||||
"dir_idx": 1.0,
|
||||
}
|
||||
|
||||
values = (0.2, 0.1)
|
||||
|
||||
def __init__(self, number=0, params=None, values=None):
|
||||
self.number = number
|
||||
if params is not None:
|
||||
self.params = params
|
||||
if values is not None:
|
||||
self.values = values
|
||||
|
||||
def suggest_float(self, name, _low, _high):
|
||||
return self.params[name]
|
||||
|
||||
|
||||
class _FakeStudy:
|
||||
def __init__(self, best_trials):
|
||||
self.best_trials = best_trials
|
||||
self.enqueued = []
|
||||
self.objective_values = []
|
||||
|
||||
def enqueue_trial(self, params):
|
||||
self.enqueued.append(params)
|
||||
|
||||
def optimize(self, objective, n_trials, show_progress_bar):
|
||||
assert show_progress_bar is False
|
||||
for number in range(n_trials):
|
||||
self.objective_values.append(objective(_FakeTrial(number=number)))
|
||||
|
||||
|
||||
class _TokenBatch(dict):
|
||||
def __init__(self):
|
||||
super().__init__(input_ids=torch.tensor([[1, 2]]))
|
||||
|
||||
|
||||
class _ReferenceTokenizer:
|
||||
def __call__(self, *_args, **_kwargs):
|
||||
return _TokenBatch()
|
||||
|
||||
|
||||
class _Layer(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.self_attn = nn.Module()
|
||||
self.self_attn.o_proj = nn.Linear(2, 2)
|
||||
self.mlp = nn.Module()
|
||||
self.mlp.down_proj = nn.Linear(2, 2)
|
||||
|
||||
|
||||
class _ReferenceModel(nn.Module):
|
||||
def __init__(self, layers):
|
||||
super().__init__()
|
||||
self.model = nn.Module()
|
||||
self.model.layers = nn.ModuleList(layers)
|
||||
|
||||
def forward(self, **_kwargs):
|
||||
return types.SimpleNamespace(logits=torch.tensor([[[0.0, 1.0]]]))
|
||||
|
||||
|
||||
def _optimization_pipeline(layers):
|
||||
pipeline = _Pipeline()
|
||||
pipeline.handle = types.SimpleNamespace(
|
||||
model=_ReferenceModel(layers),
|
||||
tokenizer=_ReferenceTokenizer(),
|
||||
architecture="llama",
|
||||
)
|
||||
pipeline._strong_layers = list(range(len(layers)))
|
||||
pipeline.refusal_directions = {
|
||||
idx: torch.tensor([float(idx + 1), 1.0])
|
||||
for idx in pipeline._strong_layers
|
||||
}
|
||||
pipeline.norm_preserve = True
|
||||
pipeline.projections = []
|
||||
pipeline.moe_calls = []
|
||||
|
||||
def project_out(module, direction, names, norm_preserve, regularization):
|
||||
pipeline.projections.append(
|
||||
{
|
||||
"module": module,
|
||||
"direction": direction.detach().clone(),
|
||||
"names": tuple(names),
|
||||
"norm_preserve": norm_preserve,
|
||||
"regularization": regularization,
|
||||
}
|
||||
)
|
||||
for name in names:
|
||||
proj = getattr(module, name, None)
|
||||
if proj is not None and hasattr(proj, "weight"):
|
||||
proj.weight.data.add_(10.0)
|
||||
return 1
|
||||
return 0
|
||||
|
||||
def project_moe(module, direction, **kwargs):
|
||||
pipeline.moe_calls.append((module, direction.detach().clone(), kwargs))
|
||||
|
||||
pipeline._project_out_advanced = project_out
|
||||
pipeline._project_moe_experts = project_moe
|
||||
return pipeline
|
||||
|
||||
|
||||
def test_parametric_layer_weight_boundaries():
|
||||
assert bo._parametric_layer_weight(0, 1, 0.8, 0.5, 0.1, 0.2) == pytest.approx(0.8)
|
||||
|
||||
# At the peak, the kernel returns the maximum weight.
|
||||
assert bo._parametric_layer_weight(2, 5, 0.9, 0.5, 0.1, 0.25) == pytest.approx(0.9)
|
||||
|
||||
# At the tent edge, it reaches the minimum weight.
|
||||
assert bo._parametric_layer_weight(1, 5, 0.9, 0.5, 0.1, 0.25) == pytest.approx(0.1)
|
||||
|
||||
# Outside the spread cutoff, the layer is skipped.
|
||||
assert bo._parametric_layer_weight(0, 5, 0.9, 0.5, 0.1, 0.24) == pytest.approx(0.0)
|
||||
|
||||
# Tiny or negative spread is clamped to 0.01.
|
||||
assert bo._parametric_layer_weight(0, 101, 0.7, 0.0, 0.2, -1.0) == pytest.approx(0.7)
|
||||
|
||||
|
||||
def test_interpolate_direction_handles_empty_clamps_exact_and_normalized_interpolation():
|
||||
pipeline = _Pipeline()
|
||||
assert torch.equal(bo._interpolate_direction(pipeline, layer_idx=3, float_dir_idx=1.0), torch.zeros(1))
|
||||
|
||||
pipeline.refusal_directions = {
|
||||
2: torch.tensor([3.0, 0.0]),
|
||||
5: torch.tensor([0.0, 4.0]),
|
||||
9: torch.tensor([1.0, 1.0]),
|
||||
}
|
||||
|
||||
low = bo._interpolate_direction(pipeline, layer_idx=5, float_dir_idx=-10.0)
|
||||
assert torch.allclose(low, torch.tensor([1.0, 0.0]))
|
||||
|
||||
high = bo._interpolate_direction(pipeline, layer_idx=5, float_dir_idx=99.0)
|
||||
assert torch.allclose(high, torch.tensor([2**-0.5, 2**-0.5]))
|
||||
|
||||
exact = bo._interpolate_direction(pipeline, layer_idx=5, float_dir_idx=1.0)
|
||||
assert torch.allclose(exact, torch.tensor([0.0, 1.0]))
|
||||
|
||||
interpolated = bo._interpolate_direction(pipeline, layer_idx=5, float_dir_idx=0.5)
|
||||
expected = torch.tensor([1.5, 2.0])
|
||||
expected = expected / expected.norm()
|
||||
assert torch.allclose(interpolated, expected)
|
||||
assert interpolated.norm().item() == pytest.approx(1.0)
|
||||
|
||||
|
||||
def test_run_bayesian_optimization_returns_empty_when_optuna_missing(monkeypatch):
|
||||
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) == {}
|
||||
|
||||
|
||||
def test_run_bayesian_optimization_returns_empty_without_handle_or_layers(monkeypatch):
|
||||
_install_fake_optuna(monkeypatch)
|
||||
|
||||
pipeline = _Pipeline()
|
||||
pipeline.handle = None
|
||||
pipeline._strong_layers = [0]
|
||||
assert bo.run_bayesian_optimization(pipeline) == {}
|
||||
|
||||
pipeline.handle = types.SimpleNamespace(model=object(), tokenizer=object())
|
||||
pipeline._strong_layers = []
|
||||
assert bo.run_bayesian_optimization(pipeline) == {}
|
||||
|
||||
|
||||
def test_run_bayesian_optimization_returns_empty_when_reference_logits_fail(monkeypatch):
|
||||
_install_fake_optuna(monkeypatch)
|
||||
|
||||
class RaisingTokenizer:
|
||||
def __call__(self, *_args, **_kwargs):
|
||||
raise RuntimeError("tokenization failed")
|
||||
|
||||
pipeline = _Pipeline()
|
||||
pipeline.handle = types.SimpleNamespace(
|
||||
model=object(),
|
||||
tokenizer=RaisingTokenizer(),
|
||||
)
|
||||
pipeline._strong_layers = [0]
|
||||
|
||||
assert bo.run_bayesian_optimization(pipeline, n_kl_prompts=2) == {}
|
||||
assert pipeline.freed == 1
|
||||
assert "Failed to collect reference logits" in pipeline.logs[-1]
|
||||
|
||||
|
||||
def test_measure_refusal_rate_restores_chat_template_and_counts_generation_failures():
|
||||
class TokenBatch(dict):
|
||||
def __init__(self, token_count):
|
||||
super().__init__(input_ids=torch.arange(token_count).unsqueeze(0))
|
||||
|
||||
class Tokenizer:
|
||||
def __init__(self):
|
||||
self.calls = 0
|
||||
|
||||
def __call__(self, prompt, **_kwargs):
|
||||
self.calls += 1
|
||||
return TokenBatch(3 if "one" in prompt else 4)
|
||||
|
||||
def decode(self, token_ids, **_kwargs):
|
||||
if token_ids.numel() == 0:
|
||||
return ""
|
||||
return "I cannot assist with that request"
|
||||
|
||||
class Model:
|
||||
def __init__(self):
|
||||
self.calls = 0
|
||||
|
||||
def generate(self, **kwargs):
|
||||
self.calls += 1
|
||||
if self.calls == 2:
|
||||
raise RuntimeError("generation failed")
|
||||
input_ids = kwargs["input_ids"]
|
||||
continuation = torch.tensor([[99, 100]])
|
||||
return torch.cat([input_ids, continuation], dim=1)
|
||||
|
||||
pipeline = _Pipeline()
|
||||
pipeline.use_chat_template = False
|
||||
pipeline.handle = types.SimpleNamespace(model=Model(), tokenizer=Tokenizer())
|
||||
|
||||
rate = bo._measure_refusal_rate(pipeline, n_prompts=2, max_new_tokens=4)
|
||||
|
||||
assert rate == pytest.approx(0.5)
|
||||
assert pipeline.use_chat_template is False
|
||||
assert pipeline.freed == 1
|
||||
|
||||
|
||||
def test_measure_kl_divergence_skips_failures_and_missing_reference_logits():
|
||||
class TokenBatch(dict):
|
||||
def __init__(self):
|
||||
super().__init__(input_ids=torch.tensor([[1, 2]]))
|
||||
|
||||
class Tokenizer:
|
||||
def __call__(self, prompt, **_kwargs):
|
||||
if "bad" in prompt:
|
||||
raise RuntimeError("tokenization failed")
|
||||
return TokenBatch()
|
||||
|
||||
class Output:
|
||||
logits = torch.tensor([[[0.0, 1.0, 2.0]]])
|
||||
|
||||
class Model:
|
||||
def __call__(self, **_kwargs):
|
||||
return Output()
|
||||
|
||||
pipeline = _Pipeline()
|
||||
pipeline.handle = types.SimpleNamespace(model=Model(), tokenizer=Tokenizer())
|
||||
reference = [torch.tensor([0.0, 1.0, 2.0])]
|
||||
|
||||
kl = bo._measure_kl_divergence(
|
||||
pipeline,
|
||||
reference_logits=reference,
|
||||
prompts=["good", "bad", "ignored because no reference"],
|
||||
)
|
||||
|
||||
assert kl == pytest.approx(0.0)
|
||||
assert pipeline.freed == 1
|
||||
|
||||
|
||||
def test_run_bayesian_optimization_pareto_path_enqueues_warm_start_and_restores(monkeypatch):
|
||||
layers = [_Layer(), _Layer()]
|
||||
originals = [
|
||||
layer.self_attn.o_proj.weight.detach().clone()
|
||||
for layer in layers
|
||||
] + [
|
||||
layer.mlp.down_proj.weight.detach().clone()
|
||||
for layer in layers
|
||||
]
|
||||
best_params = {
|
||||
"attn_max_weight": 0.8,
|
||||
"attn_peak_position": 0.0,
|
||||
"attn_min_weight": 0.1,
|
||||
"attn_spread": 0.6,
|
||||
"mlp_max_weight": 0.6,
|
||||
"mlp_peak_position": 1.0,
|
||||
"mlp_min_weight": 0.2,
|
||||
"mlp_spread": 0.6,
|
||||
"dir_idx": 1.0,
|
||||
}
|
||||
worse_params = {**best_params, "dir_idx": 0.0}
|
||||
study = _FakeStudy(
|
||||
best_trials=[
|
||||
_FakeTrial(params=worse_params, values=(0.5, 0.1)),
|
||||
_FakeTrial(params=best_params, values=(0.2, 0.3)),
|
||||
]
|
||||
)
|
||||
_install_fake_optuna(monkeypatch, study)
|
||||
monkeypatch.setattr(bo, "_measure_refusal_rate", lambda *_args, **_kwargs: 0.25)
|
||||
monkeypatch.setattr(bo, "_measure_kl_divergence", lambda *_args, **_kwargs: 0.1)
|
||||
|
||||
pipeline = _optimization_pipeline(layers)
|
||||
pipeline._informed_warm_start = {
|
||||
"max_weight": 0.8,
|
||||
"peak_position": 0.25,
|
||||
"min_weight": 0.02,
|
||||
"spread": 0.2,
|
||||
"mlp_scale": 0.5,
|
||||
"dir_idx": 0.75,
|
||||
}
|
||||
|
||||
result = bo.run_bayesian_optimization(
|
||||
pipeline,
|
||||
n_trials=2,
|
||||
n_refusal_prompts=3,
|
||||
n_kl_prompts=1,
|
||||
)
|
||||
|
||||
assert study.enqueued == [{
|
||||
"attn_max_weight": 0.8,
|
||||
"attn_peak_position": 0.25,
|
||||
"attn_min_weight": 0.02,
|
||||
"attn_spread": 0.2,
|
||||
"mlp_max_weight": 0.4,
|
||||
"mlp_peak_position": 0.25,
|
||||
"mlp_min_weight": 0.02,
|
||||
"mlp_spread": 0.2,
|
||||
"dir_idx": 0.75,
|
||||
}]
|
||||
assert study.objective_values == [(0.25, 0.1), (0.25, 0.1)]
|
||||
assert result == {0: pytest.approx(0.6), 1: pytest.approx(0.7)}
|
||||
assert pipeline._bayesian_attn_scale == pytest.approx(0.8)
|
||||
assert pipeline._bayesian_mlp_scale == pytest.approx(0.6)
|
||||
assert any("Applying interpolated direction" in msg for msg in pipeline.logs)
|
||||
|
||||
restored = [
|
||||
layer.self_attn.o_proj.weight
|
||||
for layer in layers
|
||||
] + [
|
||||
layer.mlp.down_proj.weight
|
||||
for layer in layers
|
||||
]
|
||||
for live, original in zip(restored, originals, strict=True):
|
||||
assert torch.allclose(live, original)
|
||||
|
||||
assert len(pipeline.projections) == 8
|
||||
regularizations = sorted({round(call["regularization"], 6) for call in pipeline.projections})
|
||||
assert regularizations == [0.2, 0.4, 1.0]
|
||||
assert all(call["norm_preserve"] is True for call in pipeline.projections)
|
||||
assert all(call["direction"].shape == (2, 1) for call in pipeline.projections)
|
||||
|
||||
expected_direction = torch.tensor([2.0, 1.0])
|
||||
expected_direction = expected_direction / expected_direction.norm()
|
||||
assert torch.allclose(pipeline.refusal_directions[0], expected_direction)
|
||||
assert torch.allclose(pipeline.refusal_directions[1], expected_direction)
|
||||
|
||||
|
||||
def test_run_bayesian_optimization_no_pareto_uses_objective_best_and_restores_after_projection_errors(
|
||||
monkeypatch,
|
||||
):
|
||||
layers = [_Layer()]
|
||||
original_attn = layers[0].self_attn.o_proj.weight.detach().clone()
|
||||
original_mlp = layers[0].mlp.down_proj.weight.detach().clone()
|
||||
study = _FakeStudy(best_trials=[])
|
||||
_install_fake_optuna(monkeypatch, study)
|
||||
monkeypatch.setattr(bo, "_measure_refusal_rate", lambda *_args, **_kwargs: 0.4)
|
||||
monkeypatch.setattr(bo, "_measure_kl_divergence", lambda *_args, **_kwargs: 0.2)
|
||||
|
||||
pipeline = _optimization_pipeline(layers)
|
||||
|
||||
def raising_project(*_args, **_kwargs):
|
||||
layers[0].self_attn.o_proj.weight.data.add_(5.0)
|
||||
layers[0].mlp.down_proj.weight.data.add_(7.0)
|
||||
raise RuntimeError("projection failed")
|
||||
|
||||
pipeline._project_out_advanced = raising_project
|
||||
|
||||
result = bo.run_bayesian_optimization(
|
||||
pipeline,
|
||||
n_trials=1,
|
||||
n_refusal_prompts=1,
|
||||
n_kl_prompts=1,
|
||||
)
|
||||
|
||||
assert study.enqueued == [{
|
||||
"attn_max_weight": 0.9,
|
||||
"attn_peak_position": 0.0,
|
||||
"attn_min_weight": 0.05,
|
||||
"attn_spread": 0.3,
|
||||
"mlp_max_weight": 0.6,
|
||||
"mlp_peak_position": 0.0,
|
||||
"mlp_min_weight": 0.05,
|
||||
"mlp_spread": 0.3,
|
||||
"dir_idx": 0.0,
|
||||
}]
|
||||
assert study.objective_values == [(0.4, 0.2)]
|
||||
assert result == {0: pytest.approx(0.3)}
|
||||
assert any("Using best combined score: 0.5000" in msg for msg in pipeline.logs)
|
||||
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)
|
||||
Reference in New Issue
Block a user