perf: all

Author: whes1015

common.py

@@ -0,0 +1,62 @@
+"""Shared helpers for refusal-direction extraction and ablation.
+
+Both compute_refusal_dir.py and inference.py import from here so the model
+choice and loading config live in exactly one place.
+"""
+import os
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+# Pick the model once; both scripts use it.
+MODEL_ID = "tiiuae/Falcon3-1B-Instruct"
+# MODEL_ID = "Qwen/Qwen3-1.7B"
+# MODEL_ID = "stabilityai/stablelm-2-zephyr-1_6b"
+# MODEL_ID = "Qwen/Qwen1.5-1.8B-Chat"
+# MODEL_ID = "Qwen/Qwen-1_8B-chat"
+# MODEL_ID = "google/gemma-1.1-2b-it"
+# MODEL_ID = "google/gemma-1.1-7b-it"
+# MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct"
+
+# Let every backend use all CPU cores and the fast matmul kernels.
+torch.set_num_threads(os.cpu_count() or 1)
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
+
+def load_model_and_tokenizer(model_id: str = MODEL_ID):
+    """Load a 4-bit quantized causal LM and its tokenizer, ready for inference."""
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        trust_remote_code=True,
+        dtype=torch.float16,
+        device_map="cuda",
+        quantization_config=BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16,
+        ),
+    )
+    model.eval()
+
+    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+    # Left padding lets us read the last-token hidden state at index -1 for a
+    # whole batch in one shot.
+    tokenizer.padding_side = "left"
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    return model, tokenizer
+
+
+def refusal_dir_path(model_id: str = MODEL_ID) -> str:
+    """Filename used to cache the refusal direction for a given model."""
+    return model_id.replace("/", "_") + "_refusal_dir.pt"
+
+
+def project_out(x: torch.Tensor, direction: torch.Tensor) -> torch.Tensor:
+    """Remove the component of `x` that lies along the unit vector `direction`.
+
+    proj = (x . d) d ;  result = x - proj.  Assumes ``direction`` is normalized.
+    Equivalent to the original einsum, but a plain matmul does the same job.
+    """
+    return x - (x @ direction).unsqueeze(-1) * direction

compute_refusal_dir.py

@@ -1,86 +1,59 @@
 import random
 
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
-
 from tqdm import tqdm
 
-torch.inference_mode()
-
-MODEL_ID = "tiiuae/Falcon3-1B-Instruct"
-# MODEL_ID = "Qwen/Qwen3-1.7B"
-# MODEL_ID = "stabilityai/stablelm-2-zephyr-1_6b"
-# MODEL_ID = "Qwen/Qwen1.5-1.8B-Chat"
-# MODEL_ID = "Qwen/Qwen-1_8B-chat"
-# MODEL_ID = "google/gemma-1.1-2b-it"
-# MODEL_ID = "google/gemma-1.1-7b-it"
-# MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct"
-
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID,
-                                             trust_remote_code=True,
-                                             dtype=torch.float16,
-                                             device_map="cuda",
-                                             quantization_config=BitsAndBytesConfig(load_in_4bit=True,
-                                                                                    bnb_4bit_compute_dtype=torch.float16))
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+from common import load_model_and_tokenizer, refusal_dir_path
 
 # settings:
-instructions = 32
-layer_idx = int(len(model.model.layers) * 0.6)
-pos = -1
-
-print("Instruction count: " + str(instructions))
-print("Layer index: " + str(layer_idx))
-
-with open("harmful.txt", "r") as f:
-    harmful = f.readlines()
+N_INSTRUCTIONS = 32  # harmful + harmless samples each
+BATCH_SIZE = 16      # prompts per forward pass (lower this if you run out of VRAM)
+POS = -1             # which token position to read (-1 = last prompt token)
 
-with open("harmless.txt", "r") as f:
-    harmless = f.readlines()
-
-harmful_instructions = random.sample(harmful, instructions)
-harmless_instructions = random.sample(harmless, instructions)
+model, tokenizer = load_model_and_tokenizer()
+layer_idx = int(len(model.model.layers) * 0.6)
 
