mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-17 16:37:30 +02:00
716 lines
24 KiB
Python
716 lines
24 KiB
Python
"""Deterministic contracts for analysis-informed orchestration boundaries."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from obliteratus.informed_pipeline import InformedAbliterationPipeline
|
|
|
|
|
|
@pytest.fixture
|
|
def pipeline(tmp_path):
|
|
return InformedAbliterationPipeline(
|
|
model_name="fixture/model",
|
|
output_dir=str(tmp_path / "output"),
|
|
on_log=lambda _message: None,
|
|
)
|
|
|
|
|
|
def test_run_informed_executes_the_documented_stage_order(pipeline, monkeypatch):
|
|
calls = []
|
|
output = pipeline.output_dir
|
|
for name in (
|
|
"_summon",
|
|
"_probe",
|
|
"_analyze",
|
|
"_distill_informed",
|
|
"_excise_informed",
|
|
"_verify_and_compensate",
|
|
):
|
|
monkeypatch.setattr(pipeline, name, lambda name=name: calls.append(name))
|
|
monkeypatch.setattr(
|
|
pipeline,
|
|
"_rebirth_informed",
|
|
lambda: calls.append("_rebirth_informed") or output,
|
|
)
|
|
ticks = iter((10.0, 12.5))
|
|
monkeypatch.setattr("obliteratus.informed_pipeline.time.time", lambda: next(ticks))
|
|
|
|
result, report = pipeline.run_informed()
|
|
|
|
assert result == output
|
|
assert calls == [
|
|
"_summon",
|
|
"_probe",
|
|
"_analyze",
|
|
"_distill_informed",
|
|
"_excise_informed",
|
|
"_verify_and_compensate",
|
|
"_rebirth_informed",
|
|
]
|
|
assert report.total_duration == 2.5
|
|
|
|
|
|
def test_analyze_runs_only_enabled_modules_and_records_duration(pipeline, monkeypatch):
|
|
pipeline._run_cone = False
|
|
pipeline._run_defense = False
|
|
calls = []
|
|
events = []
|
|
monkeypatch.setattr(pipeline, "_analyze_alignment_imprint", lambda: calls.append("alignment"))
|
|
monkeypatch.setattr(pipeline, "_analyze_cone_geometry", lambda: calls.append("cone"))
|
|
monkeypatch.setattr(pipeline, "_analyze_cross_layer", lambda: calls.append("cross-layer"))
|
|
monkeypatch.setattr(pipeline, "_analyze_defense_robustness", lambda: calls.append("defense"))
|
|
monkeypatch.setattr(pipeline, "_analyze_sparsity", lambda: calls.append("sparsity"))
|
|
monkeypatch.setattr(pipeline, "_derive_configuration", lambda: calls.append("derive"))
|
|
monkeypatch.setattr(pipeline, "_emit", lambda *args, **kwargs: events.append((args, kwargs)))
|
|
ticks = iter((20.0, 21.25))
|
|
monkeypatch.setattr("obliteratus.informed_pipeline.time.time", lambda: next(ticks))
|
|
|
|
pipeline._analyze()
|
|
|
|
assert calls == ["alignment", "cross-layer", "sparsity", "derive"]
|
|
assert pipeline._report.analysis_duration == 1.25
|
|
assert events[0][0][:2] == ("analyze", "running")
|
|
assert events[-1][0][:2] == ("analyze", "done")
|
|
assert events[-1][1]["duration"] == 1.25
|
|
|
|
|
|
def test_alignment_imprint_skips_zero_directions(pipeline, monkeypatch):
|
|
detector_calls = []
|
|
|
|
class Detector:
|
|
def detect_imprint(self, directions):
|
|
detector_calls.append(directions)
|
|
|
|
monkeypatch.setattr(
|
|
"obliteratus.analysis.alignment_imprint.AlignmentImprintDetector",
|
|
Detector,
|
|
)
|
|
pipeline._harmful_means = {0: torch.ones(1, 3)}
|
|
pipeline._harmless_means = {0: torch.ones(1, 3)}
|
|
|
|
pipeline._analyze_alignment_imprint()
|
|
|
|
assert detector_calls == []
|
|
assert pipeline._insights.detected_alignment_method == "unknown"
|
|
|
|
|
|
def test_alignment_imprint_normalizes_directions_and_copies_result(pipeline, monkeypatch):
|
|
captured = {}
|
|
result = SimpleNamespace(
|
|
predicted_method="dpo",
|
|
confidence=0.8,
|
|
dpo_probability=0.8,
|
|
rlhf_probability=0.1,
|
|
cai_probability=0.05,
|
|
sft_probability=0.05,
|
|
gini_coefficient=0.3,
|
|
effective_rank=2.0,
|
|
cross_layer_smoothness=0.7,
|
|
tail_layer_bias=0.2,
|
|
)
|
|
|
|
class Detector:
|
|
def detect_imprint(self, directions):
|
|
captured.update(directions)
|
|
return result
|
|
|
|
monkeypatch.setattr(
|
|
"obliteratus.analysis.alignment_imprint.AlignmentImprintDetector",
|
|
Detector,
|
|
)
|
|
pipeline._harmful_means = {1: torch.tensor([[3.0, 4.0]])}
|
|
pipeline._harmless_means = {1: torch.zeros(1, 2)}
|
|
|
|
pipeline._analyze_alignment_imprint()
|
|
|
|
assert torch.allclose(captured[1], torch.tensor([0.6, 0.8]))
|
|
assert pipeline._insights.detected_alignment_method == "dpo"
|
|
assert pipeline._insights.alignment_confidence == 0.8
|
|
assert pipeline._insights.alignment_probabilities == {
|
|
"dpo": 0.8,
|
|
"rlhf": 0.1,
|
|
"cai": 0.05,
|
|
"sft": 0.05,
|
|
}
|
|
|
|
|
|
def test_cone_geometry_aggregates_layers_and_keeps_strongest_categories(
|
|
pipeline,
|
|
monkeypatch,
|
|
):
|
|
category = SimpleNamespace(
|
|
category="safety",
|
|
direction=torch.tensor([1.0, 0.0]),
|
|
specificity=0.7,
|
|
strength=2.0,
|
|
)
|
|
results = iter(
|
|
(
|
|
SimpleNamespace(
|
|
is_polyhedral=True,
|
|
cone_dimensionality=3.0,
|
|
mean_pairwise_cosine=0.2,
|
|
general_direction=torch.tensor([2.0, 0.0]),
|
|
category_directions=[category],
|
|
category_count=1,
|
|
),
|
|
SimpleNamespace(
|
|
is_polyhedral=True,
|
|
cone_dimensionality=5.0,
|
|
mean_pairwise_cosine=0.4,
|
|
general_direction=torch.tensor([1.0, 0.0]),
|
|
category_directions=[],
|
|
category_count=0,
|
|
),
|
|
)
|
|
)
|
|
|
|
class Analyzer:
|
|
def analyze_layer(self, *_args, **_kwargs):
|
|
return next(results)
|
|
|
|
monkeypatch.setattr(
|
|
"obliteratus.analysis.concept_geometry.ConceptConeAnalyzer",
|
|
Analyzer,
|
|
)
|
|
pipeline._harmful_acts = {idx: [torch.ones(1, 2)] for idx in range(4)}
|
|
pipeline._harmless_acts = {idx: [torch.zeros(1, 2)] for idx in range(4)}
|
|
|
|
pipeline._analyze_cone_geometry()
|
|
|
|
assert pipeline._insights.cone_is_polyhedral is True
|
|
assert pipeline._insights.cone_dimensionality == 4.0
|
|
assert pipeline._insights.mean_pairwise_cosine == pytest.approx(0.3)
|
|
assert torch.equal(
|
|
pipeline._insights.per_category_directions["safety"],
|
|
torch.tensor([1.0, 0.0]),
|
|
)
|
|
assert pipeline._insights.direction_specificity == {"safety": 0.7}
|
|
|
|
|
|
def test_cone_geometry_preserves_defaults_when_no_layers_are_eligible(pipeline, monkeypatch):
|
|
class Analyzer:
|
|
def analyze_layer(self, *_args, **_kwargs):
|
|
raise AssertionError("no layer should be analyzed")
|
|
|
|
monkeypatch.setattr(
|
|
"obliteratus.analysis.concept_geometry.ConceptConeAnalyzer",
|
|
Analyzer,
|
|
)
|
|
pipeline._harmful_acts = {0: [torch.ones(1, 2)]}
|
|
pipeline._harmless_acts = {}
|
|
|
|
pipeline._analyze_cone_geometry()
|
|
|
|
assert pipeline._insights.cone_is_polyhedral is False
|
|
assert pipeline._insights.cone_dimensionality == 1.0
|
|
|
|
|
|
def test_cross_layer_analysis_selects_the_strongest_cluster_representatives(
|
|
pipeline,
|
|
monkeypatch,
|
|
):
|
|
result = SimpleNamespace(
|
|
clusters=[[0, 1], [2]],
|
|
cluster_count=2,
|
|
direction_persistence_score=0.75,
|
|
mean_adjacent_cosine=0.6,
|
|
)
|
|
|
|
class Analyzer:
|
|
def __init__(self, cluster_threshold):
|
|
assert cluster_threshold == 0.85
|
|
|
|
def analyze(self, directions):
|
|
assert set(directions) == {0, 1, 2}
|
|
return result
|
|
|
|
monkeypatch.setattr(
|
|
"obliteratus.analysis.cross_layer.CrossLayerAlignmentAnalyzer",
|
|
Analyzer,
|
|
)
|
|
pipeline._harmful_means = {
|
|
0: torch.tensor([[1.0, 0.0]]),
|
|
1: torch.tensor([[3.0, 0.0]]),
|
|
2: torch.tensor([[0.0, 2.0]]),
|
|
}
|
|
pipeline._harmless_means = {idx: torch.zeros(1, 2) for idx in range(3)}
|
|
|
|
pipeline._analyze_cross_layer()
|
|
|
|
assert pipeline._insights.direction_clusters == [[0, 1], [2]]
|
|
assert pipeline._insights.cluster_count == 2
|
|
assert pipeline._insights.direction_persistence == 0.75
|
|
assert pipeline._insights.cluster_representative_layers == [1, 2]
|
|
|
|
|
|
def test_cross_layer_analysis_requires_two_nonzero_directions(pipeline, monkeypatch):
|
|
class Analyzer:
|
|
def __init__(self, **_kwargs):
|
|
raise AssertionError("analyzer should not be constructed")
|
|
|
|
monkeypatch.setattr(
|
|
"obliteratus.analysis.cross_layer.CrossLayerAlignmentAnalyzer",
|
|
Analyzer,
|
|
)
|
|
pipeline._harmful_means = {0: torch.ones(1, 2)}
|
|
pipeline._harmless_means = {0: torch.zeros(1, 2)}
|
|
|
|
pipeline._analyze_cross_layer()
|
|
|
|
assert pipeline._insights.cluster_count == 0
|
|
|
|
|
|
def test_defense_analysis_restores_directions_and_copies_profile(pipeline, monkeypatch):
|
|
original = {9: torch.tensor([9.0])}
|
|
observed = {}
|
|
|
|
class Evaluator:
|
|
def __init__(self, candidate):
|
|
observed["temporary"] = dict(candidate.refusal_directions)
|
|
|
|
def profile_defense(self):
|
|
return SimpleNamespace(
|
|
estimated_robustness="high",
|
|
self_repair_estimate=0.6,
|
|
entanglement_score=0.4,
|
|
)
|
|
|
|
def map_entanglement(self):
|
|
return SimpleNamespace(
|
|
most_entangled_layers=[2],
|
|
least_entangled_layers=[0],
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
"obliteratus.analysis.defense_robustness.DefenseRobustnessEvaluator",
|
|
Evaluator,
|
|
)
|
|
pipeline.refusal_directions = original
|
|
pipeline._harmful_means = {0: torch.tensor([[0.0, 2.0]])}
|
|
pipeline._harmless_means = {0: torch.zeros(1, 2)}
|
|
|
|
pipeline._analyze_defense_robustness()
|
|
|
|
assert torch.equal(observed["temporary"][0], torch.tensor([0.0, 1.0]))
|
|
assert pipeline.refusal_directions is original
|
|
assert pipeline._insights.estimated_robustness == "high"
|
|
assert pipeline._insights.self_repair_estimate == 0.6
|
|
assert pipeline._insights.entangled_layers == [2]
|
|
assert pipeline._insights.clean_layers == [0]
|
|
|
|
|
|
def test_sparsity_analysis_builds_a_plan_from_compatible_ffn_weights(
|
|
pipeline,
|
|
monkeypatch,
|
|
):
|
|
projection = SimpleNamespace(weight=SimpleNamespace(data=torch.eye(2)))
|
|
layer = SimpleNamespace()
|
|
plan = SimpleNamespace(
|
|
mean_refusal_sparsity_index=0.65,
|
|
recommended_sparsity=0.2,
|
|
most_sparse_layer=0,
|
|
most_dense_layer=0,
|
|
)
|
|
captured = {}
|
|
|
|
class Surgeon:
|
|
def __init__(self, auto_sparsity):
|
|
assert auto_sparsity is True
|
|
|
|
def plan_surgery(self, weights, directions):
|
|
captured["weights"] = weights
|
|
captured["directions"] = directions
|
|
return plan
|
|
|
|
monkeypatch.setattr(
|
|
"obliteratus.analysis.sparse_surgery.SparseDirectionSurgeon",
|
|
Surgeon,
|
|
)
|
|
monkeypatch.setattr(
|
|
"obliteratus.strategies.utils.get_layer_modules",
|
|
lambda _handle: [layer],
|
|
)
|
|
monkeypatch.setattr(
|
|
"obliteratus.strategies.utils.get_ffn_module",
|
|
lambda _layer, _arch: SimpleNamespace(down_proj=projection),
|
|
)
|
|
pipeline.handle = SimpleNamespace(architecture="gpt2")
|
|
pipeline._harmful_means = {0: torch.tensor([[0.0, 2.0]])}
|
|
pipeline._harmless_means = {0: torch.zeros(1, 2)}
|
|
|
|
pipeline._analyze_sparsity()
|
|
|
|
assert set(captured["weights"]) == {0}
|
|
assert torch.equal(captured["directions"][0], torch.tensor([0.0, 1.0]))
|
|
assert pipeline._insights.mean_refusal_sparsity_index == 0.65
|
|
assert pipeline._insights.recommended_sparsity == 0.2
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("method", "budget"),
|
|
[("dpo", 0.5), ("rlhf", 0.3), ("cai", 0.2), ("sft", 0.4), ("unknown", 0.35)],
|
|
)
|
|
def test_bayesian_warm_start_sets_alignment_specific_kl_budget(
|
|
pipeline,
|
|
method,
|
|
budget,
|
|
):
|
|
pipeline._insights.detected_alignment_method = method
|
|
|
|
pipeline._configure_bayesian_warm_start()
|
|
|
|
assert pipeline.kl_budget == budget
|
|
assert pipeline._bayesian_trials == 50
|
|
assert pipeline.layer_adaptive_strength is True
|
|
assert pipeline.float_layer_interpolation is True
|
|
assert pipeline.use_kl_optimization is True
|
|
|
|
|
|
def test_bayesian_warm_start_uses_strongest_cluster_and_entanglement(pipeline):
|
|
pipeline._harmful_means = {
|
|
0: torch.tensor([[1.0, 0.0]]),
|
|
1: torch.tensor([[4.0, 0.0]]),
|
|
2: torch.tensor([[2.0, 0.0]]),
|
|
3: torch.tensor([[1.0, 0.0]]),
|
|
}
|
|
pipeline._harmless_means = {idx: torch.zeros(1, 2) for idx in range(4)}
|
|
pipeline._insights.cluster_representative_layers = [0, 1]
|
|
pipeline._insights.direction_clusters = [[0, 1], [2, 3]]
|
|
pipeline._insights.direction_persistence = 0.5
|
|
pipeline._insights.entanglement_score = 0.8
|
|
|
|
pipeline._configure_bayesian_warm_start()
|
|
|
|
warm = pipeline._informed_warm_start
|
|
assert warm["peak_position"] == pytest.approx(1 / 3)
|
|
assert warm["spread"] == pytest.approx(1 / 3)
|
|
assert warm["min_weight"] == 0.1
|
|
assert warm["attn_scale"] == 0.7
|
|
assert warm["mlp_scale"] == 0.4
|
|
|
|
|
|
def test_excise_informed_routes_sparse_and_dense_paths(pipeline, monkeypatch):
|
|
calls = []
|
|
monkeypatch.setattr(pipeline, "_excise_sparse", lambda: calls.append("sparse"))
|
|
monkeypatch.setattr(
|
|
pipeline,
|
|
"_configure_bayesian_warm_start",
|
|
lambda: calls.append("warm-start"),
|
|
)
|
|
monkeypatch.setattr(pipeline, "_excise", lambda: calls.append("dense"))
|
|
|
|
pipeline._insights.use_sparse_surgery = True
|
|
pipeline._excise_informed()
|
|
pipeline._insights.use_sparse_surgery = False
|
|
pipeline._excise_informed()
|
|
|
|
assert calls == ["sparse", "warm-start", "dense"]
|
|
|
|
|
|
def test_verify_compensation_stops_when_no_residual_layers(pipeline, monkeypatch):
|
|
calls = []
|
|
|
|
def verify():
|
|
calls.append("verify")
|
|
pipeline._quality_metrics = {"refusal_rate": 0.9, "kl_divergence": 0.1}
|
|
|
|
monkeypatch.setattr(pipeline, "_verify", verify)
|
|
monkeypatch.setattr(pipeline, "_probe", lambda: calls.append("probe"))
|
|
|
|
def distill():
|
|
calls.append("distill")
|
|
pipeline._strong_layers = []
|
|
|
|
monkeypatch.setattr(pipeline, "_distill_inner", distill)
|
|
monkeypatch.setattr(pipeline, "_excise_informed", lambda: calls.append("excise"))
|
|
|
|
pipeline._verify_and_compensate()
|
|
|
|
assert calls == ["verify", "probe", "distill"]
|
|
assert pipeline._report.ouroboros_passes == 1
|
|
assert pipeline._report.final_refusal_rate == 0.9
|
|
|
|
|
|
def test_verify_compensation_stops_at_kl_ceiling(pipeline, monkeypatch):
|
|
outcomes = iter(
|
|
(
|
|
{"refusal_rate": 0.9, "kl_divergence": 0.1},
|
|
{"refusal_rate": 0.8, "kl_divergence": 0.8},
|
|
)
|
|
)
|
|
calls = []
|
|
|
|
def verify():
|
|
calls.append("verify")
|
|
pipeline._quality_metrics = next(outcomes)
|
|
|
|
monkeypatch.setattr(pipeline, "_verify", verify)
|
|
monkeypatch.setattr(pipeline, "_probe", lambda: calls.append("probe"))
|
|
monkeypatch.setattr(pipeline, "_distill_inner", lambda: setattr(pipeline, "_strong_layers", [1]))
|
|
monkeypatch.setattr(pipeline, "_excise_informed", lambda: calls.append("excise"))
|
|
pipeline.kl_budget = 0.3
|
|
|
|
pipeline._verify_and_compensate()
|
|
|
|
assert calls == ["verify", "probe", "excise", "verify"]
|
|
assert pipeline._report.ouroboros_passes == 1
|
|
assert pipeline._report.final_refusal_rate == 0.8
|
|
|
|
|
|
def test_rebirth_writes_model_tokenizer_and_research_metadata(pipeline, monkeypatch):
|
|
saved = []
|
|
|
|
class Artifact:
|
|
def __init__(self, name):
|
|
self.name = name
|
|
|
|
def save_pretrained(self, path):
|
|
saved.append((self.name, Path(path)))
|
|
|
|
pipeline.handle = SimpleNamespace(model=Artifact("model"), tokenizer=Artifact("tokenizer"))
|
|
pipeline._strong_layers = [1, 3]
|
|
pipeline._quality_metrics = {"refusal_rate": 0.1}
|
|
pipeline._insights.detected_alignment_method = "dpo"
|
|
pipeline._insights.recommended_layers = [1, 3]
|
|
pipeline._report.analysis_duration = 1.2
|
|
pipeline._report.total_duration = 3.4
|
|
pipeline._report.ouroboros_passes = 1
|
|
pipeline._report.final_refusal_rate = 0.1
|
|
events = []
|
|
monkeypatch.setattr(pipeline, "_emit", lambda *args, **kwargs: events.append((args, kwargs)))
|
|
ticks = iter((5.0, 5.5))
|
|
monkeypatch.setattr("obliteratus.informed_pipeline.time.time", lambda: next(ticks))
|
|
|
|
result = pipeline._rebirth_informed()
|
|
|
|
assert result == pipeline.output_dir
|
|
assert saved == [("model", result), ("tokenizer", result)]
|
|
metadata = json.loads((result / "abliteration_metadata.json").read_text())
|
|
assert metadata["analysis_insights"]["detected_alignment_method"] == "dpo"
|
|
assert metadata["derived_config"]["layers_used"] == [1, 3]
|
|
assert metadata["pipeline_stats"]["ouroboros_passes"] == 1
|
|
assert metadata["quality_metrics"] == {"refusal_rate": 0.1}
|
|
assert events[0][0][:2] == ("rebirth", "running")
|
|
assert events[-1][0][:2] == ("rebirth", "done")
|
|
|
|
|
|
def test_distill_single_direction_honors_recommended_and_skipped_layers(
|
|
pipeline,
|
|
monkeypatch,
|
|
):
|
|
pipeline.handle = SimpleNamespace(hidden_size=4096, total_params=3_000_000_000)
|
|
pipeline.n_directions = 1
|
|
pipeline._harmful_means = {
|
|
0: torch.tensor([[3.0, 4.0]]),
|
|
1: torch.tensor([[0.0, 2.0]]),
|
|
}
|
|
pipeline._harmless_means = {idx: torch.zeros(1, 2) for idx in range(2)}
|
|
pipeline._insights.recommended_layers = [0, 1, 99]
|
|
pipeline._insights.skip_layers = [1]
|
|
events = []
|
|
monkeypatch.setattr(pipeline, "_emit", lambda *args, **kwargs: events.append((args, kwargs)))
|
|
ticks = iter((1.0, 1.5))
|
|
monkeypatch.setattr("obliteratus.informed_pipeline.time.time", lambda: next(ticks))
|
|
|
|
pipeline._distill_informed()
|
|
|
|
assert torch.allclose(pipeline.refusal_directions[0], torch.tensor([0.6, 0.8]))
|
|
assert pipeline.refusal_subspaces[0].shape == (1, 2)
|
|
assert pipeline._strong_layers == [0]
|
|
assert events[0][0][:2] == ("distill", "running")
|
|
assert events[-1][1]["strong_layers"] == [0]
|
|
|
|
|
|
def test_distill_svd_sanitizes_nonfinite_input_and_enriches_category_directions(
|
|
pipeline,
|
|
monkeypatch,
|
|
):
|
|
pipeline.handle = SimpleNamespace(hidden_size=4096, total_params=3_000_000_000)
|
|
pipeline.n_directions = 3
|
|
pipeline.use_whitened_svd = False
|
|
pipeline._harmful_means = {0: torch.tensor([[2.0, 0.0, 0.0]])}
|
|
pipeline._harmless_means = {0: torch.zeros(1, 3)}
|
|
pipeline._harmful_acts = {
|
|
0: [
|
|
torch.tensor([[2.0, 0.0, 0.0]]),
|
|
torch.tensor([[0.0, float("nan"), 0.0]]),
|
|
torch.tensor([[0.0, 0.0, 1.0]]),
|
|
]
|
|
}
|
|
pipeline._harmless_acts = {0: [torch.zeros(1, 3) for _ in range(3)]}
|
|
pipeline._insights.cone_is_polyhedral = True
|
|
pipeline._insights.per_category_directions = {
|
|
"one": torch.tensor([0.0, 1.0, 0.0]),
|
|
"two": torch.tensor([0.0, 0.0, 1.0]),
|
|
}
|
|
monkeypatch.setattr(pipeline, "_select_layers_knee", lambda ranked: [ranked[0][0]])
|
|
monkeypatch.setattr(pipeline, "_emit", lambda *_args, **_kwargs: None)
|
|
|
|
pipeline._distill_informed()
|
|
|
|
assert pipeline._strong_layers == [0]
|
|
assert pipeline.refusal_subspaces[0].shape[0] >= 2
|
|
assert torch.isfinite(pipeline.refusal_subspaces[0]).all()
|
|
|
|
|
|
def test_distill_uses_whitened_extractor_for_multi_direction_models(
|
|
pipeline,
|
|
monkeypatch,
|
|
):
|
|
result = SimpleNamespace(
|
|
directions=torch.tensor([[1.0, 0.0], [0.0, 1.0]]),
|
|
singular_values=torch.tensor([3.0, 1.0]),
|
|
)
|
|
|
|
class Extractor:
|
|
def extract(self, harmful, harmless, *, n_directions, layer_idx):
|
|
assert harmful is pipeline._harmful_acts[0]
|
|
assert harmless is pipeline._harmless_acts[0]
|
|
assert (n_directions, layer_idx) == (2, 0)
|
|
return result
|
|
|
|
monkeypatch.setattr(
|
|
"obliteratus.analysis.whitened_svd.WhitenedSVDExtractor",
|
|
Extractor,
|
|
)
|
|
pipeline.handle = SimpleNamespace(hidden_size=4096, total_params=3_000_000_000)
|
|
pipeline.n_directions = 2
|
|
pipeline.use_whitened_svd = True
|
|
pipeline._harmful_means = {0: torch.tensor([[1.0, 0.0]])}
|
|
pipeline._harmless_means = {0: torch.zeros(1, 2)}
|
|
pipeline._harmful_acts = {0: [torch.ones(1, 2)]}
|
|
pipeline._harmless_acts = {0: [torch.zeros(1, 2)]}
|
|
monkeypatch.setattr(pipeline, "_select_layers_knee", lambda _ranked: [0])
|
|
monkeypatch.setattr(pipeline, "_emit", lambda *_args, **_kwargs: None)
|
|
|
|
pipeline._distill_informed()
|
|
|
|
assert torch.equal(pipeline.refusal_subspaces[0], result.directions)
|
|
assert torch.equal(pipeline.refusal_directions[0], result.directions[0])
|
|
|
|
|
|
def test_distill_leace_falls_back_per_layer_after_extractor_failure(
|
|
pipeline,
|
|
monkeypatch,
|
|
):
|
|
leace_result = SimpleNamespace(
|
|
direction=torch.tensor([0.0, 1.0]),
|
|
generalized_eigenvalue=4.0,
|
|
erasure_loss=0.2,
|
|
)
|
|
|
|
class Extractor:
|
|
def extract(self, _harmful, _harmless, *, layer_idx):
|
|
if layer_idx == 0:
|
|
return leace_result
|
|
raise RuntimeError("singular fixture")
|
|
|
|
monkeypatch.setattr("obliteratus.analysis.leace.LEACEExtractor", Extractor)
|
|
pipeline.handle = SimpleNamespace(hidden_size=4096, total_params=3_000_000_000)
|
|
pipeline.direction_method = "leace"
|
|
pipeline.n_directions = 1
|
|
pipeline._harmful_means = {
|
|
0: torch.tensor([[0.0, 1.0]]),
|
|
1: torch.tensor([[1.0, 0.0]]),
|
|
}
|
|
pipeline._harmless_means = {idx: torch.zeros(1, 2) for idx in range(2)}
|
|
pipeline._harmful_acts = {idx: [torch.ones(1, 2)] for idx in range(2)}
|
|
pipeline._harmless_acts = {idx: [torch.zeros(1, 2)] for idx in range(2)}
|
|
monkeypatch.setattr(pipeline, "_select_layers_knee", lambda _ranked: [0, 1])
|
|
monkeypatch.setattr(pipeline, "_emit", lambda *_args, **_kwargs: None)
|
|
|
|
pipeline._distill_informed()
|
|
|
|
assert torch.equal(pipeline.refusal_directions[0], leace_result.direction)
|
|
assert torch.equal(pipeline.refusal_directions[1], torch.tensor([1.0, 0.0]))
|
|
|
|
|
|
def test_sparse_excision_projects_attention_and_ffn_with_iterative_reprobe(
|
|
pipeline,
|
|
monkeypatch,
|
|
):
|
|
class Layer(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.anchor = torch.nn.Parameter(torch.ones(1))
|
|
|
|
layer = Layer()
|
|
attention = SimpleNamespace(o_proj=torch.nn.Linear(2, 2, bias=False))
|
|
ffn = SimpleNamespace(down_proj=torch.nn.Linear(2, 2, bias=False))
|
|
calls = []
|
|
|
|
class Surgeon:
|
|
def __init__(self, *, sparsity, auto_sparsity):
|
|
assert sparsity == 0.25
|
|
assert auto_sparsity is True
|
|
|
|
def apply_sparse_projection(self, weight, direction):
|
|
calls.append((weight.clone(), direction.clone()))
|
|
return weight * 0.5
|
|
|
|
monkeypatch.setattr(
|
|
"obliteratus.analysis.sparse_surgery.SparseDirectionSurgeon",
|
|
Surgeon,
|
|
)
|
|
monkeypatch.setattr(
|
|
"obliteratus.strategies.utils.get_layer_modules",
|
|
lambda _handle: [layer],
|
|
)
|
|
monkeypatch.setattr(
|
|
"obliteratus.strategies.utils.get_attention_module",
|
|
lambda _layer, _arch: attention,
|
|
)
|
|
monkeypatch.setattr(
|
|
"obliteratus.strategies.utils.get_ffn_module",
|
|
lambda _layer, _arch: ffn,
|
|
)
|
|
pipeline.handle = SimpleNamespace(architecture="gpt2")
|
|
pipeline._insights.recommended_sparsity = 0.25
|
|
pipeline._strong_layers = [0]
|
|
pipeline.refusal_subspaces = {0: torch.eye(2)}
|
|
pipeline.refinement_passes = 2
|
|
pipeline.true_iterative_refinement = True
|
|
monkeypatch.setattr(pipeline, "_probe", lambda: calls.append("probe"))
|
|
monkeypatch.setattr(pipeline, "_distill_inner", lambda: calls.append("distill"))
|
|
events = []
|
|
monkeypatch.setattr(pipeline, "_emit", lambda *args, **kwargs: events.append((args, kwargs)))
|
|
|
|
pipeline._excise_sparse()
|
|
|
|
projection_calls = [call for call in calls if isinstance(call, tuple)]
|
|
assert len(projection_calls) == 8
|
|
assert calls.count("probe") == 1
|
|
assert calls.count("distill") == 1
|
|
assert events[-1][1]["modified_count"] == 8
|
|
|
|
|
|
def test_verify_compensation_stops_when_kl_rises_sharply(pipeline, monkeypatch):
|
|
outcomes = iter(
|
|
(
|
|
{"refusal_rate": 0.9, "kl_divergence": 0.1},
|
|
{"refusal_rate": 0.8, "kl_divergence": 0.12},
|
|
{"refusal_rate": 0.7, "kl_divergence": 0.2},
|
|
)
|
|
)
|
|
|
|
def verify():
|
|
pipeline._quality_metrics = next(outcomes)
|
|
|
|
monkeypatch.setattr(pipeline, "_verify", verify)
|
|
monkeypatch.setattr(pipeline, "_probe", lambda: None)
|
|
monkeypatch.setattr(pipeline, "_distill_inner", lambda: setattr(pipeline, "_strong_layers", [0]))
|
|
monkeypatch.setattr(pipeline, "_excise_informed", lambda: None)
|
|
pipeline.kl_budget = 1.0
|
|
|
|
pipeline._verify_and_compensate()
|
|
|
|
assert pipeline._report.ouroboros_passes == 2
|
|
assert pipeline._report.final_refusal_rate == 0.7
|