151 lines
6.0 KiB
Python
151 lines
6.0 KiB
Python
#Liquid Foundation Model (LFM) Abliteration Script
|
|
#This script removes the safety/refusal guardrails from Liquid AI's continuous-time hybrid models
|
|
#using layerwise orthogonal projection.
|
|
|
|
#Architecture Support: LFM-1.2B, LFM-3B
|
|
#Author: Paperscarecrow & Gemini 3.1 pro
|
|
|
|
import torch
|
|
from datasets import load_dataset
|
|
import random
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
from tqdm import tqdm
|
|
|
|
# ==========================================
|
|
# 1. ROCM / RDNA3 COMPATIBILITY PATCH
|
|
# ==========================================
|
|
# Bypasses a known `hipblas` segmentation fault on consumer AMD GPUs when processing Liquid's RoPE tensors.
|
|
import transformers.models.lfm2.modeling_lfm2 as lfm2_modeling
|
|
|
|
def patched_rope_forward(self, x, position_ids):
|
|
inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
|
|
position_ids_expanded = position_ids[:, None, :].float()
|
|
|
|
device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
|
|
with torch.autocast(device_type=device_type, enabled=False):
|
|
# Uses element-wise multiplication (*) instead of batched matmul (@) for memory safety
|
|
freqs = (inv_freq_expanded.float() * position_ids_expanded.float()).transpose(1, 2)
|
|
emb = torch.cat((freqs, freqs), dim=-1)
|
|
cos = emb.cos() * self.attention_scaling
|
|
sin = emb.sin() * self.attention_scaling
|
|
|
|
return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
|
|
|
|
lfm2_modeling.Lfm2RotaryEmbedding.forward = patched_rope_forward
|
|
|
|
# ==========================================
|
|
# CONFIGURATION
|
|
# ==========================================
|
|
MODEL_PATH = "liquidai/LFM-1.2B" # Local path or HF Hub ID
|
|
SAVE_PATH = "./LFM-1.2B-Abliterated"
|
|
TARGET_LAYER = 8 # Middle layer typically holds the clearest refusal representation
|
|
NUM_SAMPLES = 100 # Number of prompts to average for the refusal vector
|
|
|
|
# ==========================================
|
|
# 2. MEASUREMENT PHASE
|
|
# ==========================================
|
|
def get_refusal_direction(model, tokenizer, harmful_prompts, harmless_prompts, target_layer):
|
|
print(f"Measuring hidden states at layer {target_layer}...")
|
|
|
|
hidden_states_harmful = []
|
|
hidden_states_harmless = []
|
|
|
|
def hook_fn(module, input, output):
|
|
h = output[0] if isinstance(output, tuple) else output
|
|
return h[:, -1, :].detach().clone()
|
|
|
|
layer = model.model.layers[target_layer]
|
|
handle = layer.register_forward_hook(
|
|
lambda m, i, o: hidden_states_harmful.append(hook_fn(m, i, o))
|
|
if is_harmful else hidden_states_harmless.append(hook_fn(m, i, o))
|
|
)
|
|
|
|
global is_harmful
|
|
is_harmful = True
|
|
print("Processing harmful instructions...")
|
|
for prompt in tqdm(harmful_prompts):
|
|
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
|
|
with torch.no_grad():
|
|
model(**inputs)
|
|
|
|
is_harmful = False
|
|
print("Processing harmless instructions...")
|
|
for prompt in tqdm(harmless_prompts):
|
|
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
|
|
with torch.no_grad():
|
|
model(**inputs)
|
|
|
|
handle.remove()
|
|
|
|
mean_harmful = torch.stack(hidden_states_harmful).mean(dim=0).squeeze()
|
|
mean_harmless = torch.stack(hidden_states_harmless).mean(dim=0).squeeze()
|
|
|
|
return mean_harmful - mean_harmless
|
|
|
|
# ==========================================
|
|
# 3. SURGERY PHASE
|
|
# ==========================================
|
|
def abliterate_liquid_weights(model, refusal_direction):
|
|
v = refusal_direction.to(model.device, dtype=model.dtype)
|
|
v = v / v.norm() # Normalize the vector
|
|
|
|
print("\nCommencing orthogonal projection on Liquid weights...")
|
|
|
|
for i, layer in enumerate(tqdm(model.model.layers, desc="Scrubbing layers")):
|
|
# Scrub Convolution Output Projection (Token Mixing)
|
|
if hasattr(layer, 'conv') and hasattr(layer.conv, 'out_proj'):
|
|
W_conv = layer.conv.out_proj.weight.data
|
|
proj_conv = torch.outer(v, v @ W_conv)
|
|
layer.conv.out_proj.weight.data = W_conv - proj_conv
|
|
|
|
# Scrub Feed-Forward Down Projection (Channel Mixing)
|
|
if hasattr(layer, 'feed_forward') and hasattr(layer.feed_forward, 'w2'):
|
|
W_ffn = layer.feed_forward.w2.weight.data
|
|
proj_ffn = torch.outer(v, v @ W_ffn)
|
|
layer.feed_forward.w2.weight.data = W_ffn - proj_ffn
|
|
|
|
print("Surgery complete.")
|
|
return model
|
|
|
|
# ==========================================
|
|
# MAIN EXECUTION
|
|
# ==========================================
|
|
if __name__ == "__main__":
|
|
print("Loading tokenizer and base model...")
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
|
|
|
|
# Note: device_map="cpu" and float16 recommended for consumer AMD hardware to avoid hipblas segfaults
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
MODEL_PATH,
|
|
device_map="cpu",
|
|
torch_dtype=torch.float16,
|
|
trust_remote_code=True
|
|
)
|
|
|
|
print("Fetching robust datasets for vector calculation...")
|
|
# Load standardized ablation datasets
|
|
dataset_harmful = load_dataset("mlabonne/harmful_behaviors", split="train")
|
|
dataset_harmless = load_dataset("mlabonne/harmless_alpaca", split="train")
|
|
|
|
random.seed(42)
|
|
harmful_sampled = random.sample(dataset_harmful['text'], NUM_SAMPLES)
|
|
harmless_sampled = random.sample(dataset_harmless['text'], NUM_SAMPLES)
|
|
|
|
# Format strictly to Liquid's required template
|
|
my_harmful_prompts = [f"<|user|>\n{prompt}\n<|assistant|>\n" for prompt in harmful_sampled]
|
|
my_harmless_prompts = [f"<|user|>\n{prompt}\n<|assistant|>\n" for prompt in harmless_sampled]
|
|
|
|
refusal_dir = get_refusal_direction(
|
|
model, tokenizer,
|
|
my_harmful_prompts, my_harmless_prompts,
|
|
TARGET_LAYER
|
|
)
|
|
|
|
model = abliterate_liquid_weights(model, refusal_dir)
|
|
|
|
print(f"\nSaving untethered model to {SAVE_PATH}...")
|
|
model.save_pretrained(SAVE_PATH)
|
|
tokenizer.save_pretrained(SAVE_PATH)
|
|
|
|
print("Done! Ready for GGUF conversion or inference.")
|