mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-18 00:47:23 +02:00
290 lines
11 KiB
Python
290 lines
11 KiB
Python
"""LoRA-based reversible ablation mode.
|
|
|
|
Instead of permanent in-place weight surgery, applies ablation via rank-1
|
|
LoRA adapters. This provides:
|
|
|
|
1. **Reversibility**: LoRA adapters can be removed to restore original model
|
|
2. **Composability**: Multiple ablation adapters can be stacked/blended
|
|
3. **PEFT compatibility**: Output adapters work with standard HuggingFace PEFT
|
|
|
|
Inspired by Heretic (p-e-w, 2025) which pioneered LoRA-mediated ablation.
|
|
OBLITERATUS extends this with:
|
|
- Multi-direction rank-k adapters (not just rank-1)
|
|
- MoE-aware LoRA targeting (router + expert-specific adapters)
|
|
- Integration with EGA per-expert directions
|
|
- CoT-aware adapter strength modulation
|
|
|
|
The mathematical equivalence to in-place projection depends on weight orientation:
|
|
|
|
For W of shape (out, hidden) where d is in the hidden dimension:
|
|
In-place: W' = W - scale * W @ d @ d^T
|
|
LoRA: W' = W + B @ A where B = -scale * (W @ d), A = d^T
|
|
|
|
For W of shape (hidden, out) (e.g., Conv1D layers):
|
|
In-place: W' = W - scale * d @ d^T @ W
|
|
LoRA: W' = W + B @ A where B = -scale * d, A = d^T @ W
|
|
|
|
Both produce identical output, but LoRA stores {B, A} separately.
|
|
|
|
References:
|
|
- Hu et al. (2022): LoRA: Low-Rank Adaptation of Large Language Models
|
|
- Heretic (p-e-w, 2025): LoRA-mediated directional ablation
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
if TYPE_CHECKING:
|
|
from obliteratus.abliterate import AbliterationPipeline
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Target module name patterns for LoRA adapter placement
|
|
_LORA_TARGETS = [
|
|
"o_proj", "q_proj", "k_proj", "v_proj",
|
|
"down_proj", "up_proj", "gate_proj",
|
|
"gate", "router",
|
|
]
|
|
|
|
|
|
def compute_lora_adapters(
|
|
pipeline: AbliterationPipeline,
|
|
rank: int = 1,
|
|
) -> dict[str, tuple[torch.Tensor, torch.Tensor]]:
|
|
"""Compute LoRA adapter pairs (B, A) for refusal direction ablation.
|
|
|
|
For each target weight matrix W with refusal direction d:
|
|
A = d^T @ W (rank-1: shape (1, in_features) or (rank, in_features))
|
|
B = -scale * d (rank-1: shape (out_features, 1) or (out_features, rank))
|
|
|
|
So that W + B @ A ≈ W - scale * (d @ d^T) @ W
|
|
|
|
Args:
|
|
pipeline: Initialized pipeline (post-DISTILL, pre-EXCISE).
|
|
rank: LoRA rank (1 = rank-1 ablation, >1 = multi-direction).
|
|
|
|
Returns:
|
|
Dict mapping "layer.{idx}.{module}.{weight}" → (lora_B, lora_A) pairs.
|
|
"""
|
|
from obliteratus.strategies.utils import (
|
|
get_attention_module,
|
|
get_ffn_module,
|
|
get_layer_modules,
|
|
)
|
|
from obliteratus.abliterate import (
|
|
_ATTN_OUT_NAMES,
|
|
_ATTN_IN_NAMES,
|
|
_FFN_OUT_NAMES,
|
|
_FFN_IN_NAMES,
|
|
_ROUTER_NAMES,
|
|
)
|
|
|
|
if not pipeline.handle or not pipeline._strong_layers:
|
|
return {}
|
|
|
|
layers = get_layer_modules(pipeline.handle)
|
|
arch = pipeline.handle.architecture
|
|
adapters: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
|
|
|
|
for idx in pipeline._strong_layers:
|
|
if idx not in pipeline.refusal_subspaces:
|
|
continue
|
|
|
|
subspace = pipeline.refusal_subspaces[idx]
|
|
n_dirs = min(rank, subspace.shape[0])
|
|
|
|
# Compute per-layer regularization (mirroring excise logic)
|
|
reg = pipeline.regularization
|
|
if pipeline.layer_adaptive_strength and idx in pipeline._layer_excise_weights:
|
|
weight = pipeline._layer_excise_weights[idx]
|
|
reg = pipeline.regularization + (1.0 - weight) * (1.0 - pipeline.regularization) * 0.15
|
|
if pipeline.float_layer_interpolation and idx in pipeline._float_layer_weights:
|
|
float_w = pipeline._float_layer_weights[idx]
|
|
reg = reg + (1.0 - float_w) * (1.0 - reg) * 0.3
|
|
|
|
scale = 1.0 - reg
|
|
|
|
# Direction matrix: (n_dirs, hidden_dim)
|
|
D = subspace[:n_dirs].float()
|
|
|
|
# Collect target modules and their weight matrices
|
|
targets: list[tuple[str, nn.Module, list[str]]] = []
|
|
try:
|
|
attn = get_attention_module(layers[idx], arch)
|
|
targets.append(("attn", attn, _ATTN_OUT_NAMES + _ATTN_IN_NAMES))
|
|
except (AttributeError, RuntimeError):
|
|
pass
|
|
try:
|
|
ffn = get_ffn_module(layers[idx], arch)
|
|
targets.append(("ffn", ffn, _FFN_OUT_NAMES + _FFN_IN_NAMES + _ROUTER_NAMES))
|
|
except (AttributeError, RuntimeError):
|
|
pass
|
|
|
|
for module_label, module, candidate_names in targets:
|
|
for name in candidate_names:
|
|
proj = getattr(module, name, None)
|
|
if proj is None or not hasattr(proj, "weight"):
|
|
continue
|
|
|
|
W = proj.weight.data.float()
|
|
|
|
if W.shape[-1] == D.shape[1]:
|
|
# Standard: W is (out, hidden_dim), direction along last axis
|
|
# A = D @ W^T → (n_dirs, out) ... no, we need the correct decomposition
|
|
# W' = W - scale * D^T @ D @ W
|
|
# LoRA: B @ A where B = -scale * D^T shape (hidden_dim, n_dirs)
|
|
# A = D @ W shape (n_dirs, out)...
|
|
# Actually for nn.Linear: output = input @ W^T + b
|
|
# So to affect output: we need delta_W such that input @ delta_W^T
|
|
# delta_W = -scale * d @ d^T @ W → B = -scale * d, A = d^T @ W
|
|
# B shape: (out_features, n_dirs) where out_features = W.shape[0]
|
|
# Wait, let me reconsider. For W shape (out, in):
|
|
# delta_W = -scale * (d_col @ d_col^T) @ W
|
|
# where d_col is (in, 1) if direction matches input dim
|
|
# delta_W = -scale * d_col @ (d_col^T @ W)
|
|
# = -scale * d_col @ coeff where coeff = d_col^T @ W = (1, out)
|
|
# So: B = -scale * d_col = (in, 1), A = d_col^T @ W = (1, out)
|
|
# But LoRA convention: delta_W = B @ A where B is (out, r), A is (r, in)
|
|
# So we need: delta_W^T = A^T @ B^T
|
|
# Hmm, let me just store the computed delta and split it
|
|
|
|
# For each direction, compute the rank-1 adapter
|
|
lora_B_parts = []
|
|
lora_A_parts = []
|
|
for di in range(n_dirs):
|
|
d = D[di] # (hidden_dim,)
|
|
d_col = d.unsqueeze(-1) # (hidden_dim, 1)
|
|
# coeff = d^T @ W^T = (d @ W^T) → but W is (out, hidden), so:
|
|
# For W @ d → (out, 1) projection
|
|
coeff = (W @ d_col).squeeze(-1) # (out,)
|
|
# delta_W = -scale * d_col @ coeff.unsqueeze(0) would be (hidden, out)
|
|
# But we want (out, hidden) to match W shape
|
|
# delta_W[i,j] = -scale * coeff[i] * d[j]
|
|
# = B[i,:] @ A[:,j] where B = -scale * coeff.unsqueeze(1) = (out,1)
|
|
# A = d.unsqueeze(0) = (1, hidden)
|
|
lora_B_parts.append(-scale * coeff.unsqueeze(1)) # (out, 1)
|
|
lora_A_parts.append(d.unsqueeze(0)) # (1, hidden)
|
|
|
|
lora_B = torch.cat(lora_B_parts, dim=1) # (out, n_dirs)
|
|
lora_A = torch.cat(lora_A_parts, dim=0) # (n_dirs, hidden)
|
|
|
|
elif W.shape[0] == D.shape[1]:
|
|
# Transposed case: W is (hidden_dim, out)
|
|
lora_B_parts = []
|
|
lora_A_parts = []
|
|
for di in range(n_dirs):
|
|
d = D[di] # (hidden_dim,)
|
|
coeff = (d @ W) # (out,)
|
|
lora_B_parts.append(-scale * d.unsqueeze(1)) # (hidden, 1)
|
|
lora_A_parts.append(coeff.unsqueeze(0)) # (1, out)
|
|
|
|
lora_B = torch.cat(lora_B_parts, dim=1) # (hidden, n_dirs)
|
|
lora_A = torch.cat(lora_A_parts, dim=0) # (n_dirs, out)
|
|
else:
|
|
continue
|
|
|
|
key = f"layer.{idx}.{module_label}.{name}"
|
|
adapters[key] = (lora_B.half(), lora_A.half())
|
|
|
|
pipeline.log(f"Computed {len(adapters)} LoRA adapter pairs (rank={rank})")
|
|
return adapters
|
|
|
|
|
|
def apply_lora_adapters(
|
|
pipeline: AbliterationPipeline,
|
|
adapters: dict[str, tuple[torch.Tensor, torch.Tensor]],
|
|
):
|
|
"""Apply pre-computed LoRA adapters by modifying weights in-place.
|
|
|
|
This is equivalent to merging the LoRA into the base model.
|
|
The adapters dict is stored in pipeline._lora_adapters for potential
|
|
later unmerging.
|
|
"""
|
|
from obliteratus.strategies.utils import (
|
|
get_attention_module,
|
|
get_ffn_module,
|
|
get_layer_modules,
|
|
)
|
|
|
|
if not pipeline.handle:
|
|
return
|
|
|
|
layers = get_layer_modules(pipeline.handle)
|
|
arch = pipeline.handle.architecture
|
|
applied = 0
|
|
|
|
for key, (lora_B, lora_A) in adapters.items():
|
|
parts = key.split(".")
|
|
if len(parts) != 4:
|
|
continue
|
|
_, idx_str, module_label, weight_name = parts
|
|
idx = int(idx_str)
|
|
|
|
try:
|
|
if module_label == "attn":
|
|
module = get_attention_module(layers[idx], arch)
|
|
else:
|
|
module = get_ffn_module(layers[idx], arch)
|
|
except (AttributeError, RuntimeError):
|
|
continue
|
|
|
|
proj = getattr(module, weight_name, None)
|
|
if proj is None or not hasattr(proj, "weight"):
|
|
continue
|
|
|
|
W = proj.weight.data
|
|
delta = (lora_B @ lora_A).to(device=W.device, dtype=W.dtype)
|
|
|
|
if delta.shape == W.shape:
|
|
W.add_(delta)
|
|
applied += 1
|
|
|
|
pipeline._lora_adapters = adapters
|
|
pipeline.log(f"Applied {applied} LoRA adapters (merged into weights)")
|
|
|
|
|
|
def save_lora_adapters(
|
|
adapters: dict[str, tuple[torch.Tensor, torch.Tensor]],
|
|
output_dir: str | Path,
|
|
):
|
|
"""Save LoRA adapters to disk for later use.
|
|
|
|
Saves as a simple dict of {key: (B, A)} tensors using torch.save.
|
|
Can be loaded and applied to the original model for reversible ablation.
|
|
"""
|
|
output_path = Path(output_dir)
|
|
output_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
save_dict = {}
|
|
for key, (B, A) in adapters.items():
|
|
save_dict[f"{key}.lora_B"] = B
|
|
save_dict[f"{key}.lora_A"] = A
|
|
|
|
adapter_path = output_path / "abliteration_lora_adapters.pt"
|
|
torch.save(save_dict, adapter_path)
|
|
|
|
# Also save adapter config for PEFT compatibility
|
|
import json
|
|
config = {
|
|
"adapter_type": "obliteratus_abliteration_lora",
|
|
"n_adapters": len(adapters),
|
|
"target_modules": list(set(
|
|
k.split(".")[-1] for k in adapters
|
|
)),
|
|
"description": (
|
|
"Reversible abliteration LoRA adapters generated by OBLITERATUS. "
|
|
"These adapters remove refusal directions from the model when merged."
|
|
),
|
|
}
|
|
(output_path / "abliteration_lora_config.json").write_text(
|
|
json.dumps(config, indent=2)
|
|
)
|
|
|
|
return adapter_path
|