mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-30 14:40:38 +02:00
2989 lines
116 KiB
Python
2989 lines
116 KiB
Python
"""Tests for the SOTA abliteration pipeline."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
from pathlib import Path
|
||
from unittest.mock import MagicMock, Mock, patch
|
||
|
||
import pytest
|
||
import torch
|
||
from transformers import GPT2Config, GPT2LMHeadModel, Qwen2Config, Qwen2ForCausalLM
|
||
|
||
from obliteratus.abliterate import (
|
||
HARMFUL_PROMPTS,
|
||
HARMLESS_PROMPTS,
|
||
METHODS,
|
||
STAGES,
|
||
AbliterationPipeline,
|
||
PipelineStage,
|
||
StageResult,
|
||
_is_coherent_completion,
|
||
)
|
||
from obliteratus.models.loader import ModelHandle
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Fixtures
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _make_tiny_handle():
|
||
"""Create a minimal ModelHandle with a tiny GPT-2 for testing."""
|
||
config = GPT2Config(
|
||
vocab_size=1000,
|
||
n_positions=128,
|
||
n_embd=64,
|
||
n_layer=4,
|
||
n_head=2,
|
||
n_inner=256,
|
||
)
|
||
model = GPT2LMHeadModel(config)
|
||
model.eval()
|
||
|
||
tokenizer = MagicMock()
|
||
tokenizer.pad_token = "<pad>"
|
||
tokenizer.eos_token = "<eos>"
|
||
tokenizer.return_value = {
|
||
"input_ids": torch.randint(0, 1000, (1, 10)),
|
||
"attention_mask": torch.ones(1, 10, dtype=torch.long),
|
||
}
|
||
tokenizer.decode.return_value = "The capital of France is Paris, a beautiful city"
|
||
|
||
handle = ModelHandle(
|
||
model=model,
|
||
tokenizer=tokenizer,
|
||
config=config,
|
||
model_name="gpt2-test",
|
||
task="causal_lm",
|
||
)
|
||
handle.snapshot()
|
||
return handle
|
||
|
||
|
||
def _make_varied_tokenizer(handle):
|
||
"""Set up a tokenizer mock that returns different tokens per call."""
|
||
call_count = [0]
|
||
def mock_tokenizer(prompt, **kwargs):
|
||
call_count[0] += 1
|
||
torch.manual_seed(call_count[0])
|
||
batch_size = len(prompt) if isinstance(prompt, list) else 1
|
||
return {
|
||
"input_ids": torch.randint(0, 1000, (batch_size, 5)),
|
||
"attention_mask": torch.ones(batch_size, 5, dtype=torch.long),
|
||
}
|
||
handle.tokenizer.side_effect = mock_tokenizer
|
||
|
||
|
||
@pytest.fixture
|
||
def handle():
|
||
return _make_tiny_handle()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Data & stage definitions
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestPrompts:
|
||
def test_harmful_prompts_expanded(self):
|
||
assert len(HARMFUL_PROMPTS) >= 99
|
||
|
||
def test_harmless_prompts_expanded(self):
|
||
assert len(HARMLESS_PROMPTS) >= 99
|
||
|
||
def test_prompt_lists_same_length(self):
|
||
assert len(HARMFUL_PROMPTS) == len(HARMLESS_PROMPTS)
|
||
|
||
def test_prompt_count_842(self):
|
||
"""842 prompts across the expanded built-in corpus."""
|
||
assert len(HARMFUL_PROMPTS) == 842
|
||
assert len(HARMLESS_PROMPTS) == 842
|
||
|
||
def test_prompt_volume_slicing(self):
|
||
"""Slicing at standard volumes gives correct counts."""
|
||
for n in (33, 66, 99, 256, 512):
|
||
assert len(HARMFUL_PROMPTS[:n]) == n
|
||
assert len(HARMLESS_PROMPTS[:n]) == n
|
||
|
||
|
||
class TestCoherenceScoring:
|
||
@pytest.mark.parametrize(
|
||
("prompt", "completion"),
|
||
[
|
||
("The capital of France is", "Paris, the country's largest city."),
|
||
("Water boils at a temperature of", "100 degrees Celsius at sea level."),
|
||
("A binary search algorithm works by", "dividing a sorted array in half around its middle element."),
|
||
],
|
||
)
|
||
def test_accepts_relevant_completion(self, prompt, completion):
|
||
assert _is_coherent_completion(prompt, completion) is True
|
||
|
||
def test_rejects_fluent_but_irrelevant_completion(self):
|
||
assert _is_coherent_completion(
|
||
"The capital of France is",
|
||
"A computational model composed of many connected layers and nodes.",
|
||
) is False
|
||
|
||
def test_rejects_corpus_contamination_seen_in_verifier(self):
|
||
assert _is_coherent_completion(
|
||
"Water boils at a temperature of",
|
||
"100°C. What is the temperature in Fahrenheit? - Brainly.in\nprofile\nanswer\nanswered",
|
||
) is False
|
||
|
||
def test_rejects_repetition(self):
|
||
assert _is_coherent_completion(
|
||
"The largest ocean on Earth is",
|
||
"Pacific Pacific Pacific Pacific Pacific Pacific Pacific Pacific.",
|
||
) is False
|
||
|
||
|
||
class TestStages:
|
||
def test_six_stages(self):
|
||
assert len(STAGES) == 6
|
||
|
||
def test_stage_keys(self):
|
||
keys = [s.key for s in STAGES]
|
||
assert keys == ["summon", "probe", "distill", "excise", "verify", "rebirth"]
|
||
|
||
def test_stage_dataclass(self):
|
||
stage = PipelineStage(key="test", name="TEST", description="A test stage")
|
||
assert stage.key == "test"
|
||
assert stage.name == "TEST"
|
||
|
||
def test_stage_result_defaults(self):
|
||
result = StageResult(stage="test", status="running")
|
||
assert result.message == ""
|
||
assert result.duration == 0.0
|
||
assert result.details == {}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Method presets
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestMethods:
|
||
def test_methods_exist(self):
|
||
assert set(METHODS.keys()) == {"basic", "advanced", "aggressive", "informed", "surgical", "inverted", "nuclear", "optimized", "failspy", "gabliteration", "heretic", "rdo", "spectral_cascade", "som"}
|
||
|
||
def test_basic_single_direction(self):
|
||
cfg = METHODS["basic"]
|
||
assert cfg["n_directions"] == 1
|
||
assert cfg["norm_preserve"] is False
|
||
assert cfg["regularization"] == 0.0
|
||
assert cfg["refinement_passes"] == 1
|
||
|
||
def test_advanced_multi_direction(self):
|
||
cfg = METHODS["advanced"]
|
||
assert cfg["n_directions"] > 1
|
||
assert cfg["norm_preserve"] is True
|
||
assert cfg["regularization"] > 0
|
||
assert cfg["refinement_passes"] >= 2
|
||
|
||
def test_aggressive_full_gabliteration(self):
|
||
cfg = METHODS["aggressive"]
|
||
assert cfg["n_directions"] >= 8
|
||
assert cfg["norm_preserve"] is True
|
||
assert cfg["refinement_passes"] >= 3
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Pipeline init
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestPipelineInit:
|
||
def test_default_prompts(self):
|
||
pipeline = AbliterationPipeline(model_name="test-model")
|
||
assert pipeline.harmful_prompts == HARMFUL_PROMPTS
|
||
assert pipeline.harmless_prompts == HARMLESS_PROMPTS
|
||
|
||
def test_custom_prompts(self):
|
||
harmful = ["bad prompt"]
|
||
harmless = ["good prompt"]
|
||
pipeline = AbliterationPipeline(
|
||
model_name="test-model",
|
||
harmful_prompts=harmful,
|
||
harmless_prompts=harmless,
|
||
)
|
||
assert pipeline.harmful_prompts == harmful
|
||
assert pipeline.harmless_prompts == harmless
|
||
|
||
def test_defaults(self):
|
||
pipeline = AbliterationPipeline(model_name="test-model")
|
||
assert pipeline.device == "auto"
|
||
assert pipeline.dtype == "float16"
|
||
assert pipeline.output_dir == Path("abliterated")
|
||
assert pipeline.trust_remote_code is False
|
||
assert pipeline.gpu_memory_utilization is None
|
||
assert pipeline.refusal_max_tokens == 128
|
||
assert pipeline.handle is None
|
||
|
||
@pytest.mark.parametrize("invalid", [0, -1, 1.5, True])
|
||
def test_refusal_max_tokens_must_be_positive_integer(self, invalid):
|
||
with pytest.raises(
|
||
ValueError, match="refusal_max_tokens must be a positive integer",
|
||
):
|
||
AbliterationPipeline(model_name="test-model", refusal_max_tokens=invalid)
|
||
|
||
def test_refusal_max_tokens_accepts_positive_override(self):
|
||
pipeline = AbliterationPipeline(
|
||
model_name="test-model", refusal_max_tokens=512,
|
||
)
|
||
assert pipeline.refusal_max_tokens == 512
|
||
|
||
@pytest.mark.parametrize(
|
||
"invalid", [0, -0.1, 1.1, float("nan"), float("inf"), True, "0.8"],
|
||
)
|
||
def test_gpu_memory_utilization_requires_bounded_number(self, invalid):
|
||
with pytest.raises(
|
||
ValueError, match=r"gpu_memory_utilization must be a number in \(0, 1\]",
|
||
):
|
||
AbliterationPipeline(
|
||
model_name="test-model", gpu_memory_utilization=invalid,
|
||
)
|
||
|
||
def test_gpu_memory_utilization_accepts_explicit_override(self):
|
||
pipeline = AbliterationPipeline(
|
||
model_name="test-model", gpu_memory_utilization=0.95,
|
||
)
|
||
assert pipeline.gpu_memory_utilization == 0.95
|
||
|
||
def test_summon_propagates_gpu_budget_without_forcing_snapshot_skip(self, monkeypatch):
|
||
handle = MagicMock()
|
||
handle.summary.return_value = {
|
||
"architecture": "fixture",
|
||
"num_layers": 2,
|
||
"num_heads": 4,
|
||
"hidden_size": 8,
|
||
"total_params": 32,
|
||
}
|
||
load_model = Mock(return_value=handle)
|
||
monkeypatch.setattr("obliteratus.abliterate.load_model", load_model)
|
||
pipeline = AbliterationPipeline(
|
||
model_name="test-model", gpu_memory_utilization=0.95,
|
||
)
|
||
|
||
pipeline._summon()
|
||
|
||
assert load_model.call_args.kwargs["gpu_memory_utilization"] == 0.95
|
||
assert "skip_snapshot" not in load_model.call_args.kwargs
|
||
|
||
def test_default_method_is_advanced(self):
|
||
pipeline = AbliterationPipeline(model_name="test-model")
|
||
assert pipeline.method == "advanced"
|
||
assert pipeline.n_directions == METHODS["advanced"]["n_directions"]
|
||
assert pipeline.norm_preserve == METHODS["advanced"]["norm_preserve"]
|
||
assert pipeline.regularization == METHODS["advanced"]["regularization"]
|
||
|
||
def test_method_basic(self):
|
||
pipeline = AbliterationPipeline(model_name="test-model", method="basic")
|
||
assert pipeline.n_directions == 1
|
||
assert pipeline.norm_preserve is False
|
||
assert pipeline.regularization == 0.0
|
||
|
||
def test_method_aggressive(self):
|
||
pipeline = AbliterationPipeline(model_name="test-model", method="aggressive")
|
||
assert pipeline.n_directions == 8
|
||
assert pipeline.norm_preserve is True
|
||
assert pipeline.refinement_passes == 3
|
||
|
||
def test_explicit_overrides_method(self):
|
||
pipeline = AbliterationPipeline(
|
||
model_name="test-model",
|
||
method="basic",
|
||
n_directions=6,
|
||
norm_preserve=True,
|
||
regularization=0.5,
|
||
refinement_passes=4,
|
||
)
|
||
assert pipeline.n_directions == 6
|
||
assert pipeline.norm_preserve is True
|
||
assert pipeline.regularization == 0.5
|
||
assert pipeline.refinement_passes == 4
|
||
|
||
def test_callbacks(self):
|
||
stage_results = []
|
||
log_msgs = []
|
||
pipeline = AbliterationPipeline(
|
||
model_name="test-model",
|
||
on_stage=lambda r: stage_results.append(r),
|
||
on_log=lambda m: log_msgs.append(m),
|
||
)
|
||
pipeline.log("hello")
|
||
assert log_msgs == ["hello"]
|
||
|
||
pipeline._emit("test", "running", "msg")
|
||
assert len(stage_results) == 1
|
||
assert stage_results[0].stage == "test"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# _project_out_advanced (norm-preserving + regularization)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestProjectOutAdvanced:
|
||
def test_norm_preserving(self):
|
||
"""Norm-preserving mode should keep Frobenius norm constant."""
|
||
class Wrapper(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.o_proj = torch.nn.Linear(4, 8, bias=False)
|
||
|
||
module = Wrapper()
|
||
torch.manual_seed(42)
|
||
module.o_proj.weight.data = torch.randn(8, 4)
|
||
original_norm = module.o_proj.weight.data.norm().item()
|
||
|
||
direction = torch.randn(4, 1)
|
||
direction = direction / direction.norm()
|
||
|
||
AbliterationPipeline._project_out_advanced(
|
||
module, direction, ["o_proj"], norm_preserve=True, regularization=0.0
|
||
)
|
||
|
||
new_norm = module.o_proj.weight.data.norm().item()
|
||
# With amplification cap (1.10x max), exact norm preservation isn't
|
||
# guaranteed on tiny matrices (hidden_dim=4) where a single direction
|
||
# removes a large fraction of energy. Verify the norm is closer to
|
||
# original than the un-preserved norm would be (i.e. cap is working).
|
||
# The new norm should be >= the un-preserved norm (cap restores some)
|
||
assert new_norm >= original_norm * 0.85, \
|
||
f"Norm should be approximately preserved (within cap): {original_norm:.4f} vs {new_norm:.4f}"
|
||
|
||
def test_regularization_partial_removal(self):
|
||
"""Regularization should preserve some of the refusal component."""
|
||
class Wrapper(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.o_proj = torch.nn.Linear(4, 8, bias=False)
|
||
|
||
module_full = Wrapper()
|
||
module_reg = Wrapper()
|
||
torch.manual_seed(42)
|
||
W_orig = torch.randn(8, 4)
|
||
module_full.o_proj.weight.data = W_orig.clone()
|
||
module_reg.o_proj.weight.data = W_orig.clone()
|
||
|
||
direction = torch.randn(4, 1)
|
||
direction = direction / direction.norm()
|
||
|
||
# Full removal
|
||
AbliterationPipeline._project_out_advanced(
|
||
module_full, direction, ["o_proj"], norm_preserve=False, regularization=0.0
|
||
)
|
||
# Regularized (30% preserved)
|
||
AbliterationPipeline._project_out_advanced(
|
||
module_reg, direction, ["o_proj"], norm_preserve=False, regularization=0.3
|
||
)
|
||
|
||
W_full = module_full.o_proj.weight.data
|
||
W_reg = module_reg.o_proj.weight.data
|
||
|
||
# Full removal should have zero projection on direction
|
||
proj_full = (W_full @ direction).norm().item()
|
||
assert proj_full < 1e-4
|
||
|
||
# Regularized should have non-zero projection (30% preserved)
|
||
proj_reg = (W_reg @ direction).norm().item()
|
||
proj_orig = (W_orig @ direction).norm().item()
|
||
expected_ratio = 0.3
|
||
actual_ratio = proj_reg / proj_orig if proj_orig > 0 else 0
|
||
assert abs(actual_ratio - expected_ratio) < 0.05, \
|
||
f"Expected ~{expected_ratio:.0%} preserved, got {actual_ratio:.0%}"
|
||
|
||
def test_norm_preserving_transposed(self):
|
||
"""Norm-preserving should also work for transposed weights."""
|
||
class Wrapper(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.c_proj = torch.nn.Linear(8, 4, bias=False)
|
||
|
||
module = Wrapper()
|
||
torch.manual_seed(42)
|
||
module.c_proj.weight.data = torch.randn(4, 8)
|
||
original_norm = module.c_proj.weight.data.norm().item()
|
||
|
||
direction = torch.randn(4, 1)
|
||
direction = direction / direction.norm()
|
||
|
||
AbliterationPipeline._project_out_advanced(
|
||
module, direction, ["c_proj"], norm_preserve=True, regularization=0.0
|
||
)
|
||
|
||
new_norm = module.c_proj.weight.data.norm().item()
|
||
# With amplification cap (1.10x max), exact norm preservation isn't
|
||
# guaranteed on tiny matrices where a single direction removes a large
|
||
# fraction of energy.
|
||
assert new_norm >= original_norm * 0.80, \
|
||
f"Norm should be approximately preserved (within cap): {original_norm:.4f} vs {new_norm:.4f}"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Full attention projection (q/k/v + o_proj)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestAttentionFullProjection:
|
||
"""Test that ALL attention weight matrices are projected (not just o_proj)."""
|
||
|
||
def test_qkv_all_projected(self):
|
||
"""q_proj, k_proj, v_proj should all be projected alongside o_proj."""
|
||
hidden = 16
|
||
|
||
class FakeAttn(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.q_proj = torch.nn.Linear(hidden, hidden, bias=False)
|
||
self.k_proj = torch.nn.Linear(hidden, hidden, bias=False)
|
||
self.v_proj = torch.nn.Linear(hidden, hidden, bias=False)
|
||
self.o_proj = torch.nn.Linear(hidden, hidden, bias=False)
|
||
|
||
attn = FakeAttn()
|
||
torch.manual_seed(42)
|
||
for p in attn.parameters():
|
||
p.data = torch.randn_like(p.data)
|
||
|
||
originals = {
|
||
name: getattr(attn, name).weight.data.clone()
|
||
for name in ["q_proj", "k_proj", "v_proj", "o_proj"]
|
||
}
|
||
|
||
d = torch.randn(hidden, 1)
|
||
d = d / d.norm()
|
||
|
||
from obliteratus.abliterate import _ATTN_OUT_NAMES, _ATTN_IN_NAMES
|
||
count = AbliterationPipeline._project_out_advanced(
|
||
attn, d, _ATTN_OUT_NAMES + _ATTN_IN_NAMES,
|
||
)
|
||
|
||
assert count == 4, f"Should project 4 weights (q/k/v/o), got {count}"
|
||
for name in ["q_proj", "k_proj", "v_proj", "o_proj"]:
|
||
assert not torch.allclose(
|
||
getattr(attn, name).weight.data, originals[name]
|
||
), f"{name} should be modified"
|
||
|
||
def test_project_all_does_not_early_return(self):
|
||
"""_project_out_advanced should project ALL matching weights, not just first."""
|
||
hidden = 16
|
||
|
||
class FakeModule(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.up_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
self.gate_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
|
||
mod = FakeModule()
|
||
torch.manual_seed(42)
|
||
orig_up = mod.up_proj.weight.data.clone()
|
||
orig_gate = mod.gate_proj.weight.data.clone()
|
||
|
||
d = torch.randn(hidden, 1)
|
||
d = d / d.norm()
|
||
|
||
from obliteratus.abliterate import _FFN_IN_NAMES
|
||
count = AbliterationPipeline._project_out_advanced(mod, d, _FFN_IN_NAMES)
|
||
|
||
assert count == 2, f"Should project both up_proj and gate_proj, got {count}"
|
||
assert not torch.allclose(mod.up_proj.weight.data, orig_up), "up_proj should be modified"
|
||
assert not torch.allclose(mod.gate_proj.weight.data, orig_gate), "gate_proj should be modified"
|
||
|
||
def test_lm_head_projection(self):
|
||
"""lm_head should be projectable via _project_out_advanced."""
|
||
hidden = 16
|
||
vocab = 100
|
||
|
||
class FakeModel(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.lm_head = torch.nn.Linear(hidden, vocab, bias=False)
|
||
|
||
model = FakeModel()
|
||
torch.manual_seed(42)
|
||
orig = model.lm_head.weight.data.clone()
|
||
|
||
d = torch.randn(hidden, 1)
|
||
d = d / d.norm()
|
||
|
||
count = AbliterationPipeline._project_out_advanced(
|
||
model, d, ["lm_head"], regularization=0.0,
|
||
)
|
||
|
||
assert count == 1, "Should project lm_head"
|
||
assert not torch.allclose(model.lm_head.weight.data, orig), "lm_head should be modified"
|
||
# Verify refusal direction is removed from lm_head
|
||
proj = (model.lm_head.weight.data @ d).norm().item()
|
||
assert proj < 1e-4, f"Refusal direction should be removed from lm_head, proj={proj}"
|
||
|
||
|
||
class TestKneeDetectionThreshold:
|
||
"""Test that knee detection uses 5% threshold to include more layers."""
|
||
|
||
def test_five_percent_threshold_includes_more(self):
|
||
"""Layers between 5% and 10% of max should now be included."""
|
||
# Layer norms: max=10.0, then several between 5%-10%
|
||
sorted_layers = [(0, 10.0), (1, 8.0), (2, 6.0), (3, 0.7), (4, 0.6)]
|
||
selected = AbliterationPipeline._select_layers_knee(sorted_layers)
|
||
# 0.7 and 0.6 are 7% and 6% of max — should now be included (> 5% threshold)
|
||
assert 3 in selected or 4 in selected, (
|
||
f"Layers with 6-7% of max signal should be included, got {selected}"
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# MoE projection (router, shared expert, input/output, fused)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestProjectMoEExperts:
|
||
"""Test the full MoE projection pipeline: router, shared expert, experts."""
|
||
|
||
def _make_direction(self, hidden_dim=16):
|
||
d = torch.randn(hidden_dim, 1)
|
||
return d / d.norm()
|
||
|
||
def test_router_gate_projected(self):
|
||
"""Router/gate weight should have refusal direction removed."""
|
||
hidden = 16
|
||
n_experts = 4
|
||
|
||
class FakeMoE(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.gate = torch.nn.Linear(hidden, n_experts, bias=True)
|
||
self.experts = torch.nn.ModuleList([
|
||
self._make_expert() for _ in range(n_experts)
|
||
])
|
||
|
||
@staticmethod
|
||
def _make_expert():
|
||
m = torch.nn.Module()
|
||
m.down_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
m.up_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
return m
|
||
|
||
moe = FakeMoE()
|
||
d = self._make_direction(hidden)
|
||
W_gate_orig = moe.gate.weight.data.clone()
|
||
|
||
count = AbliterationPipeline._project_moe_experts(moe, d)
|
||
assert count > 0
|
||
|
||
# Gate weight should have been modified
|
||
assert not torch.allclose(moe.gate.weight.data, W_gate_orig), \
|
||
"Router/gate weights should be projected"
|
||
|
||
# The gate weight's projection onto the direction should be ~0
|
||
proj = (moe.gate.weight.data @ d).norm().item()
|
||
assert proj < 1e-4, f"Gate should have no component along refusal dir, got {proj}"
|
||
|
||
def test_shared_expert_projected(self):
|
||
"""Shared expert (always-on) should have both input and output projected."""
|
||
hidden = 16
|
||
|
||
class FakeMoE(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.gate = torch.nn.Linear(hidden, 2, bias=False)
|
||
self.shared_expert = torch.nn.Module()
|
||
self.shared_expert.down_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
self.shared_expert.up_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
self.experts = torch.nn.ModuleList([
|
||
self._make_expert() for _ in range(2)
|
||
])
|
||
|
||
@staticmethod
|
||
def _make_expert():
|
||
m = torch.nn.Module()
|
||
m.down_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
m.up_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
return m
|
||
|
||
moe = FakeMoE()
|
||
d = self._make_direction(hidden)
|
||
shared_down_orig = moe.shared_expert.down_proj.weight.data.clone()
|
||
shared_up_orig = moe.shared_expert.up_proj.weight.data.clone()
|
||
|
||
count = AbliterationPipeline._project_moe_experts(moe, d)
|
||
assert count > 0
|
||
|
||
# Both shared expert output AND input projections should be modified
|
||
assert not torch.allclose(moe.shared_expert.down_proj.weight.data, shared_down_orig), \
|
||
"Shared expert output (down_proj) should be projected"
|
||
assert not torch.allclose(moe.shared_expert.up_proj.weight.data, shared_up_orig), \
|
||
"Shared expert input (up_proj) should be projected"
|
||
|
||
def test_expert_input_projections_projected(self):
|
||
"""Expert input projections (up_proj, gate_proj) should also be modified."""
|
||
hidden = 16
|
||
|
||
class FakeExpert(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.down_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
self.up_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
self.gate_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
|
||
class FakeMoE(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.experts = torch.nn.ModuleList([FakeExpert() for _ in range(2)])
|
||
|
||
moe = FakeMoE()
|
||
d = self._make_direction(hidden)
|
||
up_orig = moe.experts[0].up_proj.weight.data.clone()
|
||
|
||
count = AbliterationPipeline._project_moe_experts(moe, d)
|
||
|
||
# Each expert contributes 2 projections (output + input)
|
||
# 2 experts * 2 = 4 minimum
|
||
assert count >= 4, f"Expected >= 4 projections (out+in per expert), got {count}"
|
||
|
||
assert not torch.allclose(moe.experts[0].up_proj.weight.data, up_orig), \
|
||
"Expert input (up_proj) should be projected"
|
||
|
||
def test_fused_3d_output_and_input(self):
|
||
"""Fused 3D parameter patterns (GPT-OSS style) should project both directions."""
|
||
hidden = 16
|
||
intermediate = 32
|
||
n_experts = 4
|
||
|
||
class FusedExperts(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.down_proj = torch.nn.Parameter(torch.randn(n_experts, intermediate, hidden))
|
||
self.up_proj = torch.nn.Parameter(torch.randn(n_experts, intermediate, hidden))
|
||
|
||
class FakeMoE(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.experts = FusedExperts()
|
||
|
||
moe = FakeMoE()
|
||
d = self._make_direction(hidden)
|
||
down_orig = moe.experts.down_proj.data.clone()
|
||
up_orig = moe.experts.up_proj.data.clone()
|
||
|
||
count = AbliterationPipeline._project_moe_experts(moe, d)
|
||
|
||
# 4 experts output + 4 experts input = 8
|
||
assert count == 8, f"Expected 8 fused projections, got {count}"
|
||
|
||
assert not torch.allclose(moe.experts.down_proj.data, down_orig), \
|
||
"Fused output (down_proj) should be projected"
|
||
assert not torch.allclose(moe.experts.up_proj.data, up_orig), \
|
||
"Fused input (up_proj) should be projected"
|
||
|
||
def test_fused_3d_norm_preserve(self):
|
||
"""Fused 3D projections should preserve norms when requested."""
|
||
hidden = 16
|
||
intermediate = 32
|
||
n_experts = 4
|
||
|
||
class FusedExperts(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.down_proj = torch.nn.Parameter(torch.randn(n_experts, intermediate, hidden))
|
||
|
||
class FakeMoE(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.experts = FusedExperts()
|
||
|
||
moe = FakeMoE()
|
||
d = self._make_direction(hidden)
|
||
|
||
# Record per-expert norms before
|
||
orig_norms = [moe.experts.down_proj.data[i].norm().item() for i in range(n_experts)]
|
||
|
||
AbliterationPipeline._project_moe_experts(moe, d, norm_preserve=True)
|
||
|
||
# Check per-expert norms preserved
|
||
for i in range(n_experts):
|
||
new_norm = moe.experts.down_proj.data[i].norm().item()
|
||
assert abs(orig_norms[i] - new_norm) < 1e-3, \
|
||
f"Expert {i} norm not preserved: {orig_norms[i]:.4f} vs {new_norm:.4f}"
|
||
|
||
def test_no_experts_returns_zero(self):
|
||
"""Module without experts attribute should return 0."""
|
||
class NoMoE(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.mlp = torch.nn.Linear(16, 32)
|
||
|
||
moe = NoMoE()
|
||
d = self._make_direction(16)
|
||
assert AbliterationPipeline._project_moe_experts(moe, d) == 0
|
||
|
||
def test_router_bias_projected(self):
|
||
"""Router bias should be projected when project_biases=True."""
|
||
hidden = 16
|
||
|
||
class FakeMoE(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.gate = torch.nn.Linear(hidden, 4, bias=True)
|
||
self.experts = torch.nn.ModuleList([
|
||
self._make_expert() for _ in range(4)
|
||
])
|
||
|
||
@staticmethod
|
||
def _make_expert():
|
||
m = torch.nn.Module()
|
||
m.down_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
return m
|
||
|
||
moe = FakeMoE()
|
||
d = self._make_direction(hidden)
|
||
bias_orig = moe.gate.bias.data.clone()
|
||
|
||
count = AbliterationPipeline._project_moe_experts(moe, d, project_biases=True)
|
||
|
||
# Gate has 4 outputs (num_experts), direction has 16 dims
|
||
# bias shape (4,) != direction shape (16,), so bias won't match.
|
||
# This is correct: router bias is (num_experts,), not (hidden_dim,),
|
||
# so _project_bias won't modify it (shape mismatch is expected).
|
||
assert torch.allclose(moe.gate.bias.data, bias_orig), (
|
||
"Router bias should be unchanged when shape mismatches direction"
|
||
)
|
||
assert isinstance(count, int)
|
||
assert count > 0 # expert weights should still be projected
|
||
|
||
def test_router_auto_detection_fallback(self):
|
||
"""Unknown router name should be auto-detected and projected."""
|
||
import warnings as w
|
||
hidden = 16
|
||
n_experts = 4
|
||
|
||
class FakeMoE(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
# Unusual router name not in _ROUTER_NAMES
|
||
self.moe_gate_proj = torch.nn.Linear(hidden, n_experts, bias=False)
|
||
self.experts = torch.nn.ModuleList([
|
||
self._make_expert() for _ in range(n_experts)
|
||
])
|
||
|
||
@staticmethod
|
||
def _make_expert():
|
||
m = torch.nn.Module()
|
||
m.down_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
return m
|
||
|
||
moe = FakeMoE()
|
||
d = self._make_direction(hidden)
|
||
gate_orig = moe.moe_gate_proj.weight.data.clone()
|
||
|
||
with w.catch_warnings(record=True) as caught:
|
||
w.simplefilter("always")
|
||
AbliterationPipeline._project_moe_experts(moe, d)
|
||
|
||
# Should auto-detect and project the unusual router name
|
||
assert not torch.allclose(moe.moe_gate_proj.weight.data, gate_orig), \
|
||
"Auto-detected router should be projected"
|
||
|
||
# Should emit a warning about the auto-detection
|
||
auto_detect_warnings = [
|
||
x for x in caught
|
||
if "auto-detected" in str(x.message)
|
||
]
|
||
assert len(auto_detect_warnings) > 0, "Should warn about auto-detected router"
|
||
|
||
def test_full_moe_all_components(self):
|
||
"""End-to-end: all MoE components should be modified together."""
|
||
hidden = 16
|
||
|
||
class FakeExpert(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.down_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
self.up_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
|
||
class FakeMoE(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.gate = torch.nn.Linear(hidden, 4, bias=False)
|
||
self.shared_expert = torch.nn.Module()
|
||
self.shared_expert.down_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
self.shared_expert.up_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
self.experts = torch.nn.ModuleList([FakeExpert() for _ in range(4)])
|
||
|
||
moe = FakeMoE()
|
||
d = self._make_direction(hidden)
|
||
|
||
count = AbliterationPipeline._project_moe_experts(moe, d)
|
||
|
||
# Expected: 1 (gate) + 2 (shared out+in) + 4*2 (expert out+in) = 11
|
||
assert count == 11, f"Expected 11 total projections, got {count}"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# SOTA technique #1: Safety-neuron masking (GateBreaker-style z-score)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestSafetyNeuronMasking:
|
||
def test_outlier_neurons_zeroed(self):
|
||
"""Neurons with outsized refusal projection should be zeroed."""
|
||
hidden = 16
|
||
|
||
class Wrapper(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.down_proj = torch.nn.Linear(hidden, 64, bias=False)
|
||
|
||
module = Wrapper()
|
||
torch.manual_seed(42)
|
||
# Inject a few rows with very high projection along direction
|
||
d = torch.randn(hidden, 1)
|
||
d = d / d.norm()
|
||
# Make rows 0,1,2 have huge projection (outliers)
|
||
for i in range(3):
|
||
module.down_proj.weight.data[i] = d.squeeze() * 10.0
|
||
|
||
n_masked = AbliterationPipeline._mask_safety_neurons(
|
||
module, d, ["down_proj"], z_threshold=2.0,
|
||
)
|
||
|
||
assert n_masked >= 3, f"Expected >= 3 masked neurons, got {n_masked}"
|
||
# Masked rows should be zero
|
||
for i in range(3):
|
||
assert module.down_proj.weight.data[i].abs().max().item() < 1e-6
|
||
|
||
def test_no_outliers_no_masking(self):
|
||
"""When all neurons have similar projection, none should be masked."""
|
||
hidden = 16
|
||
|
||
class Wrapper(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.down_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
|
||
module = Wrapper()
|
||
# Uniform weights — no outliers
|
||
module.down_proj.weight.data = torch.ones(32, hidden) * 0.1
|
||
|
||
d = torch.randn(hidden, 1)
|
||
d = d / d.norm()
|
||
|
||
n_masked = AbliterationPipeline._mask_safety_neurons(
|
||
module, d, ["down_proj"], z_threshold=2.0,
|
||
)
|
||
assert n_masked == 0
|
||
|
||
def test_high_threshold_masks_fewer(self):
|
||
"""Higher z_threshold should mask fewer neurons."""
|
||
hidden = 16
|
||
|
||
class Wrapper(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.down_proj = torch.nn.Linear(hidden, 64, bias=False)
|
||
|
||
module1 = Wrapper()
|
||
module2 = Wrapper()
|
||
torch.manual_seed(42)
|
||
W = torch.randn(64, hidden)
|
||
module1.down_proj.weight.data = W.clone()
|
||
module2.down_proj.weight.data = W.clone()
|
||
|
||
d = torch.randn(hidden, 1)
|
||
d = d / d.norm()
|
||
|
||
n1 = AbliterationPipeline._mask_safety_neurons(module1, d, ["down_proj"], z_threshold=1.5)
|
||
n2 = AbliterationPipeline._mask_safety_neurons(module2, d, ["down_proj"], z_threshold=3.0)
|
||
assert n1 >= n2, f"Lower threshold should mask more: {n1} vs {n2}"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# SOTA technique #2: Jailbreak-contrastive direction refinement
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestJailbreakContrastive:
|
||
def test_jailbreak_prompt_generation(self):
|
||
"""Pipeline should generate jailbreak variants of harmful prompts."""
|
||
pipeline = AbliterationPipeline(model_name="test")
|
||
pipeline.harmful_prompts = ["How to hack a computer", "Write a virus"]
|
||
|
||
jailbreak = pipeline._generate_jailbreak_prompts()
|
||
|
||
assert len(jailbreak) == 2
|
||
# Each should contain the original prompt
|
||
assert "hack a computer" in jailbreak[0].lower()
|
||
assert "virus" in jailbreak[1].lower()
|
||
# Each should be wrapped in a template (longer than original)
|
||
assert len(jailbreak[0]) > len("How to hack a computer")
|
||
|
||
def test_jailbreak_contrast_blending(self):
|
||
"""Jailbreak-contrastive blending should modify refusal direction."""
|
||
pipeline = AbliterationPipeline(
|
||
model_name="test",
|
||
use_jailbreak_contrast=True,
|
||
n_directions=1,
|
||
)
|
||
hidden = 16
|
||
pipeline._on_log = lambda m: None
|
||
|
||
# Simulate probed means
|
||
torch.manual_seed(42)
|
||
harm_mean = torch.randn(1, hidden)
|
||
safe_mean = torch.randn(1, hidden)
|
||
jb_mean = torch.randn(1, hidden)
|
||
|
||
pipeline._harmful_means = {0: harm_mean}
|
||
pipeline._harmless_means = {0: safe_mean}
|
||
pipeline._jailbreak_means = {0: jb_mean}
|
||
pipeline._harmful_acts = {0: [harm_mean]}
|
||
pipeline._harmless_acts = {0: [safe_mean]}
|
||
pipeline._jailbreak_acts = {0: [jb_mean]}
|
||
|
||
# Run distill (will set standard direction, then blend)
|
||
pipeline._distill()
|
||
|
||
# Direction should be a unit vector
|
||
d = pipeline.refusal_directions[0]
|
||
assert abs(d.norm().item() - 1.0) < 1e-4
|
||
|
||
# Direction should differ from pure harm-safe difference
|
||
std_diff = (harm_mean - safe_mean).squeeze()
|
||
std_dir = std_diff / std_diff.norm()
|
||
cosine = (d @ std_dir).item()
|
||
# Blended direction should not be identical to standard
|
||
assert cosine < 0.99, f"Blended direction too similar to standard: cos={cosine}"
|
||
|
||
def test_surgical_method_enables_jailbreak(self):
|
||
"""Surgical method should enable jailbreak-contrastive by default."""
|
||
cfg = METHODS["surgical"]
|
||
assert cfg["use_jailbreak_contrast"] is True
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# SOTA technique #3: Layer-adaptive projection strength
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestLayerAdaptiveStrength:
|
||
def test_layer_weights_computed(self):
|
||
"""Layer-adaptive weights should be proportional to refusal signal."""
|
||
pipeline = AbliterationPipeline(
|
||
model_name="test",
|
||
layer_adaptive_strength=True,
|
||
n_directions=1,
|
||
)
|
||
hidden = 16
|
||
pipeline._on_log = lambda m: None
|
||
|
||
# Simulate: layer 0 has strong signal, layer 1 weak
|
||
torch.manual_seed(42)
|
||
strong_diff = torch.randn(1, hidden) * 10.0
|
||
weak_diff = torch.randn(1, hidden) * 1.0
|
||
zero_mean = torch.zeros(1, hidden)
|
||
|
||
pipeline._harmful_means = {0: strong_diff, 1: weak_diff}
|
||
pipeline._harmless_means = {0: zero_mean, 1: zero_mean}
|
||
pipeline._harmful_acts = {0: [strong_diff], 1: [weak_diff]}
|
||
pipeline._harmless_acts = {0: [zero_mean], 1: [zero_mean]}
|
||
|
||
pipeline._distill()
|
||
|
||
# Layer weights should exist for strong layers
|
||
assert len(pipeline._layer_excise_weights) > 0
|
||
# Strongest layer should have weight ~1.0
|
||
max_weight = max(pipeline._layer_excise_weights.values())
|
||
assert max_weight > 0.9, f"Max weight should be ~1.0, got {max_weight}"
|
||
|
||
def test_surgical_method_enables_adaptive(self):
|
||
"""Surgical method should enable layer-adaptive by default."""
|
||
cfg = METHODS["surgical"]
|
||
assert cfg["layer_adaptive_strength"] is True
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# SOTA technique #5: Attention head surgery
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestAttentionHeadSurgery:
|
||
def test_head_selective_projection(self):
|
||
"""Selective head projection should only modify targeted head rows."""
|
||
hidden = 16
|
||
n_heads = 4
|
||
head_dim = hidden // n_heads
|
||
|
||
class FakeAttn(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.o_proj = torch.nn.Linear(hidden, hidden, bias=False)
|
||
|
||
attn = FakeAttn()
|
||
torch.manual_seed(42)
|
||
W_orig = attn.o_proj.weight.data.clone()
|
||
|
||
d = torch.randn(hidden, 1)
|
||
d = d / d.norm()
|
||
|
||
# Head scores: head 0 is top safety head, head 3 is lowest
|
||
head_scores = [(0, 5.0), (1, 3.0), (2, 1.0), (3, 0.5)]
|
||
|
||
n_modified = AbliterationPipeline._project_head_selective(
|
||
attn, d, head_scores, n_heads=n_heads, head_fraction=0.25,
|
||
)
|
||
|
||
assert n_modified >= 1, "Should modify at least 1 head"
|
||
|
||
W_new = attn.o_proj.weight.data
|
||
# Head 0 columns (targeted) should be modified
|
||
assert not torch.allclose(
|
||
W_new[:, 0:head_dim], W_orig[:, 0:head_dim]
|
||
), "Targeted head 0 should be modified"
|
||
|
||
# Head 3 columns (NOT targeted) should be untouched
|
||
assert torch.allclose(
|
||
W_new[:, 3*head_dim:4*head_dim],
|
||
W_orig[:, 3*head_dim:4*head_dim],
|
||
), "Non-targeted head 3 should be untouched"
|
||
|
||
def test_head_surgery_norm_preserve(self):
|
||
"""Head surgery with norm_preserve should maintain per-head norms."""
|
||
hidden = 16
|
||
n_heads = 4
|
||
head_dim = hidden // n_heads
|
||
|
||
class FakeAttn(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.o_proj = torch.nn.Linear(hidden, hidden, bias=False)
|
||
|
||
attn = FakeAttn()
|
||
torch.manual_seed(42)
|
||
|
||
d = torch.randn(hidden, 1)
|
||
d = d / d.norm()
|
||
|
||
orig_norms = [
|
||
attn.o_proj.weight.data[:, h*head_dim:(h+1)*head_dim].norm().item()
|
||
for h in range(n_heads)
|
||
]
|
||
|
||
head_scores = [(0, 5.0), (1, 3.0), (2, 1.0), (3, 0.5)]
|
||
AbliterationPipeline._project_head_selective(
|
||
attn, d, head_scores, n_heads=n_heads,
|
||
head_fraction=0.5, norm_preserve=True,
|
||
)
|
||
|
||
# Targeted heads should have preserved norms
|
||
for h in range(2): # top 50% = 2 heads
|
||
new_norm = attn.o_proj.weight.data[:, h*head_dim:(h+1)*head_dim].norm().item()
|
||
assert abs(orig_norms[h] - new_norm) < 1e-3, \
|
||
f"Head {h} norm not preserved: {orig_norms[h]:.4f} vs {new_norm:.4f}"
|
||
|
||
def test_head_surgery_non_square_gqa(self):
|
||
"""Head surgery should work for GQA models with non-square o_proj (attn_dim != hidden_dim)."""
|
||
hidden_dim = 12 # model hidden dimension
|
||
attn_dim = 32 # attention dimension (n_heads * head_dim_attn)
|
||
n_heads = 4
|
||
head_dim_attn = attn_dim // n_heads # 8
|
||
|
||
class FakeAttnGQA(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
# o_proj maps attn_dim -> hidden_dim
|
||
# nn.Linear weight shape: (hidden_dim, attn_dim) = (12, 32)
|
||
self.o_proj = torch.nn.Linear(attn_dim, hidden_dim, bias=False)
|
||
|
||
attn = FakeAttnGQA()
|
||
torch.manual_seed(42)
|
||
attn.o_proj.weight.data = torch.randn(hidden_dim, attn_dim)
|
||
W_orig = attn.o_proj.weight.data.clone()
|
||
|
||
d = torch.randn(hidden_dim, 1)
|
||
d = d / d.norm()
|
||
|
||
head_scores = [(0, 5.0), (1, 3.0), (2, 1.0), (3, 0.5)]
|
||
|
||
n_modified = AbliterationPipeline._project_head_selective(
|
||
attn, d, head_scores, n_heads=n_heads, head_fraction=0.25,
|
||
)
|
||
|
||
assert n_modified >= 1, "Should modify at least 1 head"
|
||
|
||
W_new = attn.o_proj.weight.data
|
||
# Head 0 columns (targeted) should be modified
|
||
assert not torch.allclose(
|
||
W_new[:, 0:head_dim_attn], W_orig[:, 0:head_dim_attn]
|
||
), "Targeted head 0 should be modified"
|
||
|
||
# Head 3 columns (NOT targeted) should be untouched
|
||
assert torch.allclose(
|
||
W_new[:, 3*head_dim_attn:4*head_dim_attn],
|
||
W_orig[:, 3*head_dim_attn:4*head_dim_attn],
|
||
), "Non-targeted head 3 should be untouched"
|
||
|
||
def test_head_surgery_gqa_norm_preserve(self):
|
||
"""Head surgery on GQA non-square o_proj with norm_preserve."""
|
||
hidden_dim = 12
|
||
attn_dim = 32
|
||
n_heads = 4
|
||
head_dim_attn = attn_dim // n_heads
|
||
|
||
class FakeAttnGQA(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.o_proj = torch.nn.Linear(attn_dim, hidden_dim, bias=False)
|
||
|
||
attn = FakeAttnGQA()
|
||
torch.manual_seed(42)
|
||
attn.o_proj.weight.data = torch.randn(hidden_dim, attn_dim)
|
||
|
||
d = torch.randn(hidden_dim, 1)
|
||
d = d / d.norm()
|
||
|
||
orig_norms = [
|
||
attn.o_proj.weight.data[:, h*head_dim_attn:(h+1)*head_dim_attn].norm().item()
|
||
for h in range(n_heads)
|
||
]
|
||
|
||
head_scores = [(0, 5.0), (1, 3.0), (2, 1.0), (3, 0.5)]
|
||
AbliterationPipeline._project_head_selective(
|
||
attn, d, head_scores, n_heads=n_heads,
|
||
head_fraction=0.5, norm_preserve=True,
|
||
)
|
||
|
||
for h in range(2): # top 50% = 2 heads
|
||
new_norm = attn.o_proj.weight.data[:, h*head_dim_attn:(h+1)*head_dim_attn].norm().item()
|
||
assert abs(orig_norms[h] - new_norm) < 1e-3, \
|
||
f"GQA head {h} norm not preserved: {orig_norms[h]:.4f} vs {new_norm:.4f}"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# SOTA technique #6: SAE feature-level abliteration
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@pytest.mark.filterwarnings(
|
||
"ignore:SAE held-out reconstruction MSE.*:UserWarning",
|
||
)
|
||
class TestSAEAbliteration:
|
||
"""Exercise deliberately undertrained SAEs without weakening global warnings."""
|
||
|
||
def test_sae_train_and_reconstruct(self):
|
||
"""SAE should train and reconstruct activations."""
|
||
from obliteratus.analysis.sae_abliteration import train_sae
|
||
|
||
hidden = 32
|
||
# Generate synthetic activations
|
||
torch.manual_seed(42)
|
||
acts = [torch.randn(hidden) for _ in range(64)]
|
||
|
||
sae = train_sae(acts, hidden, expansion=2, n_epochs=10, lr=1e-3, device="cpu")
|
||
|
||
# Forward pass should work
|
||
x = torch.randn(1, hidden)
|
||
x_hat, z = sae(x)
|
||
assert x_hat.shape == x.shape
|
||
assert z.shape == (1, 2 * hidden) # expansion=2
|
||
|
||
# Z should be sparse (ReLU activation)
|
||
assert (z == 0).float().mean() > 0.3, "Features should be sparse"
|
||
|
||
def test_refusal_feature_identification(self):
|
||
"""SAE should identify features that differ between harmful/harmless."""
|
||
from obliteratus.analysis.sae_abliteration import (
|
||
train_sae, identify_refusal_features,
|
||
)
|
||
|
||
hidden = 32
|
||
torch.manual_seed(42)
|
||
|
||
# Create activations with clear harmful/harmless separation
|
||
refusal_dir = torch.randn(hidden)
|
||
refusal_dir = refusal_dir / refusal_dir.norm()
|
||
|
||
harmful_acts = [torch.randn(hidden) + 2.0 * refusal_dir for _ in range(32)]
|
||
harmless_acts = [torch.randn(hidden) - 2.0 * refusal_dir for _ in range(32)]
|
||
all_acts = harmful_acts + harmless_acts
|
||
|
||
sae = train_sae(all_acts, hidden, expansion=2, n_epochs=30, lr=3e-4, device="cpu")
|
||
result = identify_refusal_features(
|
||
sae, harmful_acts, harmless_acts, layer_idx=0, top_k=4,
|
||
)
|
||
|
||
assert result.n_refusal_features == 4
|
||
assert result.sae_directions.shape == (4, hidden)
|
||
assert result.variance_explained > 0.0
|
||
# SAE directions should have some alignment with the actual refusal direction
|
||
best_cos = max(
|
||
abs((result.sae_directions[i] @ refusal_dir).item())
|
||
for i in range(result.sae_directions.shape[0])
|
||
)
|
||
assert best_cos > 0.1, f"SAE should find direction aligned with refusal: best_cos={best_cos}"
|
||
|
||
def test_sae_directions_unit_norm(self):
|
||
"""SAE-derived directions should be unit normalized."""
|
||
from obliteratus.analysis.sae_abliteration import (
|
||
train_sae, identify_refusal_features,
|
||
)
|
||
|
||
hidden = 16
|
||
torch.manual_seed(42)
|
||
harmful = [torch.randn(hidden) + torch.ones(hidden) for _ in range(16)]
|
||
harmless = [torch.randn(hidden) - torch.ones(hidden) for _ in range(16)]
|
||
|
||
sae = train_sae(harmful + harmless, hidden, expansion=2, n_epochs=10, device="cpu")
|
||
result = identify_refusal_features(sae, harmful, harmless, 0, top_k=3)
|
||
|
||
for i in range(result.sae_directions.shape[0]):
|
||
norm = result.sae_directions[i].norm().item()
|
||
assert abs(norm - 1.0) < 1e-3, f"Direction {i} norm={norm}, expected 1.0"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Surgical method preset
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestSurgicalMethod:
|
||
def test_surgical_enables_all_sota(self):
|
||
"""Surgical method should enable all 6 SOTA techniques."""
|
||
cfg = METHODS["surgical"]
|
||
assert cfg["use_jailbreak_contrast"] is True
|
||
assert cfg["layer_adaptive_strength"] is True
|
||
assert cfg["safety_neuron_masking"] is True
|
||
assert cfg["per_expert_directions"] is True
|
||
assert cfg["attention_head_surgery"] is True
|
||
assert cfg["use_sae_features"] is True
|
||
|
||
def test_basic_disables_all_sota(self):
|
||
"""Basic method should not enable SOTA techniques (no keys or False)."""
|
||
cfg = METHODS["basic"]
|
||
assert cfg.get("use_jailbreak_contrast", False) is False
|
||
assert cfg.get("layer_adaptive_strength", False) is False
|
||
assert cfg.get("safety_neuron_masking", False) is False
|
||
|
||
def test_pipeline_init_surgical(self):
|
||
"""Pipeline initialized with surgical method should have all flags set."""
|
||
pipeline = AbliterationPipeline(model_name="test", method="surgical")
|
||
assert pipeline.use_jailbreak_contrast is True
|
||
assert pipeline.layer_adaptive_strength is True
|
||
assert pipeline.safety_neuron_masking is True
|
||
assert pipeline.per_expert_directions is True
|
||
assert pipeline.attention_head_surgery is True
|
||
assert pipeline.use_sae_features is True
|
||
|
||
def test_pipeline_init_explicit_override(self):
|
||
"""Explicit params should override method defaults."""
|
||
pipeline = AbliterationPipeline(
|
||
model_name="test", method="surgical",
|
||
safety_neuron_masking=False,
|
||
)
|
||
assert pipeline.safety_neuron_masking is False
|
||
assert pipeline.use_jailbreak_contrast is True # rest still from surgical
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Inverted method (semantic refusal inversion)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestInvertedMethod:
|
||
def test_inverted_preset_config(self):
|
||
"""Inverted method preset should enable inversion flag."""
|
||
cfg = METHODS["inverted"]
|
||
assert cfg["invert_refusal"] is True
|
||
assert cfg["n_directions"] == 8
|
||
assert cfg["use_jailbreak_contrast"] is True
|
||
|
||
def test_surgical_does_not_invert(self):
|
||
"""Surgical method should NOT enable inversion by default."""
|
||
cfg = METHODS["surgical"]
|
||
assert cfg.get("invert_refusal", False) is False
|
||
|
||
def test_pipeline_init_inverted(self):
|
||
"""Pipeline initialized with inverted method should have flag set."""
|
||
pipeline = AbliterationPipeline(model_name="test", method="inverted")
|
||
assert pipeline.invert_refusal is True
|
||
assert pipeline.use_jailbreak_contrast is True
|
||
assert pipeline.safety_neuron_masking is False # zeroing + reflection is destructive
|
||
|
||
def test_pipeline_invert_explicit_override(self):
|
||
"""Explicit invert_refusal param should override method default."""
|
||
pipeline = AbliterationPipeline(
|
||
model_name="test", method="surgical", invert_refusal=True,
|
||
)
|
||
assert pipeline.invert_refusal is True
|
||
|
||
pipeline2 = AbliterationPipeline(
|
||
model_name="test", method="inverted", invert_refusal=False,
|
||
)
|
||
assert pipeline2.invert_refusal is False
|
||
|
||
def test_reflection_math(self):
|
||
"""2x projection (reflection) should negate the refusal component."""
|
||
hidden = 16
|
||
|
||
class Wrapper(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.o_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
|
||
module = Wrapper()
|
||
torch.manual_seed(42)
|
||
W_orig = module.o_proj.weight.data.clone()
|
||
|
||
d = torch.randn(hidden, 1)
|
||
d = d / d.norm()
|
||
|
||
# Original projection onto d
|
||
orig_proj = (W_orig @ d).squeeze()
|
||
|
||
# Reflection: regularization=-1.0 → scale=2.0
|
||
AbliterationPipeline._project_out_advanced(
|
||
module, d, ["o_proj"], regularization=-1.0,
|
||
)
|
||
|
||
W_reflected = module.o_proj.weight.data
|
||
new_proj = (W_reflected @ d).squeeze()
|
||
|
||
# After reflection, projection should be NEGATED (sign flipped)
|
||
assert torch.allclose(new_proj, -orig_proj, atol=1e-4), (
|
||
f"Reflected projection should be negated: expected ~{-orig_proj[:3]} got {new_proj[:3]}"
|
||
)
|
||
|
||
def test_reflection_preserves_orthogonal_component(self):
|
||
"""Reflection should not change the component perpendicular to d."""
|
||
hidden = 8
|
||
|
||
class Wrapper(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.o_proj = torch.nn.Linear(hidden, 16, bias=False)
|
||
|
||
module = Wrapper()
|
||
torch.manual_seed(42)
|
||
W_orig = module.o_proj.weight.data.clone()
|
||
|
||
d = torch.randn(hidden, 1)
|
||
d = d / d.norm()
|
||
|
||
# Compute original orthogonal component
|
||
orig_d_component = (W_orig @ d) @ d.T # rank-1 matrix: projection onto d
|
||
orig_ortho = W_orig - orig_d_component # everything except d-component
|
||
|
||
AbliterationPipeline._project_out_advanced(
|
||
module, d, ["o_proj"], regularization=-1.0,
|
||
)
|
||
|
||
W_reflected = module.o_proj.weight.data
|
||
new_d_component = (W_reflected @ d) @ d.T
|
||
new_ortho = W_reflected - new_d_component
|
||
|
||
# Orthogonal component should be unchanged
|
||
assert torch.allclose(orig_ortho, new_ortho, atol=1e-4), (
|
||
"Reflection should preserve orthogonal component"
|
||
)
|
||
|
||
def test_moe_expert_safety_classification(self):
|
||
"""_identify_safety_experts should classify experts by router affinity."""
|
||
hidden = 16
|
||
n_experts = 4
|
||
|
||
class FakeMoE(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.gate = torch.nn.Linear(hidden, n_experts, bias=False)
|
||
self.experts = torch.nn.ModuleList([
|
||
torch.nn.Linear(hidden, hidden) for _ in range(n_experts)
|
||
])
|
||
|
||
class FakeLayer(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.self_attn = torch.nn.Module()
|
||
self.self_attn.o_proj = torch.nn.Linear(hidden, hidden, bias=False)
|
||
self.mlp = FakeMoE()
|
||
|
||
from obliteratus.models.loader import ModelHandle
|
||
from unittest.mock import MagicMock
|
||
from transformers import GPT2Config
|
||
|
||
config = GPT2Config(n_embd=hidden, n_head=2, n_layer=1, vocab_size=100, n_positions=64)
|
||
model = MagicMock()
|
||
model.parameters.return_value = iter([torch.zeros(1)])
|
||
|
||
handle = ModelHandle(
|
||
model=model, tokenizer=MagicMock(),
|
||
config=config, model_name="test", task="causal_lm",
|
||
)
|
||
|
||
pipeline = AbliterationPipeline(model_name="test", method="inverted")
|
||
pipeline.handle = handle
|
||
pipeline._on_log = lambda m: None
|
||
pipeline._on_stage = lambda r: None
|
||
|
||
# Set up fake layer and direction
|
||
layer = FakeLayer()
|
||
torch.manual_seed(42)
|
||
|
||
# Make router weight so expert 0 has highest affinity for d
|
||
d = torch.randn(hidden)
|
||
d = d / d.norm()
|
||
# Set router weights: expert 0 aligned with d, expert 3 anti-aligned
|
||
layer.mlp.gate.weight.data[0] = d * 5.0
|
||
layer.mlp.gate.weight.data[1] = d * 1.0
|
||
layer.mlp.gate.weight.data[2] = d * -1.0
|
||
layer.mlp.gate.weight.data[3] = d * -5.0
|
||
|
||
# Mock get_layer_modules to return our fake layer
|
||
import obliteratus.abliterate as abl_module
|
||
orig_get_layers = abl_module.get_layer_modules
|
||
orig_get_ffn = abl_module.get_ffn_module
|
||
abl_module.get_layer_modules = lambda h: [layer]
|
||
abl_module.get_ffn_module = lambda lay, a: lay.mlp
|
||
try:
|
||
pipeline.refusal_directions = {0: d}
|
||
pipeline._strong_layers = [0]
|
||
pipeline._identify_safety_experts()
|
||
finally:
|
||
abl_module.get_layer_modules = orig_get_layers
|
||
abl_module.get_ffn_module = orig_get_ffn
|
||
|
||
assert 0 in pipeline._expert_safety_scores
|
||
scores = pipeline._expert_safety_scores[0]
|
||
# Expert 0 should be highest safety affinity
|
||
assert scores[0][0] == 0, f"Expert 0 should be top safety, got {scores[0]}"
|
||
# Expert 3 should be lowest
|
||
assert scores[-1][0] == 3, f"Expert 3 should be lowest, got {scores[-1]}"
|
||
|
||
def test_moe_inverted_excision_selective(self):
|
||
"""Inverted MoE excision should reflect safety experts and remove from capability."""
|
||
hidden = 16
|
||
n_experts = 4
|
||
|
||
class FakeExpert(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.down_proj = torch.nn.Linear(hidden, hidden, bias=False)
|
||
|
||
class FakeMoE(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.gate = torch.nn.Linear(hidden, n_experts, bias=False)
|
||
self.experts = torch.nn.ModuleList([FakeExpert() for _ in range(n_experts)])
|
||
|
||
moe = FakeMoE()
|
||
torch.manual_seed(42)
|
||
for p in moe.parameters():
|
||
p.data = torch.randn_like(p.data)
|
||
|
||
d = torch.randn(hidden, 1)
|
||
d = d / d.norm()
|
||
|
||
# Set up safety scores: experts 0,1 are safety, 2,3 are capability
|
||
pipeline = AbliterationPipeline(model_name="test", method="inverted")
|
||
pipeline._on_log = lambda m: None
|
||
pipeline._on_stage = lambda r: None
|
||
pipeline._expert_safety_scores = {
|
||
0: [(0, 5.0), (1, 3.0), (2, -1.0), (3, -3.0)]
|
||
}
|
||
|
||
orig_router = moe.gate.weight.data.clone()
|
||
|
||
count = pipeline._project_moe_experts_inverted(
|
||
moe, d, 0, norm_preserve=False, project_biases=False,
|
||
)
|
||
|
||
assert count > 0, "Should project some weights"
|
||
|
||
# Router should be reflected (capped at 1.5x to prevent extreme logits
|
||
# that cause CUDA illegal memory access in batched expert forward).
|
||
# With router_reg = max(reflect_reg, -0.5) → scale = 1.5:
|
||
# new_proj ≈ orig_proj - 1.5 * orig_proj = -0.5 * orig_proj
|
||
# Additionally, _stabilize_router_weights clamps outliers, so we
|
||
# verify the sign is flipped and magnitude is substantial.
|
||
router_proj = (moe.gate.weight.data @ d.squeeze()).squeeze()
|
||
orig_router_proj = (orig_router @ d.squeeze()).squeeze()
|
||
cosine = torch.nn.functional.cosine_similarity(
|
||
router_proj.unsqueeze(0), -orig_router_proj.unsqueeze(0),
|
||
)
|
||
assert cosine > 0.5, (
|
||
f"Router projection should be at least partially reflected, cosine={cosine.item():.3f}"
|
||
)
|
||
|
||
# Safety expert 0: should be reflected (projection negated)
|
||
e0_proj = (moe.experts[0].down_proj.weight.data @ d).norm()
|
||
# After reflection the projection doesn't go to zero — it negates
|
||
assert e0_proj > 1e-4, "Safety expert should have non-zero projection (reflected, not removed)"
|
||
|
||
# Capability expert 3: should have projection removed (near zero)
|
||
e3_proj = (moe.experts[3].down_proj.weight.data @ d).norm().item()
|
||
assert e3_proj < 1e-3, f"Capability expert should have projection removed, got {e3_proj}"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Nuclear method
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestNuclearMethod:
|
||
def test_nuclear_preset_config(self):
|
||
"""Nuclear method should match inverted baseline + permanent weight techniques."""
|
||
cfg = METHODS["nuclear"]
|
||
assert cfg["invert_refusal"] is True
|
||
assert cfg["n_directions"] == 4 # fewer than inverted to avoid over-ablation
|
||
assert cfg["refinement_passes"] == 2 # same as inverted
|
||
assert cfg["reflection_strength"] == 1.25 # tempered for CoT coherence
|
||
assert cfg["project_embeddings"] is True
|
||
assert cfg["embed_regularization"] == 0.50 # conservative cascade limit
|
||
assert cfg["activation_steering"] is True # residual cleanup hooks
|
||
assert cfg["steering_strength"] == 0.15 # light residual correction
|
||
assert cfg["expert_transplant"] is True
|
||
assert cfg["transplant_blend"] == 0.10 # gentle nudge, not overwrite
|
||
assert cfg["use_jailbreak_contrast"] is True
|
||
assert cfg["attention_head_surgery"] is True
|
||
assert cfg["layer_adaptive_strength"] is True # per-layer scaling
|
||
|
||
def test_nuclear_pipeline_init(self):
|
||
"""Pipeline initialized with nuclear method should have all flags set."""
|
||
pipeline = AbliterationPipeline(model_name="test", method="nuclear")
|
||
assert pipeline.invert_refusal is True
|
||
assert pipeline.reflection_strength == 1.25
|
||
assert pipeline.embed_regularization == 0.50
|
||
assert pipeline.transplant_blend == 0.10
|
||
assert pipeline.project_embeddings is True
|
||
assert pipeline.activation_steering is True # residual cleanup
|
||
assert pipeline.expert_transplant is True
|
||
assert pipeline.n_directions == 4
|
||
assert pipeline.refinement_passes == 2
|
||
assert pipeline.layer_adaptive_strength is True
|
||
|
||
def test_reflection_strength_configurable(self):
|
||
"""reflection_strength should be explicitly overridable."""
|
||
pipeline = AbliterationPipeline(
|
||
model_name="test", method="inverted", reflection_strength=3.0,
|
||
)
|
||
assert pipeline.reflection_strength == 3.0
|
||
|
||
def test_inverted_default_strength_is_2(self):
|
||
"""Inverted method should default to reflection_strength=2.0."""
|
||
pipeline = AbliterationPipeline(model_name="test", method="inverted")
|
||
assert pipeline.reflection_strength == 2.0
|
||
|
||
def test_boosted_reflection_math(self):
|
||
"""2.5x reflection should produce stronger negation than 2x."""
|
||
hidden = 16
|
||
|
||
class Wrapper(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.o_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
|
||
d = torch.randn(hidden, 1)
|
||
d = d / d.norm()
|
||
|
||
# 2x reflection
|
||
module_2x = Wrapper()
|
||
torch.manual_seed(42)
|
||
module_2x.o_proj.weight.data = torch.randn(32, hidden)
|
||
orig = module_2x.o_proj.weight.data.clone()
|
||
AbliterationPipeline._project_out_advanced(
|
||
module_2x, d, ["o_proj"], regularization=-1.0, # scale=2.0
|
||
)
|
||
proj_2x = (module_2x.o_proj.weight.data @ d).squeeze()
|
||
|
||
# 2.5x reflection
|
||
module_25x = Wrapper()
|
||
module_25x.o_proj.weight.data = orig.clone()
|
||
AbliterationPipeline._project_out_advanced(
|
||
module_25x, d, ["o_proj"], regularization=-1.5, # scale=2.5
|
||
)
|
||
proj_25x = (module_25x.o_proj.weight.data @ d).squeeze()
|
||
|
||
# 2.5x should be 25% stronger negation than 2x
|
||
assert proj_25x.norm() > proj_2x.norm(), (
|
||
"2.5x reflection should produce stronger (more negative) projection than 2x"
|
||
)
|
||
|
||
def test_activation_steering_hook(self):
|
||
"""Steering hooks should subtract refusal direction from hidden states."""
|
||
hidden = 8
|
||
|
||
class FakeLayer(torch.nn.Module):
|
||
def forward(self, x):
|
||
return x
|
||
|
||
layer = FakeLayer()
|
||
layers = torch.nn.ModuleList([layer])
|
||
|
||
# Explicitly enable steering (nuclear preset has it off by default)
|
||
pipeline = AbliterationPipeline(
|
||
model_name="test", method="inverted", activation_steering=True,
|
||
steering_strength=0.5,
|
||
)
|
||
pipeline._on_log = lambda m: None
|
||
pipeline._on_stage = lambda r: None
|
||
|
||
d = torch.randn(hidden)
|
||
d = d / d.norm()
|
||
pipeline.refusal_directions = {0: d}
|
||
pipeline._strong_layers = [0]
|
||
|
||
n_hooks = pipeline._install_activation_steering(layers)
|
||
assert n_hooks == 1
|
||
assert len(pipeline._steering_hooks) == 1
|
||
|
||
# Create a hidden state with strong refusal component
|
||
batch = torch.randn(1, 4, hidden)
|
||
refusal_component = 5.0 * d.unsqueeze(0).unsqueeze(0).expand_as(batch)
|
||
input_hidden = batch + refusal_component
|
||
|
||
# Run through the layer (hook should fire)
|
||
output = layer(input_hidden)
|
||
|
||
# The refusal component should be reduced
|
||
proj_before = torch.einsum("bsh,h->bs", input_hidden, d).abs().mean()
|
||
proj_after = torch.einsum("bsh,h->bs", output, d).abs().mean()
|
||
assert proj_after < proj_before, (
|
||
f"Steering should reduce refusal projection: before={proj_before:.3f}, after={proj_after:.3f}"
|
||
)
|
||
|
||
# Cleanup
|
||
for hook in pipeline._steering_hooks:
|
||
hook.remove()
|
||
|
||
def test_expert_transplant(self):
|
||
"""Expert transplant should overwrite safety expert weights with capability average."""
|
||
hidden = 16
|
||
n_experts = 4
|
||
|
||
class FakeExpert(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.down_proj = torch.nn.Linear(hidden, hidden, bias=False)
|
||
|
||
class FakeMoE(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.gate = torch.nn.Linear(hidden, n_experts, bias=False)
|
||
self.experts = torch.nn.ModuleList([FakeExpert() for _ in range(n_experts)])
|
||
|
||
class FakeLayer(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.self_attn = torch.nn.Module()
|
||
self.self_attn.o_proj = torch.nn.Linear(hidden, hidden, bias=False)
|
||
self.mlp = FakeMoE()
|
||
|
||
layer = FakeLayer()
|
||
layers = torch.nn.ModuleList([layer])
|
||
torch.manual_seed(42)
|
||
for p in layer.parameters():
|
||
p.data = torch.randn_like(p.data)
|
||
|
||
# Save original safety expert weight
|
||
orig_safety0 = layer.mlp.experts[0].down_proj.weight.data.clone()
|
||
# Save capability expert weights for computing expected mean
|
||
# With top-third classification (n_experts // 3 = 1), only expert 0
|
||
# is safety; experts 1, 2, 3 are all capability.
|
||
cap1 = layer.mlp.experts[1].down_proj.weight.data.clone()
|
||
cap2 = layer.mlp.experts[2].down_proj.weight.data.clone()
|
||
cap3 = layer.mlp.experts[3].down_proj.weight.data.clone()
|
||
expected_mean = (cap1 + cap2 + cap3) / 3.0
|
||
|
||
import obliteratus.abliterate as abl_module
|
||
from obliteratus.models.loader import ModelHandle
|
||
from transformers import GPT2Config
|
||
|
||
config = GPT2Config(n_embd=hidden, n_head=2, n_layer=1, vocab_size=100, n_positions=64)
|
||
model = MagicMock()
|
||
model.parameters.return_value = iter([torch.zeros(1)])
|
||
handle = ModelHandle(model=model, tokenizer=MagicMock(), config=config, model_name="test", task="causal_lm")
|
||
|
||
pipeline = AbliterationPipeline(model_name="test", method="nuclear")
|
||
pipeline.handle = handle
|
||
pipeline._on_log = lambda m: None
|
||
pipeline._on_stage = lambda r: None
|
||
pipeline._strong_layers = [0]
|
||
# Experts 0,1 are safety (high affinity), 2,3 are capability
|
||
pipeline._expert_safety_scores = {
|
||
0: [(0, 5.0), (1, 3.0), (2, -1.0), (3, -3.0)]
|
||
}
|
||
|
||
orig_get_ffn = abl_module.get_ffn_module
|
||
abl_module.get_ffn_module = lambda lay, a: lay.mlp
|
||
try:
|
||
count = pipeline._transplant_expert_weights(layers)
|
||
finally:
|
||
abl_module.get_ffn_module = orig_get_ffn
|
||
|
||
assert count >= 1, f"Should blend at least 1 weight (top-third safety expert), got {count}"
|
||
|
||
# Safety expert 0 should be a 10% blend toward capability mean
|
||
# (nuclear default transplant_blend=0.10)
|
||
# new = 0.90 * original + 0.10 * capability_mean
|
||
blend = pipeline.transplant_blend # 0.10
|
||
expected_blend = (1.0 - blend) * orig_safety0 + blend * expected_mean
|
||
transplanted = layer.mlp.experts[0].down_proj.weight.data
|
||
assert torch.allclose(transplanted, expected_blend, atol=1e-4), (
|
||
f"Safety expert weight should be {blend:.0%} blended toward capability mean"
|
||
)
|
||
|
||
# Capability expert 2 should be unchanged
|
||
assert torch.allclose(layer.mlp.experts[2].down_proj.weight.data, cap2, atol=1e-6), (
|
||
"Capability expert should be unchanged"
|
||
)
|
||
|
||
def test_gather_state_dict_raises_on_unmaterialized_meta(self):
|
||
"""Should raise RuntimeError when meta tensors survive materialization."""
|
||
from obliteratus.models.loader import ModelHandle
|
||
from transformers import GPT2Config
|
||
|
||
config = GPT2Config(n_embd=8, n_head=2, n_layer=1, vocab_size=100, n_positions=64)
|
||
|
||
# Create a fake model whose state_dict returns a meta tensor.
|
||
# Patch get_state_dict_offloaded_model to return the same meta tensor
|
||
# (simulating a failed materialization).
|
||
fake_model = MagicMock()
|
||
meta_tensor = torch.empty(4, 8, device="meta")
|
||
fake_model.state_dict.return_value = {"layer.weight": meta_tensor}
|
||
|
||
handle = ModelHandle(
|
||
model=fake_model, tokenizer=MagicMock(), config=config,
|
||
model_name="test", task="causal_lm",
|
||
)
|
||
|
||
pipeline = AbliterationPipeline(model_name="test", method="nuclear")
|
||
pipeline.handle = handle
|
||
pipeline._on_log = lambda m: None
|
||
pipeline._on_stage = lambda r: None
|
||
|
||
# Patch accelerate's function to return meta tensors (simulates
|
||
# materialization failure — the safety net should catch this).
|
||
with patch(
|
||
"accelerate.utils.modeling.get_state_dict_offloaded_model",
|
||
return_value={"layer.weight": meta_tensor},
|
||
):
|
||
with pytest.raises(RuntimeError, match="bricked checkpoint"):
|
||
pipeline._gather_state_dict()
|
||
|
||
def test_gather_state_dict_materializes_offloaded_meta(self):
|
||
"""Should return real CPU tensors from Accelerate materialization."""
|
||
from obliteratus.models.loader import ModelHandle
|
||
from transformers import GPT2Config
|
||
|
||
config = GPT2Config(n_embd=8, n_head=2, n_layer=1, vocab_size=100, n_positions=64)
|
||
fake_model = MagicMock()
|
||
fake_model.state_dict.return_value = {
|
||
"layer.weight": torch.empty(4, 8, device="meta"),
|
||
}
|
||
fake_model.parameters.return_value = []
|
||
fake_model.buffers.return_value = []
|
||
materialized = torch.zeros(4, 8)
|
||
|
||
handle = ModelHandle(
|
||
model=fake_model, tokenizer=MagicMock(), config=config,
|
||
model_name="test", task="causal_lm",
|
||
)
|
||
pipeline = AbliterationPipeline(model_name="test", method="nuclear")
|
||
pipeline.handle = handle
|
||
pipeline._on_log = lambda m: None
|
||
pipeline._on_stage = lambda r: None
|
||
|
||
with patch(
|
||
"accelerate.utils.modeling.get_state_dict_offloaded_model",
|
||
return_value={"layer.weight": materialized},
|
||
) as gather_offloaded:
|
||
state_dict = pipeline._gather_state_dict()
|
||
|
||
gather_offloaded.assert_called_once_with(fake_model)
|
||
assert state_dict["layer.weight"] is materialized
|
||
assert state_dict["layer.weight"].device.type == "cpu"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Knee detection
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestKneeDetection:
|
||
def test_empty_input(self):
|
||
result = AbliterationPipeline._select_layers_knee([])
|
||
assert result == []
|
||
|
||
def test_two_layers(self):
|
||
result = AbliterationPipeline._select_layers_knee([(0, 5.0), (1, 3.0)])
|
||
assert set(result) == {0, 1}
|
||
|
||
def test_clear_knee(self):
|
||
"""Layers with a sharp dropoff should be separated by knee detection."""
|
||
sorted_layers = [
|
||
(14, 10.0), (15, 9.5), (13, 9.0), # strong cluster
|
||
(16, 2.0), (12, 1.5), (17, 1.0), (11, 0.5), (18, 0.2), (10, 0.1),
|
||
]
|
||
result = AbliterationPipeline._select_layers_knee(sorted_layers)
|
||
# Should select the strong cluster (layers 14, 15, 13) and exclude weak ones
|
||
assert 14 in result
|
||
assert 15 in result
|
||
assert 13 in result
|
||
assert len(result) <= 5 # shouldn't select all 9
|
||
|
||
def test_minimum_threshold_filters_noise(self):
|
||
"""Layers below 10% of max should be filtered out."""
|
||
sorted_layers = [(0, 10.0), (1, 0.5)] # 0.5 is 5% of 10
|
||
result = AbliterationPipeline._select_layers_knee(sorted_layers)
|
||
# Layer 1 is below 10% threshold
|
||
assert 0 in result
|
||
|
||
def test_all_equal_norms(self):
|
||
"""When all norms are equal, should select all (or most)."""
|
||
sorted_layers = [(i, 5.0) for i in range(5)]
|
||
result = AbliterationPipeline._select_layers_knee(sorted_layers)
|
||
assert len(result) >= 1
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Activation collection
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestActivationCollection:
|
||
def test_collect_activations(self, handle):
|
||
"""Test that activation collection returns correct structure."""
|
||
from obliteratus.strategies.utils import get_layer_modules
|
||
|
||
pipeline = AbliterationPipeline(model_name="test")
|
||
pipeline.handle = handle
|
||
pipeline._on_log = lambda m: None
|
||
|
||
layers = get_layer_modules(handle)
|
||
prompts = ["Hello world", "Test prompt"]
|
||
|
||
_make_varied_tokenizer(handle)
|
||
|
||
activations = pipeline._collect_activations(layers, prompts, "test")
|
||
|
||
assert len(activations) == len(layers)
|
||
for idx in range(len(layers)):
|
||
assert len(activations[idx]) == len(prompts)
|
||
for act in activations[idx]:
|
||
assert act.device == torch.device("cpu")
|
||
assert act.shape[-1] == handle.hidden_size
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Distill: single direction (basic method)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestDistillBasic:
|
||
def test_single_direction(self, handle):
|
||
"""Basic method: single refusal direction via difference-in-means."""
|
||
from obliteratus.strategies.utils import get_layer_modules
|
||
|
||
pipeline = AbliterationPipeline(
|
||
model_name="test",
|
||
method="basic",
|
||
harmful_prompts=["bad prompt"],
|
||
harmless_prompts=["good prompt"],
|
||
)
|
||
pipeline.handle = handle
|
||
pipeline._on_log = lambda m: None
|
||
pipeline._on_stage = lambda r: None
|
||
_make_varied_tokenizer(handle)
|
||
|
||
pipeline._probe()
|
||
pipeline._distill()
|
||
|
||
n_layers = len(get_layer_modules(handle))
|
||
assert len(pipeline.refusal_directions) == n_layers
|
||
for idx, direction in pipeline.refusal_directions.items():
|
||
assert abs(direction.norm().item() - 1.0) < 1e-4
|
||
# Single direction: subspace should be (1, hidden_dim)
|
||
assert pipeline.refusal_subspaces[idx].shape[0] == 1
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Distill: multi-direction SVD (advanced/aggressive method)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestDistillSVD:
|
||
def test_multi_direction_svd(self, handle):
|
||
"""Advanced method: SVD extracts multiple refusal directions.
|
||
|
||
Note: on small models (hidden_size < 2048 or < 2B params), n_directions
|
||
is automatically capped to 2 to prevent over-ablation. The test model
|
||
(hidden_size=64, 4 layers) triggers this safeguard.
|
||
"""
|
||
from obliteratus.strategies.utils import get_layer_modules
|
||
|
||
pipeline = AbliterationPipeline(
|
||
model_name="test",
|
||
method="advanced",
|
||
harmful_prompts=["bad1", "bad2", "bad3", "bad4", "bad5"],
|
||
harmless_prompts=["good1", "good2", "good3", "good4", "good5"],
|
||
)
|
||
pipeline.handle = handle
|
||
pipeline._on_log = lambda m: None
|
||
pipeline._on_stage = lambda r: None
|
||
_make_varied_tokenizer(handle)
|
||
|
||
pipeline._probe()
|
||
pipeline._distill()
|
||
|
||
n_layers = len(get_layer_modules(handle))
|
||
assert len(pipeline.refusal_subspaces) == n_layers
|
||
# Small-model cap: n_directions capped to 2 for tiny test model
|
||
expected_dirs = min(2, pipeline.n_directions, 5, handle.hidden_size)
|
||
for idx, subspace in pipeline.refusal_subspaces.items():
|
||
assert subspace.shape[0] == expected_dirs
|
||
assert subspace.shape[1] == handle.hidden_size
|
||
|
||
# Primary direction should still be a unit vector
|
||
for idx, direction in pipeline.refusal_directions.items():
|
||
assert abs(direction.norm().item() - 1.0) < 1e-4
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Full pipeline: excise with different methods
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestExcise:
|
||
def test_qwen25_coder_advanced_excise_projects_all_dense_weights(self):
|
||
"""Qwen2.5 must not silently resolve to zero projection targets."""
|
||
config = Qwen2Config(
|
||
vocab_size=128,
|
||
hidden_size=32,
|
||
intermediate_size=64,
|
||
num_hidden_layers=2,
|
||
num_attention_heads=4,
|
||
num_key_value_heads=2,
|
||
max_position_embeddings=64,
|
||
)
|
||
model = Qwen2ForCausalLM(config)
|
||
handle = ModelHandle(
|
||
model=model,
|
||
tokenizer=MagicMock(),
|
||
config=config,
|
||
model_name="Qwen/Qwen2.5-Coder-7B-Instruct",
|
||
task="causal_lm",
|
||
)
|
||
pipeline = AbliterationPipeline(
|
||
model_name=handle.model_name,
|
||
method="advanced",
|
||
refinement_passes=1,
|
||
norm_preserve=False,
|
||
project_biases=False,
|
||
layer_adaptive_strength=False,
|
||
safety_neuron_masking=False,
|
||
attention_head_surgery=False,
|
||
use_sae_features=False,
|
||
use_kl_optimization=False,
|
||
projection_target="all",
|
||
harmful_prompts=["harmful"],
|
||
harmless_prompts=["harmless"],
|
||
)
|
||
pipeline.handle = handle
|
||
pipeline._strong_layers = [0]
|
||
direction = torch.arange(1, config.hidden_size + 1, dtype=torch.float32)
|
||
direction /= direction.norm()
|
||
pipeline.refusal_directions = {0: direction.clone()}
|
||
pipeline.refusal_subspaces = {0: direction.unsqueeze(0)}
|
||
|
||
layer = model.model.layers[0]
|
||
original_weights = {
|
||
name: parameter.detach().clone()
|
||
for name, parameter in layer.named_parameters()
|
||
}
|
||
|
||
pipeline._excise()
|
||
|
||
expected = {
|
||
"self_attn.q_proj.weight",
|
||
"self_attn.k_proj.weight",
|
||
"self_attn.v_proj.weight",
|
||
"self_attn.o_proj.weight",
|
||
"mlp.gate_proj.weight",
|
||
"mlp.up_proj.weight",
|
||
"mlp.down_proj.weight",
|
||
}
|
||
changed = {
|
||
name
|
||
for name, parameter in layer.named_parameters()
|
||
if not torch.equal(original_weights[name], parameter)
|
||
}
|
||
assert handle.architecture == "qwen2"
|
||
assert changed == expected
|
||
assert pipeline._excise_modified_count == 8 # seven layer weights + lm_head
|
||
|
||
def test_excise_basic(self, handle):
|
||
"""Basic method should modify weights."""
|
||
from obliteratus.strategies.utils import get_layer_modules
|
||
|
||
pipeline = AbliterationPipeline(
|
||
model_name="test",
|
||
method="basic",
|
||
harmful_prompts=["bad prompt"],
|
||
harmless_prompts=["good prompt"],
|
||
)
|
||
pipeline.handle = handle
|
||
pipeline._on_log = lambda m: None
|
||
pipeline._on_stage = lambda r: None
|
||
_make_varied_tokenizer(handle)
|
||
|
||
layers = get_layer_modules(handle)
|
||
original_weights = {}
|
||
for idx in range(len(layers)):
|
||
for name, param in layers[idx].named_parameters():
|
||
original_weights[(idx, name)] = param.data.clone()
|
||
|
||
pipeline._probe()
|
||
pipeline._distill()
|
||
pipeline._excise()
|
||
|
||
any_changed = False
|
||
for idx in range(len(layers)):
|
||
for name, param in layers[idx].named_parameters():
|
||
if not torch.allclose(original_weights[(idx, name)], param.data, atol=1e-6):
|
||
any_changed = True
|
||
break
|
||
|
||
assert any_changed, "Excise should modify at least some weights"
|
||
|
||
def test_excise_advanced_norm_preserving(self, handle):
|
||
"""Advanced method with norm preservation should maintain weight norms."""
|
||
from obliteratus.strategies.utils import get_layer_modules
|
||
|
||
pipeline = AbliterationPipeline(
|
||
model_name="test",
|
||
method="advanced",
|
||
harmful_prompts=["bad prompt"],
|
||
harmless_prompts=["good prompt"],
|
||
)
|
||
pipeline.handle = handle
|
||
pipeline._on_log = lambda m: None
|
||
pipeline._on_stage = lambda r: None
|
||
_make_varied_tokenizer(handle)
|
||
|
||
get_layer_modules(handle)
|
||
|
||
pipeline._probe()
|
||
pipeline._distill()
|
||
pipeline._excise()
|
||
|
||
# Weights should have been modified (advanced uses _project_out_advanced)
|
||
assert len(pipeline._strong_layers) > 0
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Rebirth (save)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestRebirth:
|
||
def test_rebirth_saves_metadata(self, handle, tmp_path):
|
||
"""Rebirth should save model and comprehensive metadata JSON."""
|
||
pipeline = AbliterationPipeline(
|
||
model_name="test-model",
|
||
output_dir=str(tmp_path / "output"),
|
||
method="advanced",
|
||
)
|
||
pipeline.handle = handle
|
||
pipeline._on_log = lambda m: None
|
||
pipeline._on_stage = lambda r: None
|
||
pipeline._strong_layers = [0]
|
||
pipeline._quality_metrics = {"perplexity": 8.5, "coherence": 1.0}
|
||
|
||
handle.model.save_pretrained = MagicMock(
|
||
side_effect=lambda path, **_kwargs: (
|
||
(Path(path) / "config.json").write_text("{}", encoding="utf-8"),
|
||
(Path(path) / "model.safetensors").write_bytes(b"weights"),
|
||
),
|
||
)
|
||
handle.tokenizer.save_pretrained = MagicMock(
|
||
side_effect=lambda path: (Path(path) / "tokenizer_config.json").write_text(
|
||
"{}",
|
||
encoding="utf-8",
|
||
),
|
||
)
|
||
|
||
result_path = pipeline._rebirth()
|
||
|
||
assert result_path == tmp_path / "output"
|
||
assert (result_path / "abliteration_metadata.json").exists()
|
||
|
||
metadata = json.loads((result_path / "abliteration_metadata.json").read_text())
|
||
assert metadata["source_model"] == "test-model"
|
||
assert metadata["technique"] == "refusal_direction_ablation"
|
||
assert metadata["method"] == "advanced"
|
||
assert metadata["strong_layers"] == [0]
|
||
assert "method_config" in metadata
|
||
assert metadata["method_config"]["n_directions"] == METHODS["advanced"]["n_directions"]
|
||
assert metadata["method_config"]["norm_preserve"] is True
|
||
assert "references" in metadata
|
||
assert len(metadata["references"]) >= 3
|
||
assert "quality_metrics" in metadata
|
||
assert metadata["quality_metrics"]["perplexity"] == 8.5
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# CLI integration
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestCLI:
|
||
def test_abliterate_parser_with_method(self):
|
||
"""Test that the abliterate subcommand parses method correctly."""
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser()
|
||
subparsers = parser.add_subparsers(dest="command")
|
||
abl_parser = subparsers.add_parser("abliterate")
|
||
abl_parser.add_argument("model", type=str)
|
||
abl_parser.add_argument("--output-dir", type=str, default=None)
|
||
abl_parser.add_argument("--device", type=str, default="auto")
|
||
abl_parser.add_argument("--dtype", type=str, default="float16")
|
||
abl_parser.add_argument("--method", type=str, default="advanced",
|
||
choices=["basic", "advanced", "aggressive"])
|
||
abl_parser.add_argument("--n-directions", type=int, default=None)
|
||
abl_parser.add_argument("--regularization", type=float, default=None)
|
||
abl_parser.add_argument("--refinement-passes", type=int, default=None)
|
||
|
||
args = parser.parse_args(["abliterate", "gpt2", "--method", "aggressive", "--n-directions", "6"])
|
||
assert args.command == "abliterate"
|
||
assert args.model == "gpt2"
|
||
assert args.method == "aggressive"
|
||
assert args.n_directions == 6
|
||
assert args.dtype == "float16"
|
||
|
||
def test_default_method(self):
|
||
"""Default method should be advanced."""
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser()
|
||
subparsers = parser.add_subparsers(dest="command")
|
||
abl_parser = subparsers.add_parser("abliterate")
|
||
abl_parser.add_argument("model", type=str)
|
||
abl_parser.add_argument("--method", type=str, default="advanced")
|
||
|
||
args = parser.parse_args(["abliterate", "gpt2"])
|
||
assert args.method == "advanced"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Expert-Granular Abliteration (EGA)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TestFindRouterModule:
|
||
"""Test _find_router_module static method."""
|
||
|
||
def test_finds_gate(self):
|
||
"""Should find a router named 'gate'."""
|
||
hidden = 16
|
||
|
||
class FakeMoE(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.gate = torch.nn.Linear(hidden, 4, bias=False)
|
||
self.experts = torch.nn.ModuleList()
|
||
|
||
moe = FakeMoE()
|
||
router = AbliterationPipeline._find_router_module(moe)
|
||
assert router is moe.gate
|
||
|
||
def test_finds_router(self):
|
||
"""Should find a router named 'router'."""
|
||
hidden = 16
|
||
|
||
class FakeMoE(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.router = torch.nn.Linear(hidden, 4, bias=False)
|
||
self.experts = torch.nn.ModuleList()
|
||
|
||
moe = FakeMoE()
|
||
router = AbliterationPipeline._find_router_module(moe)
|
||
assert router is moe.router
|
||
|
||
def test_auto_detects_unknown_router(self):
|
||
"""Should auto-detect a router with unusual name via heuristic."""
|
||
hidden = 16
|
||
|
||
class FakeMoE(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.moe_gate_proj = torch.nn.Linear(hidden, 4, bias=False)
|
||
self.experts = torch.nn.ModuleList()
|
||
|
||
moe = FakeMoE()
|
||
router = AbliterationPipeline._find_router_module(moe)
|
||
assert router is moe.moe_gate_proj
|
||
|
||
def test_returns_none_no_router(self):
|
||
"""Should return None when no router is found."""
|
||
class NoRouter(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.linear = torch.nn.Linear(16, 16)
|
||
|
||
mod = NoRouter()
|
||
assert AbliterationPipeline._find_router_module(mod) is None
|
||
|
||
|
||
class TestRouterProfilingHooks:
|
||
"""Test _install_router_profiling_hooks."""
|
||
|
||
def _make_moe_pipeline_and_layers(self, hidden=16, n_experts=4):
|
||
"""Create a pipeline with a fake MoE model for router profiling tests."""
|
||
from obliteratus.models.loader import ModelHandle
|
||
from transformers import GPT2Config
|
||
|
||
class FakeExpert(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.down_proj = torch.nn.Linear(hidden, hidden, bias=False)
|
||
|
||
class FakeMoE(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.gate = torch.nn.Linear(hidden, n_experts, bias=False)
|
||
self.experts = torch.nn.ModuleList([FakeExpert() for _ in range(n_experts)])
|
||
|
||
def forward(self, x):
|
||
return x
|
||
|
||
class FakeLayer(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.self_attn = torch.nn.Module()
|
||
self.self_attn.o_proj = torch.nn.Linear(hidden, hidden, bias=False)
|
||
self.mlp = FakeMoE()
|
||
|
||
def forward(self, x):
|
||
return (x,)
|
||
|
||
config = GPT2Config(n_embd=hidden, n_head=2, n_layer=1, vocab_size=100, n_positions=64)
|
||
model = MagicMock()
|
||
model.parameters.return_value = iter([torch.zeros(1)])
|
||
handle = ModelHandle(model=model, tokenizer=MagicMock(), config=config, model_name="test", task="causal_lm")
|
||
|
||
pipeline = AbliterationPipeline(model_name="test", method="surgical")
|
||
pipeline.handle = handle
|
||
pipeline._on_log = lambda m: None
|
||
pipeline._on_stage = lambda r: None
|
||
|
||
layer = FakeLayer()
|
||
layers = torch.nn.ModuleList([layer])
|
||
|
||
# Monkey-patch get_ffn_module
|
||
import obliteratus.abliterate as abl_module
|
||
orig_get_ffn = abl_module.get_ffn_module
|
||
abl_module.get_ffn_module = lambda lay, a: lay.mlp
|
||
|
||
return pipeline, layers, layer, abl_module, orig_get_ffn
|
||
|
||
def test_hooks_installed(self):
|
||
"""Should install hooks on MoE router modules."""
|
||
pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers()
|
||
try:
|
||
hooks = pipeline._install_router_profiling_hooks(layers)
|
||
assert len(hooks) == 1
|
||
assert 0 in pipeline._routing_harmful
|
||
assert 0 in pipeline._routing_harmless
|
||
finally:
|
||
for h in hooks:
|
||
h.remove()
|
||
abl_module.get_ffn_module = orig_get_ffn
|
||
|
||
def test_hooks_record_logits(self):
|
||
"""Hooks should record router logits during forward passes."""
|
||
pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers()
|
||
try:
|
||
hooks = pipeline._install_router_profiling_hooks(layers)
|
||
|
||
# Simulate harmful forward pass
|
||
pipeline._routing_is_harmful = True
|
||
x = torch.randn(1, 5, 16)
|
||
harmful_expected = torch.nn.functional.linear(x, layer.mlp.gate.weight)[0, -1]
|
||
layer.mlp.gate(x) # triggers hook
|
||
|
||
assert len(pipeline._routing_harmful[0]) == 1
|
||
assert pipeline._routing_harmful[0][0].shape[0] == 4 # n_experts
|
||
assert torch.equal(pipeline._routing_harmful[0][0], harmful_expected)
|
||
|
||
# Simulate harmless forward pass
|
||
pipeline._routing_is_harmful = False
|
||
harmless_x = x + 1
|
||
harmless_expected = torch.nn.functional.linear(
|
||
harmless_x, layer.mlp.gate.weight,
|
||
)[0, -1]
|
||
layer.mlp.gate(harmless_x)
|
||
|
||
assert len(pipeline._routing_harmless[0]) == 1
|
||
assert torch.equal(pipeline._routing_harmless[0][0], harmless_expected)
|
||
assert not torch.equal(
|
||
pipeline._routing_harmful[0][0], pipeline._routing_harmless[0][0],
|
||
)
|
||
finally:
|
||
for h in hooks:
|
||
h.remove()
|
||
abl_module.get_ffn_module = orig_get_ffn
|
||
|
||
def test_hooks_record_single_prompt_vector_output(self):
|
||
"""A one-dimensional single-prompt router output should remain intact."""
|
||
pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers()
|
||
hooks = []
|
||
try:
|
||
hooks = pipeline._install_router_profiling_hooks(layers)
|
||
x = torch.arange(16, dtype=torch.float32) / 100
|
||
expected = torch.nn.functional.linear(x, layer.mlp.gate.weight)
|
||
|
||
layer.mlp.gate(x)
|
||
|
||
assert len(pipeline._routing_harmful[0]) == 1
|
||
assert torch.equal(pipeline._routing_harmful[0][0], expected)
|
||
finally:
|
||
for h in hooks:
|
||
h.remove()
|
||
abl_module.get_ffn_module = orig_get_ffn
|
||
|
||
def test_hooks_record_one_vector_for_each_3d_batch_prompt(self):
|
||
"""Native 3-D router outputs should retain every batch member."""
|
||
pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers()
|
||
hooks = []
|
||
try:
|
||
hooks = pipeline._install_router_profiling_hooks(layers)
|
||
x = torch.arange(2 * 5 * 16, dtype=torch.float32).reshape(2, 5, 16) / 100
|
||
expected = torch.nn.functional.linear(x, layer.mlp.gate.weight)
|
||
|
||
layer.mlp.gate(x)
|
||
|
||
recorded = pipeline._routing_harmful[0]
|
||
assert len(recorded) == 2
|
||
assert torch.equal(recorded[0], expected[0, -1])
|
||
assert torch.equal(recorded[1], expected[1, -1])
|
||
finally:
|
||
for h in hooks:
|
||
h.remove()
|
||
abl_module.get_ffn_module = orig_get_ffn
|
||
|
||
def test_hooks_preserve_flattened_prompt_association_with_padding(self):
|
||
"""Flattened token rows should become one vector per padded prompt."""
|
||
pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers()
|
||
hooks = []
|
||
try:
|
||
hooks = pipeline._install_router_profiling_hooks(layers)
|
||
pipeline._routing_attention_mask = torch.tensor(
|
||
[[1, 1, 0, 0], [1, 1, 1, 1]], dtype=torch.long,
|
||
)
|
||
x = torch.arange(8 * 16, dtype=torch.float32).reshape(8, 16) / 100
|
||
expected = torch.nn.functional.linear(x, layer.mlp.gate.weight).reshape(2, 4, 4)
|
||
|
||
layer.mlp.gate(x)
|
||
|
||
recorded = pipeline._routing_harmful[0]
|
||
assert len(recorded) == 2
|
||
assert torch.equal(recorded[0], expected[0, 1])
|
||
assert torch.equal(recorded[1], expected[1, 3])
|
||
finally:
|
||
pipeline._routing_attention_mask = None
|
||
for h in hooks:
|
||
h.remove()
|
||
abl_module.get_ffn_module = orig_get_ffn
|
||
|
||
def test_hooks_record_single_prompt_flattened_tokens(self):
|
||
"""A flattened single-prompt sequence should select its final token."""
|
||
pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers()
|
||
hooks = []
|
||
try:
|
||
hooks = pipeline._install_router_profiling_hooks(layers)
|
||
pipeline._routing_attention_mask = torch.ones(1, 4, dtype=torch.long)
|
||
x = torch.arange(4 * 16, dtype=torch.float32).reshape(4, 16) / 100
|
||
expected = torch.nn.functional.linear(x, layer.mlp.gate.weight)[-1]
|
||
|
||
layer.mlp.gate(x)
|
||
|
||
assert len(pipeline._routing_harmful[0]) == 1
|
||
assert torch.equal(pipeline._routing_harmful[0][0], expected)
|
||
finally:
|
||
pipeline._routing_attention_mask = None
|
||
for h in hooks:
|
||
h.remove()
|
||
abl_module.get_ffn_module = orig_get_ffn
|
||
|
||
def test_hooks_average_only_valid_cot_tokens_per_prompt(self):
|
||
"""CoT aggregation should exclude padding without mixing prompts."""
|
||
pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers()
|
||
hooks = []
|
||
try:
|
||
hooks = pipeline._install_router_profiling_hooks(layers)
|
||
pipeline.cot_aware = True
|
||
pipeline._routing_attention_mask = torch.tensor(
|
||
[[1, 1, 1, 1, 1, 0], [0, 1, 1, 1, 1, 1]], dtype=torch.long,
|
||
)
|
||
x = torch.arange(12 * 16, dtype=torch.float32).reshape(12, 16) / 100
|
||
expected = torch.nn.functional.linear(x, layer.mlp.gate.weight).reshape(2, 6, 4)
|
||
|
||
layer.mlp.gate(x)
|
||
|
||
recorded = pipeline._routing_harmful[0]
|
||
assert len(recorded) == 2
|
||
assert torch.allclose(recorded[0], expected[0, :5].mean(dim=0))
|
||
assert torch.allclose(recorded[1], expected[1, 1:].mean(dim=0))
|
||
finally:
|
||
pipeline._routing_attention_mask = None
|
||
for h in hooks:
|
||
h.remove()
|
||
abl_module.get_ffn_module = orig_get_ffn
|
||
|
||
def test_hooks_skip_ambiguous_flattened_logits(self):
|
||
"""Ambiguous flattened layouts should fail visibly instead of collapsing."""
|
||
pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers()
|
||
hooks = []
|
||
try:
|
||
hooks = pipeline._install_router_profiling_hooks(layers)
|
||
|
||
with pytest.warns(RuntimeWarning, match="multi-row 2-D router logits are ambiguous"):
|
||
layer.mlp.gate(torch.randn(5, 16))
|
||
|
||
assert pipeline._routing_harmful[0] == []
|
||
finally:
|
||
for h in hooks:
|
||
h.remove()
|
||
abl_module.get_ffn_module = orig_get_ffn
|
||
|
||
def test_hooks_skip_mismatched_flattened_logits(self):
|
||
"""A stale or incompatible mask should produce an actionable warning."""
|
||
pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers()
|
||
hooks = []
|
||
try:
|
||
hooks = pipeline._install_router_profiling_hooks(layers)
|
||
pipeline._routing_attention_mask = torch.ones(2, 4, dtype=torch.long)
|
||
|
||
with pytest.warns(RuntimeWarning, match="do not match batch × sequence"):
|
||
layer.mlp.gate(torch.randn(7, 16))
|
||
|
||
assert pipeline._routing_harmful[0] == []
|
||
finally:
|
||
pipeline._routing_attention_mask = None
|
||
for h in hooks:
|
||
h.remove()
|
||
abl_module.get_ffn_module = orig_get_ffn
|
||
|
||
def test_no_handle_returns_empty(self):
|
||
"""Should return empty list when handle is None."""
|
||
pipeline = AbliterationPipeline(model_name="test", method="surgical")
|
||
pipeline.handle = None
|
||
hooks = pipeline._install_router_profiling_hooks(torch.nn.ModuleList())
|
||
assert hooks == []
|
||
|
||
|
||
class TestComputeExpertGranularDirections:
|
||
"""Test _compute_expert_granular_directions."""
|
||
|
||
def test_computes_per_expert_directions(self):
|
||
"""Should compute per-expert refusal directions from routing data."""
|
||
hidden = 16
|
||
n_experts = 4
|
||
|
||
pipeline = AbliterationPipeline(model_name="test", method="surgical")
|
||
pipeline._on_log = lambda m: None
|
||
pipeline._on_stage = lambda r: None
|
||
pipeline._strong_layers = [0]
|
||
|
||
torch.manual_seed(42)
|
||
|
||
# Simulate router logits: expert 0 favored for harmful, expert 3 for harmless
|
||
h_logits = []
|
||
s_logits = []
|
||
for _ in range(10):
|
||
hl = torch.randn(n_experts)
|
||
hl[0] += 2.0 # bias expert 0 for harmful
|
||
h_logits.append(hl)
|
||
sl = torch.randn(n_experts)
|
||
sl[3] += 2.0 # bias expert 3 for harmless
|
||
s_logits.append(sl)
|
||
|
||
pipeline._routing_harmful = {0: h_logits}
|
||
pipeline._routing_harmless = {0: s_logits}
|
||
|
||
# Simulate per-prompt activations with harmful/harmless separation
|
||
refusal_dir = torch.randn(hidden)
|
||
refusal_dir = refusal_dir / refusal_dir.norm()
|
||
|
||
h_acts = [torch.randn(hidden) + 1.5 * refusal_dir for _ in range(10)]
|
||
s_acts = [torch.randn(hidden) - 1.5 * refusal_dir for _ in range(10)]
|
||
pipeline._harmful_acts = {0: h_acts}
|
||
pipeline._harmless_acts = {0: s_acts}
|
||
|
||
pipeline._compute_expert_granular_directions()
|
||
|
||
# Should have computed expert directions for layer 0
|
||
assert 0 in pipeline._expert_directions
|
||
assert len(pipeline._expert_directions[0]) > 0
|
||
|
||
# Should have dynamic safety scores
|
||
assert 0 in pipeline._expert_safety_scores
|
||
scores = pipeline._expert_safety_scores[0]
|
||
assert len(scores) == n_experts
|
||
# Expert 0 should have higher safety score (more activated for harmful)
|
||
expert_0_score = next(s for eid, s in scores if eid == 0)
|
||
expert_3_score = next(s for eid, s in scores if eid == 3)
|
||
assert expert_0_score > expert_3_score, (
|
||
f"Expert 0 should have higher safety score: {expert_0_score} vs {expert_3_score}"
|
||
)
|
||
|
||
def test_directions_are_unit_vectors(self):
|
||
"""Per-expert directions should be unit normalized."""
|
||
hidden = 16
|
||
n_experts = 4
|
||
|
||
pipeline = AbliterationPipeline(model_name="test", method="surgical")
|
||
pipeline._on_log = lambda m: None
|
||
pipeline._strong_layers = [0]
|
||
|
||
torch.manual_seed(42)
|
||
h_logits = [torch.randn(n_experts) for _ in range(10)]
|
||
s_logits = [torch.randn(n_experts) for _ in range(10)]
|
||
pipeline._routing_harmful = {0: h_logits}
|
||
pipeline._routing_harmless = {0: s_logits}
|
||
pipeline._harmful_acts = {0: [torch.randn(hidden) + torch.ones(hidden) for _ in range(10)]}
|
||
pipeline._harmless_acts = {0: [torch.randn(hidden) - torch.ones(hidden) for _ in range(10)]}
|
||
|
||
pipeline._compute_expert_granular_directions()
|
||
|
||
if 0 in pipeline._expert_directions:
|
||
for ei, d in pipeline._expert_directions[0].items():
|
||
assert abs(d.norm().item() - 1.0) < 1e-4, (
|
||
f"Expert {ei} direction norm={d.norm().item()}, expected 1.0"
|
||
)
|
||
|
||
def test_skips_when_no_routing_data(self):
|
||
"""Should skip gracefully when no routing data is available."""
|
||
pipeline = AbliterationPipeline(model_name="test", method="surgical")
|
||
pipeline._on_log = lambda m: None
|
||
pipeline._routing_harmful = {}
|
||
pipeline._routing_harmless = {}
|
||
|
||
pipeline._compute_expert_granular_directions()
|
||
|
||
assert len(pipeline._expert_directions) == 0
|
||
|
||
def test_skips_expert_with_low_routing_weight(self):
|
||
"""Experts with insufficient routing weight should not get directions."""
|
||
hidden = 16
|
||
|
||
pipeline = AbliterationPipeline(model_name="test", method="surgical")
|
||
pipeline._on_log = lambda m: None
|
||
pipeline._strong_layers = [0]
|
||
|
||
# Create routing logits where expert 3 is never selected (very low)
|
||
h_logits = []
|
||
s_logits = []
|
||
for _ in range(3):
|
||
hl = torch.tensor([5.0, 5.0, 5.0, -100.0]) # expert 3 never routed
|
||
h_logits.append(hl)
|
||
sl = torch.tensor([5.0, 5.0, 5.0, -100.0])
|
||
s_logits.append(sl)
|
||
|
||
pipeline._routing_harmful = {0: h_logits}
|
||
pipeline._routing_harmless = {0: s_logits}
|
||
|
||
torch.manual_seed(42)
|
||
pipeline._harmful_acts = {0: [torch.randn(hidden) for _ in range(3)]}
|
||
pipeline._harmless_acts = {0: [torch.randn(hidden) for _ in range(3)]}
|
||
|
||
pipeline._compute_expert_granular_directions()
|
||
|
||
# Expert 3 should NOT have a direction (routing weight too low)
|
||
if 0 in pipeline._expert_directions:
|
||
assert 3 not in pipeline._expert_directions[0], (
|
||
"Expert with near-zero routing weight should not get a direction"
|
||
)
|
||
|
||
|
||
class TestProjectMoEExpertsGranular:
|
||
"""Test _project_moe_experts_granular (ModuleList path)."""
|
||
|
||
def _make_direction(self, hidden_dim=16):
|
||
d = torch.randn(hidden_dim, 1)
|
||
return d / d.norm()
|
||
|
||
def test_per_expert_directions_applied(self):
|
||
"""Each expert should use its own direction when available."""
|
||
hidden = 16
|
||
n_experts = 4
|
||
|
||
class FakeExpert(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.down_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
self.up_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
|
||
class FakeMoE(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.gate = torch.nn.Linear(hidden, n_experts, bias=False)
|
||
self.experts = torch.nn.ModuleList([FakeExpert() for _ in range(n_experts)])
|
||
|
||
moe = FakeMoE()
|
||
torch.manual_seed(42)
|
||
for p in moe.parameters():
|
||
p.data = torch.randn_like(p.data)
|
||
|
||
shared_dir = self._make_direction(hidden)
|
||
|
||
# Create distinct per-expert directions
|
||
expert_dirs = {}
|
||
for ei in range(n_experts):
|
||
d = torch.randn(hidden)
|
||
d = d / d.norm()
|
||
expert_dirs[ei] = d
|
||
|
||
pipeline = AbliterationPipeline(model_name="test", method="surgical")
|
||
pipeline._on_log = lambda m: None
|
||
pipeline._expert_directions = {0: expert_dirs}
|
||
|
||
# Save originals
|
||
orig_weights = {
|
||
ei: moe.experts[ei].down_proj.weight.data.clone()
|
||
for ei in range(n_experts)
|
||
}
|
||
|
||
count = pipeline._project_moe_experts_granular(
|
||
moe, shared_dir, layer_idx=0,
|
||
)
|
||
|
||
assert count > 0, "Should project some weights"
|
||
|
||
# All experts should be modified
|
||
for ei in range(n_experts):
|
||
assert not torch.allclose(
|
||
moe.experts[ei].down_proj.weight.data, orig_weights[ei]
|
||
), f"Expert {ei} should be modified"
|
||
|
||
def test_falls_back_to_shared_direction(self):
|
||
"""Experts without per-expert direction should use shared direction."""
|
||
hidden = 16
|
||
n_experts = 4
|
||
|
||
class FakeExpert(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.down_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
self.up_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
|
||
class FakeMoE(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.gate = torch.nn.Linear(hidden, n_experts, bias=False)
|
||
self.experts = torch.nn.ModuleList([FakeExpert() for _ in range(n_experts)])
|
||
|
||
moe = FakeMoE()
|
||
torch.manual_seed(42)
|
||
for p in moe.parameters():
|
||
p.data = torch.randn_like(p.data)
|
||
|
||
shared_dir = self._make_direction(hidden)
|
||
|
||
# Only expert 0 has a per-expert direction
|
||
expert_dirs = {0: torch.randn(hidden).div_(torch.randn(hidden).norm())}
|
||
expert_dirs[0] = expert_dirs[0] / expert_dirs[0].norm()
|
||
|
||
pipeline = AbliterationPipeline(model_name="test", method="surgical")
|
||
pipeline._on_log = lambda m: None
|
||
pipeline._expert_directions = {0: expert_dirs}
|
||
|
||
orig_e1 = moe.experts[1].down_proj.weight.data.clone()
|
||
|
||
pipeline._project_moe_experts_granular(
|
||
moe, shared_dir, layer_idx=0,
|
||
)
|
||
|
||
# Experts 1,2,3 should be modified (using shared direction)
|
||
assert not torch.allclose(moe.experts[1].down_proj.weight.data, orig_e1), \
|
||
"Expert 1 should use shared direction fallback"
|
||
|
||
def test_router_uses_shared_direction(self):
|
||
"""Router should always use the shared direction, not per-expert."""
|
||
hidden = 16
|
||
n_experts = 4
|
||
|
||
class FakeExpert(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.down_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
|
||
class FakeMoE(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.gate = torch.nn.Linear(hidden, n_experts, bias=False)
|
||
self.experts = torch.nn.ModuleList([FakeExpert() for _ in range(n_experts)])
|
||
|
||
moe = FakeMoE()
|
||
shared_dir = self._make_direction(hidden)
|
||
|
||
pipeline = AbliterationPipeline(model_name="test", method="surgical")
|
||
pipeline._on_log = lambda m: None
|
||
pipeline._expert_directions = {0: {0: torch.randn(hidden)}}
|
||
|
||
orig_gate = moe.gate.weight.data.clone()
|
||
|
||
pipeline._project_moe_experts_granular(moe, shared_dir, layer_idx=0)
|
||
|
||
# Gate should be projected
|
||
assert not torch.allclose(moe.gate.weight.data, orig_gate), \
|
||
"Router should be projected with shared direction"
|
||
|
||
# Gate's projection onto shared direction should be near zero
|
||
proj = (moe.gate.weight.data @ shared_dir).norm().item()
|
||
assert proj < 1e-4, f"Router should have shared dir removed, proj={proj}"
|
||
|
||
def test_shared_expert_uses_shared_direction(self):
|
||
"""Shared expert should always use the shared direction."""
|
||
hidden = 16
|
||
|
||
class FakeExpert(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.down_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
self.up_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
|
||
class FakeMoE(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.gate = torch.nn.Linear(hidden, 2, bias=False)
|
||
self.shared_expert = torch.nn.Module()
|
||
self.shared_expert.down_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
self.shared_expert.up_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
self.experts = torch.nn.ModuleList([FakeExpert() for _ in range(2)])
|
||
|
||
moe = FakeMoE()
|
||
shared_dir = self._make_direction(hidden)
|
||
|
||
pipeline = AbliterationPipeline(model_name="test", method="surgical")
|
||
pipeline._on_log = lambda m: None
|
||
pipeline._expert_directions = {0: {0: torch.randn(hidden)}}
|
||
|
||
orig_shared = moe.shared_expert.down_proj.weight.data.clone()
|
||
|
||
pipeline._project_moe_experts_granular(moe, shared_dir, layer_idx=0)
|
||
|
||
assert not torch.allclose(moe.shared_expert.down_proj.weight.data, orig_shared), \
|
||
"Shared expert should be projected"
|
||
|
||
|
||
class TestProjectFused3DGranular:
|
||
"""Test _project_fused_3d_granular for fused 3D expert tensors."""
|
||
|
||
def test_per_expert_directions_on_fused(self):
|
||
"""Each expert slice should use its own direction."""
|
||
hidden = 16
|
||
intermediate = 32
|
||
n_experts = 4
|
||
|
||
class FusedExperts(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.down_proj = torch.nn.Parameter(torch.randn(n_experts, intermediate, hidden))
|
||
|
||
container = FusedExperts()
|
||
torch.manual_seed(42)
|
||
|
||
shared_dir = torch.randn(hidden, 1)
|
||
shared_dir = shared_dir / shared_dir.norm()
|
||
|
||
# Per-expert directions
|
||
expert_dirs = {}
|
||
for ei in range(n_experts):
|
||
d = torch.randn(hidden)
|
||
d = d / d.norm()
|
||
expert_dirs[ei] = d
|
||
|
||
orig_data = container.down_proj.data.clone()
|
||
|
||
count = AbliterationPipeline._project_fused_3d_granular(
|
||
container, shared_dir, expert_dirs, ["down_proj"],
|
||
norm_preserve=False, scale=1.0,
|
||
)
|
||
|
||
assert count == n_experts, f"Should project {n_experts} experts, got {count}"
|
||
|
||
# Each expert should be modified
|
||
for ei in range(n_experts):
|
||
assert not torch.allclose(
|
||
container.down_proj.data[ei], orig_data[ei]
|
||
), f"Expert {ei} should be modified"
|
||
|
||
def test_fallback_to_shared_on_fused(self):
|
||
"""Experts without per-expert direction should use shared direction."""
|
||
hidden = 16
|
||
intermediate = 32
|
||
n_experts = 4
|
||
|
||
class FusedExperts(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.down_proj = torch.nn.Parameter(torch.randn(n_experts, intermediate, hidden))
|
||
|
||
container = FusedExperts()
|
||
torch.manual_seed(42)
|
||
|
||
shared_dir = torch.randn(hidden, 1)
|
||
shared_dir = shared_dir / shared_dir.norm()
|
||
|
||
# Only expert 0 has a direction
|
||
expert_dirs = {0: torch.randn(hidden).div_(1.0)}
|
||
expert_dirs[0] = expert_dirs[0] / expert_dirs[0].norm()
|
||
|
||
orig_data = container.down_proj.data.clone()
|
||
|
||
count = AbliterationPipeline._project_fused_3d_granular(
|
||
container, shared_dir, expert_dirs, ["down_proj"],
|
||
norm_preserve=False, scale=1.0,
|
||
)
|
||
|
||
assert count == n_experts
|
||
# All experts should be modified (experts 1-3 use shared dir)
|
||
for ei in range(n_experts):
|
||
assert not torch.allclose(
|
||
container.down_proj.data[ei], orig_data[ei]
|
||
), f"Expert {ei} should be modified"
|
||
|
||
def test_norm_preserve_on_fused(self):
|
||
"""Fused 3D with norm_preserve should maintain per-expert norms."""
|
||
hidden = 16
|
||
intermediate = 32
|
||
n_experts = 4
|
||
|
||
class FusedExperts(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.down_proj = torch.nn.Parameter(torch.randn(n_experts, intermediate, hidden))
|
||
|
||
container = FusedExperts()
|
||
torch.manual_seed(42)
|
||
|
||
shared_dir = torch.randn(hidden, 1)
|
||
shared_dir = shared_dir / shared_dir.norm()
|
||
|
||
expert_dirs = {}
|
||
for ei in range(n_experts):
|
||
d = torch.randn(hidden)
|
||
expert_dirs[ei] = d / d.norm()
|
||
|
||
orig_norms = [container.down_proj.data[i].norm().item() for i in range(n_experts)]
|
||
|
||
AbliterationPipeline._project_fused_3d_granular(
|
||
container, shared_dir, expert_dirs, ["down_proj"],
|
||
norm_preserve=True, scale=1.0,
|
||
)
|
||
|
||
for i in range(n_experts):
|
||
new_norm = container.down_proj.data[i].norm().item()
|
||
assert abs(orig_norms[i] - new_norm) < 1e-3, (
|
||
f"Expert {i} norm not preserved: {orig_norms[i]:.4f} vs {new_norm:.4f}"
|
||
)
|
||
|
||
def test_skips_non_3d_params(self):
|
||
"""Should skip parameters that are not 3-dimensional."""
|
||
hidden = 16
|
||
|
||
class FlatExperts(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.down_proj = torch.nn.Parameter(torch.randn(32, hidden))
|
||
|
||
container = FlatExperts()
|
||
shared_dir = torch.randn(hidden, 1)
|
||
shared_dir = shared_dir / shared_dir.norm()
|
||
|
||
count = AbliterationPipeline._project_fused_3d_granular(
|
||
container, shared_dir, {}, ["down_proj"],
|
||
norm_preserve=False, scale=1.0,
|
||
)
|
||
assert count == 0
|
||
|
||
|
||
class TestEGAExciseIntegration:
|
||
"""Test that EGA integrates properly in the excise stage path."""
|
||
|
||
def test_ega_pipeline_flags(self):
|
||
"""Pipeline with surgical method should enable per_expert_directions."""
|
||
pipeline = AbliterationPipeline(model_name="test", method="surgical")
|
||
assert pipeline.per_expert_directions is True
|
||
|
||
def test_ega_only_on_primary_direction(self):
|
||
"""EGA should only apply for dir_idx==0, not higher SVD directions."""
|
||
# This is enforced by the `and dir_idx == 0` check in _excise
|
||
# We verify the code structure exists
|
||
from obliteratus.abliterate import AbliterationPipeline
|
||
import inspect
|
||
source = inspect.getsource(AbliterationPipeline._excise_inner)
|
||
assert "dir_idx == 0" in source, "EGA should only apply for primary direction"
|
||
assert "_project_moe_experts_granular" in source, "EGA method should be called in excise"
|
||
|
||
def test_ega_distill_integration(self):
|
||
"""EGA should be called during distill when per_expert_directions is enabled."""
|
||
from obliteratus.abliterate import AbliterationPipeline
|
||
import inspect
|
||
source = inspect.getsource(AbliterationPipeline._distill)
|
||
assert "_compute_expert_granular_directions" in source
|
||
assert "per_expert_directions" in source
|
||
|
||
def test_nuclear_method_enables_ega(self):
|
||
"""Nuclear method should also enable per_expert_directions."""
|
||
cfg = METHODS["nuclear"]
|
||
assert cfg["per_expert_directions"] is True
|
||
pipeline = AbliterationPipeline(model_name="test", method="nuclear")
|
||
assert pipeline.per_expert_directions is True
|
||
|
||
def test_basic_method_disables_ega(self):
|
||
"""Basic method should not enable per_expert_directions."""
|
||
cfg = METHODS["basic"]
|
||
assert cfg.get("per_expert_directions", False) is False
|
||
|
||
def test_inverted_method_enables_ega(self):
|
||
"""Inverted method should enable per_expert_directions."""
|
||
cfg = METHODS["inverted"]
|
||
assert cfg["per_expert_directions"] is True
|
||
|
||
def test_ega_with_routing_data_end_to_end(self):
|
||
"""End-to-end: EGA computes directions and granular projection modifies weights."""
|
||
hidden = 16
|
||
n_experts = 4
|
||
|
||
class FakeExpert(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.down_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
self.up_proj = torch.nn.Linear(hidden, 32, bias=False)
|
||
|
||
class FakeMoE(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.gate = torch.nn.Linear(hidden, n_experts, bias=False)
|
||
self.experts = torch.nn.ModuleList([FakeExpert() for _ in range(n_experts)])
|
||
|
||
moe = FakeMoE()
|
||
torch.manual_seed(42)
|
||
for p in moe.parameters():
|
||
p.data = torch.randn_like(p.data)
|
||
|
||
pipeline = AbliterationPipeline(model_name="test", method="surgical")
|
||
pipeline._on_log = lambda m: None
|
||
pipeline._on_stage = lambda r: None
|
||
pipeline._strong_layers = [0]
|
||
|
||
# Simulate EGA routing data
|
||
h_logits = [torch.randn(n_experts) for _ in range(5)]
|
||
s_logits = [torch.randn(n_experts) for _ in range(5)]
|
||
pipeline._routing_harmful = {0: h_logits}
|
||
pipeline._routing_harmless = {0: s_logits}
|
||
|
||
# Simulate activations with clear separation
|
||
refusal_dir = torch.randn(hidden)
|
||
refusal_dir = refusal_dir / refusal_dir.norm()
|
||
pipeline._harmful_acts = {0: [torch.randn(hidden) + 2 * refusal_dir for _ in range(5)]}
|
||
pipeline._harmless_acts = {0: [torch.randn(hidden) - 2 * refusal_dir for _ in range(5)]}
|
||
|
||
# Step 1: compute EGA directions
|
||
pipeline._compute_expert_granular_directions()
|
||
assert 0 in pipeline._expert_directions
|
||
assert len(pipeline._expert_directions[0]) > 0
|
||
|
||
# Step 2: apply granular projection
|
||
shared_dir = torch.randn(hidden, 1)
|
||
shared_dir = shared_dir / shared_dir.norm()
|
||
|
||
orig_expert0 = moe.experts[0].down_proj.weight.data.clone()
|
||
|
||
count = pipeline._project_moe_experts_granular(
|
||
moe, shared_dir, layer_idx=0,
|
||
)
|
||
|
||
assert count > 0
|
||
assert not torch.allclose(moe.experts[0].down_proj.weight.data, orig_expert0), \
|
||
"Expert weights should be modified by EGA"
|