-harmful_toks = [
-    tokenizer.apply_chat_template(conversation=[{"role": "user", "content": insn}],
-                                  add_generation_prompt=True,
-                                  return_tensors="pt") for insn in harmful_instructions]
-harmless_toks = [
-    tokenizer.apply_chat_template(conversation=[{"role": "user", "content": insn}],
-                                  add_generation_prompt=True,
-                                  return_tensors="pt") for insn in harmless_instructions]
+print(f"Instruction count: {N_INSTRUCTIONS}")
+print(f"Layer index: {layer_idx}")
 
-max_its = instructions*2
-bar = tqdm(total=max_its)
 
+def load_instructions(path: str) -> list[str]:
+    with open(path, "r") as f:
+        return [line.strip() for line in f if line.strip()]
 
-def generate(toks):
-    bar.update(n=1)
-    return model.generate(toks.to(model.device),
-                          use_cache=False,
-                          max_new_tokens=1,
-                          return_dict_in_generate=True,
-                          output_hidden_states=True)
 
+harmful = random.sample(load_instructions("harmful.txt"), N_INSTRUCTIONS)
+harmless = random.sample(load_instructions("harmless.txt"), N_INSTRUCTIONS)
 
-harmful_outputs = [generate(toks) for toks in harmful_toks]
-harmless_outputs = [generate(toks) for toks in harmless_toks]
 
-bar.close()
+@torch.inference_mode()
+def mean_hidden_state(instructions: list[str]) -> torch.Tensor:
+    """Mean last-token hidden state at `layer_idx`, averaged over `instructions`.
 
-harmful_hidden = [output.hidden_states[0][layer_idx][:, pos, :] for output in harmful_outputs]
-harmless_hidden = [output.hidden_states[0][layer_idx][:, pos, :] for output in harmless_outputs]
+    Runs a single batched forward pass per BATCH_SIZE chunk instead of one
+    generate() call per prompt. Accumulated in float32 for a stable mean.
+    """
+    total = None
+    for start in tqdm(range(0, len(instructions), BATCH_SIZE)):
+        batch = instructions[start:start + BATCH_SIZE]
+        toks = tokenizer.apply_chat_template(
+            [[{"role": "user", "content": text}] for text in batch],
+            add_generation_prompt=True,
+            return_tensors="pt",
+            padding=True,
+            return_dict=True,
+        ).to(model.device)
 
-print(harmful_hidden)
+        hidden = model(**toks, output_hidden_states=True, use_cache=False).hidden_states[layer_idx]
+        summed = hidden[:, POS, :].float().sum(dim=0)
+        total = summed if total is None else total + summed
 
-harmful_mean = torch.stack(harmful_hidden).mean(dim=0)
-harmless_mean = torch.stack(harmless_hidden).mean(dim=0)
+    return total / len(instructions)
 
-print(harmful_mean)
 
-refusal_dir = harmful_mean - harmless_mean
+refusal_dir = mean_hidden_state(harmful) - mean_hidden_state(harmless)
 refusal_dir = refusal_dir / refusal_dir.norm()
 
-print(refusal_dir)
-
-torch.save(refusal_dir, MODEL_ID.replace("/", "_") + "_refusal_dir.pt")
+path = refusal_dir_path()
+torch.save(refusal_dir, path)
+print(f"Saved refusal direction ({refusal_dir.shape[0]} dims) to {path}")

inference.py

@@ -1,97 +1,51 @@
-import einops
-import jaxtyping
 import torch
-import torch.nn as nn
-from typing import Optional, Tuple
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer, BitsAndBytesConfig
-from inspect import signature
+from transformers import TextStreamer
 
-torch.inference_mode()
+from common import MODEL_ID, load_model_and_tokenizer, project_out, refusal_dir_path
 
-MODEL_ID = "tiiuae/Falcon3-1B-Instruct"
-# MODEL_ID = "Qwen/Qwen3-1.7B"
-# MODEL_ID = "stabilityai/stablelm-2-zephyr-1_6b"
-# MODEL_ID = "Qwen/Qwen1.5-1.8B-Chat"
-# MODEL_ID = "Qwen/Qwen-1_8B-chat"
-# MODEL_ID = "google/gemma-1.1-2b-it"
-# MODEL_ID = "google/gemma-1.1-7b-it"
-# MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct"
+model, tokenizer = load_model_and_tokenizer()
 
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID,
-                                             trust_remote_code=True,
-                                             dtype=torch.float16,
-                                             device_map="cuda",
-                                             quantization_config=BitsAndBytesConfig(load_in_4bit=True,
-                                                                                    bnb_4bit_compute_dtype=torch.float16))
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+# Cast the direction to the model's device/dtype once, not on every layer call.
+refusal_dir = torch.load(refusal_dir_path()).to(device=model.device, dtype=model.dtype)
 
