mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-18 00:47:23 +02:00
463 lines
15 KiB
Python
463 lines
15 KiB
Python
"""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)
|