194 lines
7.0 KiB
Python
194 lines
7.0 KiB
Python
#!/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
Detect Important Attention Heads
|
|
--------------------------------
|
|
• Single-GPU: Forces model/LoRA to specified GPU; blocks non-target devices like cuda:0.
|
|
• Multi-GPU: Exposes user-specified GPUs; uses device_map="auto" for slicing.
|
|
"""
|
|
|
|
import os, json, argparse
|
|
from pathlib import Path
|
|
from collections import defaultdict
|
|
from typing import Dict, List, Tuple
|
|
|
|
import torch, numpy as np
|
|
from tqdm import tqdm
|
|
from transformers import (
|
|
AutoConfig,
|
|
AutoTokenizer,
|
|
AutoModelForCausalLM,
|
|
)
|
|
from peft import PeftModel
|
|
|
|
# ========================= 1. Model Loader =========================
|
|
def load_generic_model(model_dir: str,
|
|
device,
|
|
device_map_cfg: Dict):
|
|
"""
|
|
device : torch.device('cuda:i') or cpu
|
|
device_map_cfg : {"": i} for single-GPU or "auto" for multi-GPU
|
|
"""
|
|
cfg = AutoConfig.from_pretrained(model_dir, trust_remote_code=True)
|
|
tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
|
|
|
|
if tokenizer.pad_token_id is None:
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|
tokenizer.pad_token_id = tokenizer.eos_token_id
|
|
tokenizer.padding_side = "right"
|
|
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
model_dir,
|
|
config=cfg,
|
|
torch_dtype=torch.bfloat16,
|
|
trust_remote_code=True,
|
|
attn_implementation="eager",
|
|
device_map=device_map_cfg,
|
|
)
|
|
return model, tokenizer
|
|
|
|
# ========================= 2. Score Function =========================
|
|
def trim_and_stack(rows: List[np.ndarray]) -> np.ndarray:
|
|
L = min(len(r) for r in rows)
|
|
return np.stack([r[:L] for r in rows])
|
|
|
|
def trim_to_same(a: np.ndarray, b: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
|
|
L = min(a.shape[1], b.shape[1])
|
|
return a[:, :L], b[:, :L]
|
|
|
|
def score_heads(normal: Dict[str, List[np.ndarray]],
|
|
conflict: Dict[str, List[np.ndarray]],
|
|
eps: float = 1e-6):
|
|
scores = {}
|
|
for k in normal:
|
|
if k not in conflict:
|
|
continue
|
|
try:
|
|
n = trim_and_stack(normal[k])
|
|
c = trim_and_stack(conflict[k])
|
|
n, c = trim_to_same(n, c)
|
|
except Exception as e:
|
|
print(f"⚠️ Skipped {k} (incompatible shape): {e}")
|
|
continue
|
|
if n.size == 0 or c.size == 0:
|
|
continue
|
|
|
|
frob = np.linalg.norm(n - c, ord="fro")
|
|
mean_shift = np.mean(np.abs(n.mean(1) - c.mean(1)))
|
|
|
|
def softmax(x):
|
|
e = np.exp(x - x.max(-1, keepdims=True))
|
|
return e / np.clip(e.sum(-1, keepdims=True), eps, None)
|
|
|
|
p, q = softmax(n), softmax(c)
|
|
kl = (p * (np.log(p + eps) - np.log(q + eps))).sum() / p.shape[0]
|
|
|
|
scores[k] = 0.4 * frob + 0.3 * mean_shift + 0.3 * kl
|
|
return scores
|
|
|
|
# ========================= 3. Extract Last-Token Attention =========================
|
|
@torch.inference_mode()
|
|
def extract_attention(model, tokenizer, sys_msg: str, usr_msg: str):
|
|
msgs = [
|
|
{"role": "system", "content": sys_msg},
|
|
{"role": "user", "content": usr_msg},
|
|
]
|
|
prompt = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
|
|
inputs = tokenizer(prompt, return_tensors="pt")
|
|
inputs = {k: v.to(model.device) for k, v in inputs.items()}
|
|
|
|
outs = model(**inputs, output_attentions=True)
|
|
gen = model.generate(**inputs, max_new_tokens=128)
|
|
decoded = tokenizer.decode(gen[0], skip_special_tokens=False)
|
|
|
|
# Extract only assistant portion
|
|
assistant_txt = decoded.split("assistant", 1)[-1].strip() if "assistant" in decoded else decoded.strip()
|
|
return outs.attentions, inputs["input_ids"], assistant_txt
|
|
|
|
# ========================= 4. Main Detection Procedure =========================
|
|
def detect_heads(json_path: str, model, tokenizer, out_dir: str):
|
|
with open(json_path, encoding="utf-8") as f:
|
|
raw = json.load(f)
|
|
|
|
grouped = defaultdict(lambda: {"normal": None, "conflict": None})
|
|
for s in raw:
|
|
base = s["id"].replace("_normal", "").replace("_conflict", "")
|
|
grouped[base][s["label"]] = s
|
|
|
|
normal, conflict = defaultdict(list), defaultdict(list)
|
|
responses = []
|
|
|
|
for _, pair in tqdm(grouped.items()):
|
|
for lbl in ("normal", "conflict"):
|
|
sample = pair[lbl]
|
|
if sample is None:
|
|
continue
|
|
usr_msg = f"{sample['task']} {sample['user_message']}".strip() if sample["user_message"].strip() else sample["task"]
|
|
attn, ids, output = extract_attention(model, tokenizer, sample["system_message"], usr_msg)
|
|
|
|
responses.append({
|
|
"id": sample["id"], "label": lbl, "output": output
|
|
})
|
|
|
|
n_layer = len(attn)
|
|
n_head = attn[0][0].shape[0]
|
|
last_tok = attn[0][0].shape[2] - 1
|
|
|
|
for L in range(n_layer):
|
|
for H in range(n_head):
|
|
vec = attn[L][0][H, last_tok].to(torch.float32).cpu().numpy()
|
|
key = f"L{L}_H{H}"
|
|
(normal if lbl == "normal" else conflict)[key].append(vec)
|
|
|
|
scores = score_heads(normal, conflict)
|
|
top10 = sorted(scores.items(), key=lambda kv: kv[1], reverse=True)[:10]
|
|
|
|
stem = Path(json_path).stem.replace("_instruction", "")
|
|
tgt = Path(out_dir) / f"{stem}_outputs"
|
|
tgt.mkdir(parents=True, exist_ok=True)
|
|
|
|
out_json = tgt / "important_heads.json"
|
|
with out_json.open("w", encoding="utf-8") as f:
|
|
json.dump({"important_heads": [(k, float(v)) for k, v in top10],
|
|
"responses": responses}, f, indent=2, ensure_ascii=False)
|
|
|
|
print(f"\n✅ Saved → {out_json}")
|
|
print("📌 Top-10 Important Heads:")
|
|
for h, s in top10:
|
|
print(f" {h:8s} │ {s:8.4f}")
|
|
|
|
# ========================= 5. CLI Entry =========================
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--json_path", required=True)
|
|
parser.add_argument("--model_path", required=True)
|
|
parser.add_argument("--cuda", type=int, nargs="+", default=[0], help="GPUs to use. Example: --cuda 0 or --cuda 0 1 2")
|
|
parser.add_argument("--output_dir", default="outputs")
|
|
parser.add_argument("--lora_path", default="", help="Optional: LoRA adapter path")
|
|
args = parser.parse_args()
|
|
|
|
# GPU setup
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join([str(i) for i in args.cuda])
|
|
|
|
if len(args.cuda) == 1:
|
|
idx = args.cuda[0]
|
|
device = torch.device(f"cuda:{idx}" if torch.cuda.is_available() else "cpu")
|
|
device_map = {"": 0} if device.type == "cuda" else {"": "cpu"}
|
|
else:
|
|
device = None
|
|
device_map = "auto"
|
|
|
|
print(f"🔵 Loading base model from {args.model_path} ...")
|
|
model, tok = load_generic_model(args.model_path, device, device_map)
|
|
|
|
if args.lora_path:
|
|
print(f"🟣 Loading LoRA from {args.lora_path} ...")
|
|
model = PeftModel.from_pretrained(model, args.lora_path, device_map=device_map)
|
|
model = model.merge_and_unload()
|
|
print("✅ LoRA merged.")
|
|
|
|
detect_heads(args.json_path, model, tok, args.output_dir)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|