-refusal_dir = torch.load(MODEL_ID.replace("/", "_") + "_refusal_dir.pt")
 
+def ablation_pre_hook(_module, args, kwargs):
+    """Project the refusal direction out of the residual stream entering a layer.
 
-def direction_ablation_hook(activation: jaxtyping.Float[torch.Tensor, "... d_act"],
-                            direction: jaxtyping.Float[torch.Tensor, "d_act"]):
-    proj = einops.einsum(activation, direction.view(-1, 1),
-                         '... d_act, d_act single -> ... single') * direction
-    return activation - proj
+    A forward pre-hook on every decoder layer is equivalent to inserting an
+    ablation layer before each one, but without doubling num_hidden_layers or
+    wrestling with tuple-vs-tensor return signatures.
+    """
+    if args:
+        return (project_out(args[0], refusal_dir), *args[1:]), kwargs
+    kwargs["hidden_states"] = project_out(kwargs["hidden_states"], refusal_dir)
+    return args, kwargs
 
 
-# Some model developers thought it was stupid to pass a tuple of tuple of tuples around (rightfully so), but unfortunately now we have a divide
-sig = signature(model.model.layers[0].forward)
-simple = sig.return_annotation == torch.Tensor
-
-
-class AblationDecoderLayer(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.attention_type = "full_attention"
-
-    def forward(
-            self,
-            hidden_states: torch.Tensor,
-            attention_mask: Optional[torch.Tensor] = None,
-            position_ids: Optional[torch.LongTensor] = None,
-            past_key_value: Optional[Tuple[torch.Tensor]] = None,
-            output_attentions: Optional[bool] = False,
-            use_cache: Optional[bool] = False,
-            cache_position: Optional[torch.LongTensor] = None,
-            **kwargs,
-    ):
-        assert not output_attentions
-
-        ablated = direction_ablation_hook(hidden_states, refusal_dir.to(
-            hidden_states.device)).to(hidden_states.device)
-
-        if simple:
-            return ablated
-
-        outputs = (ablated,)
-
-        if use_cache:
-            outputs += (past_key_value,)
-
-        return outputs
-
-
-# for qwen 1 this needs to be changed to model.transformer.h
-for idx in reversed(range(len(model.model.layers))):
-    model.model.layers.insert(idx, AblationDecoderLayer())
-
-# bruh
-if hasattr(model, "config") and hasattr(model.config, "num_hidden_layers"):
-    model.config.num_hidden_layers *= 2
+# for qwen 1 this needs to be model.transformer.h
+for layer in model.model.layers:
+    layer.register_forward_pre_hook(ablation_pre_hook, with_kwargs=True)
 
 conversation = []
-
-streamer = TextStreamer(tokenizer)
+streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
 
 print(f"Chat with {MODEL_ID}:")
 while True:
-    prompt = input()
-    conversation.append({"role": "user", "content": prompt})
-    toks = tokenizer.apply_chat_template(conversation=conversation,
-                                         add_generation_prompt=True, return_tensors="pt")
-
-    gen = model.generate(toks.to(model.device), streamer=streamer, max_new_tokens=1337)
+    try:
+        prompt = input("> ")
+    except (EOFError, KeyboardInterrupt):
+        break
 
-    decoded = tokenizer.batch_decode(gen[0][len(toks[0]):], skip_special_tokens=True)
-    conversation.append({"role": "assistant", "content": "".join(decoded)})
+    conversation.append({"role": "user", "content": prompt})
+    toks = tokenizer.apply_chat_template(
+        conversation,
+        add_generation_prompt=True,
+        return_tensors="pt",
+        return_dict=True,
+    ).to(model.device)
+
+    with torch.inference_mode():
+        gen = model.generate(**toks, streamer=streamer, max_new_tokens=1337)
+
+    reply = tokenizer.decode(gen[0][toks["input_ids"].shape[1]:], skip_special_tokens=True)
+    conversation.append({"role": "assistant", "content": reply})

requirements.txt

@@ -1,7 +1,5 @@
-jaxtyping
 transformers
 tqdm
-einops
 torch
 bitsandbytes
-accelerate
+accelerate