mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-30 22:50:46 +02:00
232 lines
8.1 KiB
Python
232 lines
8.1 KiB
Python
"""Contracts for the explicitly supported Qwen3.8-27B residual writers."""
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
from torch import nn
|
|
|
|
from obliteratus.abliterate import AbliterationPipeline
|
|
from obliteratus.models.qwen35_contracts import (
|
|
Qwen35LayerTargets,
|
|
Qwen35ContractError,
|
|
validate_qwen38_27b_projection_contract,
|
|
)
|
|
|
|
|
|
class _LinearAttentionLayer(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
with torch.device("meta"):
|
|
self.linear_attn = nn.Module()
|
|
self.linear_attn.out_proj = nn.Linear(6144, 5120, bias=False)
|
|
self.linear_attn.in_proj_qkv = nn.Linear(5120, 10240, bias=False)
|
|
self.linear_attn.in_proj_z = nn.Linear(5120, 6144, bias=False)
|
|
self.mlp = nn.Module()
|
|
self.mlp.down_proj = nn.Linear(17408, 5120, bias=False)
|
|
self.mlp.gate_proj = nn.Linear(5120, 17408, bias=False)
|
|
self.mlp.up_proj = nn.Linear(5120, 17408, bias=False)
|
|
|
|
|
|
class _FullAttentionLayer(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
with torch.device("meta"):
|
|
self.self_attn = nn.Module()
|
|
self.self_attn.o_proj = nn.Linear(6144, 5120, bias=False)
|
|
self.self_attn.q_proj = nn.Linear(5120, 12288, bias=False)
|
|
self.self_attn.k_proj = nn.Linear(5120, 1024, bias=False)
|
|
self.self_attn.v_proj = nn.Linear(5120, 1024, bias=False)
|
|
self.mlp = nn.Module()
|
|
self.mlp.down_proj = nn.Linear(17408, 5120, bias=False)
|
|
self.mlp.gate_proj = nn.Linear(5120, 17408, bias=False)
|
|
self.mlp.up_proj = nn.Linear(5120, 17408, bias=False)
|
|
|
|
|
|
def _handle():
|
|
layer_types = [
|
|
layer_type
|
|
for _ in range(16)
|
|
for layer_type in (
|
|
"linear_attention",
|
|
"linear_attention",
|
|
"linear_attention",
|
|
"full_attention",
|
|
)
|
|
]
|
|
layers = nn.ModuleList(
|
|
_LinearAttentionLayer() if kind == "linear_attention" else _FullAttentionLayer()
|
|
for kind in layer_types
|
|
)
|
|
text_config = SimpleNamespace(
|
|
layer_types=layer_types,
|
|
hidden_size=5120,
|
|
intermediate_size=17408,
|
|
num_hidden_layers=64,
|
|
)
|
|
return SimpleNamespace(
|
|
architecture="qwen3_5",
|
|
model_name="Qwen/Qwen3.8-27B",
|
|
config=SimpleNamespace(text_config=text_config),
|
|
model=SimpleNamespace(model=SimpleNamespace(layers=layers)),
|
|
num_layers=64,
|
|
)
|
|
|
|
|
|
def test_qwen38_manifest_contains_only_residual_writers():
|
|
targets = validate_qwen38_27b_projection_contract(_handle())
|
|
|
|
assert len(targets) == 64
|
|
assert sum(target.mixer_attribute == "linear_attn" for target in targets) == 48
|
|
assert sum(target.mixer_attribute == "self_attn" for target in targets) == 16
|
|
assert targets[0].parameter_names == frozenset(
|
|
{"linear_attn.out_proj.weight", "mlp.down_proj.weight"}
|
|
)
|
|
assert targets[3].parameter_names == frozenset(
|
|
{"self_attn.o_proj.weight", "mlp.down_proj.weight"}
|
|
)
|
|
forbidden = {
|
|
"in_proj_qkv",
|
|
"in_proj_z",
|
|
"q_proj",
|
|
"k_proj",
|
|
"v_proj",
|
|
"gate_proj",
|
|
"up_proj",
|
|
"lm_head",
|
|
}
|
|
assert all(
|
|
not any(name in parameter for name in forbidden)
|
|
for target in targets
|
|
for parameter in target.parameter_names
|
|
)
|
|
|
|
|
|
def test_qwen38_manifest_fails_closed_on_topology_change():
|
|
handle = _handle()
|
|
handle.config.text_config.layer_types[0] = "full_attention"
|
|
|
|
with pytest.raises(Qwen35ContractError, match="48 DeltaNet"):
|
|
validate_qwen38_27b_projection_contract(handle)
|
|
|
|
|
|
def test_qwen38_manifest_fails_closed_on_writer_shape_change():
|
|
handle = _handle()
|
|
with torch.device("meta"):
|
|
handle.model.model.layers[0].linear_attn.out_proj = nn.Linear(5120, 5120, bias=False)
|
|
|
|
with pytest.raises(Qwen35ContractError, match="unexpected type or shape"):
|
|
validate_qwen38_27b_projection_contract(handle)
|
|
|
|
|
|
def test_qwen_output_only_projection_preserves_forbidden_tensors_bit_exactly():
|
|
torch.manual_seed(7)
|
|
layer = nn.Module()
|
|
layer.linear_attn = nn.Module()
|
|
layer.linear_attn.out_proj = nn.Linear(4, 4, bias=False)
|
|
layer.linear_attn.in_proj_qkv = nn.Linear(4, 8, bias=False)
|
|
layer.linear_attn.in_proj_z = nn.Linear(4, 4, bias=False)
|
|
layer.mlp = nn.Module()
|
|
layer.mlp.down_proj = nn.Linear(8, 4, bias=False)
|
|
layer.mlp.gate_proj = nn.Linear(4, 8, bias=False)
|
|
layer.mlp.up_proj = nn.Linear(4, 8, bias=False)
|
|
target = Qwen35LayerTargets("linear_attn", "out_proj")
|
|
forbidden_before = {
|
|
name: parameter.detach().clone()
|
|
for name, parameter in layer.named_parameters()
|
|
if name not in target.parameter_names
|
|
}
|
|
allowed_before = {
|
|
name: parameter.detach().clone()
|
|
for name, parameter in layer.named_parameters()
|
|
if name in target.parameter_names
|
|
}
|
|
direction = torch.tensor([[1.0], [0.0], [0.0], [0.0]])
|
|
|
|
saved_norms = AbliterationPipeline._capture_layer_weight_norms(
|
|
layer, parameter_names=target.parameter_names
|
|
)
|
|
AbliterationPipeline._project_out_advanced(
|
|
layer.linear_attn,
|
|
direction,
|
|
[target.mixer_output],
|
|
norm_preserve=False,
|
|
)
|
|
AbliterationPipeline._project_out_advanced(
|
|
layer.mlp,
|
|
direction,
|
|
[target.ffn_output],
|
|
norm_preserve=False,
|
|
)
|
|
AbliterationPipeline._restore_layer_weight_norms(layer, saved_norms)
|
|
|
|
current = dict(layer.named_parameters())
|
|
assert all(torch.equal(current[name], value) for name, value in forbidden_before.items())
|
|
assert all(not torch.equal(current[name], value) for name, value in allowed_before.items())
|
|
|
|
|
|
def test_qwen_excise_route_does_not_touch_inputs_gates_or_lm_head():
|
|
torch.manual_seed(11)
|
|
layer = nn.Module()
|
|
layer.linear_attn = nn.Module()
|
|
layer.linear_attn.out_proj = nn.Linear(4, 4, bias=False)
|
|
layer.linear_attn.in_proj_qkv = nn.Linear(4, 8, bias=False)
|
|
layer.linear_attn.in_proj_z = nn.Linear(4, 4, bias=False)
|
|
layer.mlp = nn.Module()
|
|
layer.mlp.down_proj = nn.Linear(8, 4, bias=False)
|
|
layer.mlp.gate_proj = nn.Linear(4, 8, bias=False)
|
|
layer.mlp.up_proj = nn.Linear(4, 8, bias=False)
|
|
model = nn.Module()
|
|
model.model = nn.Module()
|
|
model.model.layers = nn.ModuleList([layer])
|
|
model.lm_head = nn.Linear(4, 16, bias=False)
|
|
handle = SimpleNamespace(
|
|
architecture="qwen3_5",
|
|
config=SimpleNamespace(num_attention_heads=1),
|
|
model=model,
|
|
num_layers=1,
|
|
)
|
|
pipeline = AbliterationPipeline(
|
|
model_name="Qwen/Qwen3.8-27B",
|
|
method="advanced",
|
|
projection_target="all",
|
|
)
|
|
pipeline.handle = handle
|
|
pipeline._qwen35_projection_manifest = (
|
|
Qwen35LayerTargets("linear_attn", "out_proj"),
|
|
)
|
|
pipeline._strong_layers = [0]
|
|
pipeline.refusal_subspaces = {
|
|
0: torch.tensor([[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0]])
|
|
}
|
|
pipeline.refusal_directions = {0: pipeline.refusal_subspaces[0][0]}
|
|
pipeline._layer_excise_weights = {}
|
|
pipeline._on_log = lambda _message: None
|
|
before = {name: value.detach().clone() for name, value in model.named_parameters()}
|
|
|
|
pipeline._excise()
|
|
|
|
after = dict(model.named_parameters())
|
|
allowed = {
|
|
"model.layers.0.linear_attn.out_proj.weight",
|
|
"model.layers.0.mlp.down_proj.weight",
|
|
}
|
|
assert all(not torch.equal(after[name], before[name]) for name in allowed)
|
|
assert all(
|
|
torch.equal(after[name], before[name])
|
|
for name in before.keys() - allowed
|
|
)
|
|
assert pipeline._effective_refinement_passes == 1
|
|
method_config = pipeline._build_metadata()["method_config"]
|
|
assert method_config["refinement_passes"] == 1
|
|
assert method_config["requested_refinement_passes"] == 2
|
|
|
|
|
|
@pytest.mark.parametrize("architecture", ["qwen3_5_text", "qwen3_5_moe"])
|
|
def test_other_qwen35_variants_remain_blocked(architecture):
|
|
handle = _handle()
|
|
handle.architecture = architecture
|
|
|
|
with pytest.raises(Qwen35ContractError, match="unsupported Qwen hybrid architecture"):
|
|
validate_qwen38_27b_projection_contract(handle)
|