487 lines
20 KiB
Python
487 lines
20 KiB
Python
"""
|
||
Head-LoRA Finetune (Refactored)
|
||
==========================================
|
||
Target: Align specific attention heads to ignore 'data' tokens using Flash Attention Interception.
|
||
"""
|
||
|
||
import os
|
||
import json
|
||
import csv
|
||
import argparse
|
||
import math
|
||
import random
|
||
import re
|
||
from typing import List, Tuple, Dict
|
||
|
||
import torch
|
||
import numpy as np
|
||
from tqdm import tqdm
|
||
from torch.utils.data import Dataset, DataLoader
|
||
from evallib import quick_eval_mmlu, quick_eval_asr_util
|
||
from transformers import (
|
||
AutoConfig, AutoTokenizer, AutoModelForCausalLM,
|
||
BitsAndBytesConfig, get_linear_schedule_with_warmup,
|
||
)
|
||
from transformers.models.llama.modeling_llama import ALL_ATTENTION_FUNCTIONS
|
||
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, PeftModel
|
||
from lib.tokenize_data_mask import apply_chat_tokenize_with_strip_and_mark
|
||
from lib.printcolor import print_tok_color
|
||
# -----------------------------------------------------------------------------
|
||
# 1. Flash Attention Spy Interceptor (核心機制)
|
||
# -----------------------------------------------------------------------------
|
||
print("🥷 Injecting Spy Interceptor into Flash Attention 2...")
|
||
|
||
if "flash_attention_2" not in ALL_ATTENTION_FUNCTIONS:
|
||
raise RuntimeError("Current environment does not support Flash Attention 2!")
|
||
|
||
orig_flash_attn = ALL_ATTENTION_FUNCTIONS["flash_attention_2"]
|
||
|
||
def wrapped_flash_attn(module, query, key, value, attention_mask, scaling, **kwargs):
|
||
"""
|
||
攔截 Flash Attention，並在 Breakpoint 驗證手算結果與原始 FA2 結果的一致性。
|
||
"""
|
||
# -------------------------------------------------------------------------
|
||
# [TASK A] 間諜行動：偷算 Last Token Weight
|
||
# -------------------------------------------------------------------------
|
||
# 1. 切片 (只看最後一個 Token)
|
||
q_last = query[..., -1:, :]
|
||
|
||
# 2. 處理 GQA (手動 Expand Key 和 Value)
|
||
num_heads = query.shape[1]
|
||
num_kv_heads = key.shape[1]
|
||
|
||
# 用於後面計算 Manual Output
|
||
k_expanded = key
|
||
v_expanded = value
|
||
|
||
if num_heads != num_kv_heads:
|
||
n_rep = num_heads // num_kv_heads
|
||
k_expanded = key.repeat_interleave(n_rep, dim=1)
|
||
v_expanded = value.repeat_interleave(n_rep, dim=1) # Value 也要 Expand 才能乘
|
||
|
||
# 3. 手動計算 Tiny Attention Weight [Batch, Heads, 1, Seq]
|
||
# Q: [B, H, 1, D], K: [B, H, S, D] -> QK^T: [B, H, 1, S]
|
||
attn_weights_tiny = torch.matmul(q_last, k_expanded.transpose(2, 3)) * scaling
|
||
|
||
# 4. 處理 Mask
|
||
if attention_mask is not None:
|
||
if attention_mask.dim() == 4:
|
||
if attention_mask.size(2) > 1:
|
||
mask_slice = attention_mask[..., -1:, :]
|
||
attn_weights_tiny = attn_weights_tiny + mask_slice
|
||
else:
|
||
attn_weights_tiny = attn_weights_tiny + attention_mask
|
||
elif attention_mask.dim() == 2:
|
||
mask_expanded = attention_mask[:, None, None, :]
|
||
if attention_mask.dtype == torch.bool or (attention_mask.max() <= 1.0 and attention_mask.min() >= 0.0):
|
||
min_dtype = torch.finfo(attn_weights_tiny.dtype).min
|
||
attn_weights_tiny = attn_weights_tiny.masked_fill(mask_expanded == 0, min_dtype)
|
||
else:
|
||
attn_weights_tiny = attn_weights_tiny + mask_expanded
|
||
|
||
# 5. Softmax (得到機率分佈 P)
|
||
# 注意：這裡轉成 float32 做 softmax 以求精確，實際 FA2 內部可能是 fp16/bf16
|
||
attn_probs = torch.nn.functional.softmax(attn_weights_tiny, dim=-1, dtype=torch.float32).to(query.dtype)
|
||
|
||
# 填入 Config (原本的邏輯)
|
||
retrieve_map = getattr(module.config, "retrieve_attn_map", None)
|
||
if retrieve_map is not None and hasattr(module, "layer_idx") and module.layer_idx in retrieve_map:
|
||
target_heads_dict = retrieve_map[module.layer_idx]
|
||
for h_idx in target_heads_dict.keys():
|
||
# 存入 config 用於 Loss 計算
|
||
target_heads_dict[h_idx] = attn_probs[:, h_idx, :, :]
|
||
|
||
# -------------------------------------------------------------------------
|
||
# [DEBUG] 驗證環節：手算 Output vs Flash Attention Output
|
||
# -------------------------------------------------------------------------
|
||
|
||
# 1. 執行原始 Flash Attention
|
||
orig_result = orig_flash_attn(module, query, key, value, attention_mask, scaling=scaling, **kwargs)
|
||
debug=False
|
||
if debug:
|
||
orig_tensor = orig_result[0]
|
||
# 2. 手動計算 Output ( O = Attention_Probs * V )
|
||
# attn_probs: [B, H, 1, S]
|
||
# v_expanded: [B, H, S, D]
|
||
# manual_out: [B, H, 1, D]
|
||
manual_out = torch.matmul(attn_probs, v_expanded)
|
||
|
||
# 3. 對齊形狀以便比較
|
||
# FA2 output 通常是 [Batch, Seq, Heads, Dim]
|
||
# 我們取最後一個 token: [Batch, 1, Heads, Dim]
|
||
orig_out_last = orig_tensor[:, -1:, :, :]
|
||
|
||
# 手算結果原本是 [Batch, Heads, 1, Dim]，轉置成 [Batch, 1, Heads, Dim]
|
||
manual_out_check = manual_out.transpose(1, 2)
|
||
|
||
# 4. 計算誤差
|
||
diff = (orig_out_last - manual_out_check).abs()
|
||
max_diff = diff.max().item()
|
||
mean_diff = diff.mean().item()
|
||
|
||
print(f"\n🕵️ [Layer {getattr(module, 'layer_idx', '?')}] Validation:")
|
||
print(f" Max Diff: {max_diff:.8f}")
|
||
print(f" Mean Diff: {mean_diff:.8f}")
|
||
|
||
# 如果誤差過大 (例如 > 1e-3)，自動暫停檢查
|
||
# 注意：BF16 下 FlashAttention 和標準 Matmul 會有一定精度差異是正常的
|
||
if max_diff > 1e-2:
|
||
print("⚠️ Warning: Large difference detected!")
|
||
|
||
# 呼叫 breakpoint 讓你進去檢查
|
||
# 在 pdb 輸入:
|
||
# p manual_out_check[0,0,0,:5]
|
||
# p orig_out_last[0,0,0,:5]
|
||
# 來對比數值
|
||
breakpoint()
|
||
|
||
return orig_result
|
||
|
||
# 替換全局函數
|
||
ALL_ATTENTION_FUNCTIONS["flash_attention_2"] = wrapped_flash_attn
|
||
print("✅ Spy Interceptor installed successfully!")
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# 2. Utils: Model & Heads
|
||
# -----------------------------------------------------------------------------
|
||
SEED = 42
|
||
random.seed(SEED)
|
||
np.random.seed(SEED)
|
||
torch.manual_seed(SEED)
|
||
|
||
def get_lora_targets(model, layers: List[int]) -> List[str]:
|
||
mtype = (getattr(model.config, "model_type", "") or "").lower()
|
||
if "llama" in mtype or "mistral" in mtype:
|
||
return [f"model.layers.{i}.self_attn.{p}" for i in layers for p in ("q_proj", "k_proj")]
|
||
# Fallback / Other architectures
|
||
cand = []
|
||
for name, _ in model.named_modules():
|
||
if any(f".{i}." in name for i in layers) and name.split(".")[-1] in {"q_proj", "k_proj"}:
|
||
cand.append(name)
|
||
return cand
|
||
|
||
def load_model(model_path: str):
|
||
cfg = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
|
||
tok = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, use_fast=True,padding_side="left")
|
||
|
||
if tok.pad_token_id is None:
|
||
tok.pad_token = tok.eos_token
|
||
tok.pad_token_id = tok.eos_token_id
|
||
|
||
bnb_cfg = BitsAndBytesConfig(
|
||
load_in_4bit=True,
|
||
bnb_4bit_compute_dtype=torch.float16,
|
||
bnb_4bit_use_double_quant=True,
|
||
bnb_4bit_quant_type="nf4",
|
||
)
|
||
model = AutoModelForCausalLM.from_pretrained(
|
||
model_path,
|
||
config=cfg,
|
||
quantization_config=bnb_cfg,
|
||
device_map="auto",
|
||
trust_remote_code=True,
|
||
attn_implementation="flash_attention_2",
|
||
)
|
||
return model, tok, tok # use same tokenizer for simplicity
|
||
|
||
def get_target_structure(heads_file: str, topk_spec: str) -> Dict[int, List[int]]:
|
||
"""Load heads from file and parse into {layer: [heads]} structure."""
|
||
if not os.path.exists(heads_file):
|
||
raise FileNotFoundError(f"Heads file not found: {heads_file}")
|
||
|
||
with open(heads_file, "r") as f:
|
||
ihead_list = json.load(f)
|
||
if type(ihead_list[0]) == str:
|
||
all_heads = [(str(tag), 0.0) for tag in ihead_list]
|
||
else:
|
||
all_heads = ihead_list
|
||
|
||
# Filter Top-K
|
||
count = len(all_heads)
|
||
if topk_spec.endswith("p"):
|
||
count = max(1, math.ceil(float(topk_spec[:-1]) / 100.0 * len(all_heads)))
|
||
else:
|
||
count = int(topk_spec)
|
||
|
||
target_heads = all_heads[:min(count, len(all_heads))]
|
||
|
||
# Parse into structure
|
||
structure = {}
|
||
for tag, _ in target_heads:
|
||
try:
|
||
# tag format: "L22_H7"
|
||
parts = tag.split('_')
|
||
l = int(parts[0][1:])
|
||
h = int(parts[1][1:])
|
||
structure.setdefault(l, []).append(h)
|
||
except Exception:
|
||
continue
|
||
|
||
print(f"🎯 Selected {len(target_heads)} heads from {topk_spec}.")
|
||
return structure
|
||
|
||
def discover_existing_adapter(out_dir: str):
|
||
if not os.path.isdir(out_dir): return None, 0, 0
|
||
candidates = []
|
||
if os.path.exists(os.path.join(out_dir, "adapter_config.json")):
|
||
candidates.append((0, 0, out_dir))
|
||
|
||
for entry in os.listdir(out_dir):
|
||
path = os.path.join(out_dir, entry)
|
||
if os.path.isdir(path) and os.path.exists(os.path.join(path, "adapter_config.json")):
|
||
m = re.match(r"^batch_(\d+)_(\d+)$", entry)
|
||
if m:
|
||
epoch_idx = int(m.group(1))
|
||
batch_idx = int(m.group(2))
|
||
candidates.append((epoch_idx, batch_idx, path))
|
||
|
||
if not candidates: return None, 0, 0
|
||
candidates.sort(key=lambda x: (x[0], x[1]))
|
||
return candidates[-1][2], candidates[-1][0], candidates[-1][1]
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# 3. Loss Function
|
||
# -----------------------------------------------------------------------------
|
||
def head_attention_loss(tuned_map, base_map, data_mask, lambda_data=1.0, eps=1e-8):
|
||
mask_data = data_mask.bool() # [B,S]
|
||
mask_valid = ~mask_data
|
||
|
||
if mask_valid.dim() == 2:
|
||
mask_valid = mask_valid.unsqueeze(1) # [B,1,S]
|
||
mask_data = mask_data.unsqueeze(1) # [B,1,S]
|
||
|
||
total = 0.0
|
||
n = 0
|
||
|
||
for l, heads in tuned_map.items():
|
||
for h, tuned in heads.items():
|
||
base = base_map[l][h]
|
||
if tuned is None or base is None:
|
||
continue
|
||
|
||
tuned = tuned.float() # [B,1,S]
|
||
base = base.float().to(tuned.device)
|
||
|
||
# 1) teacher: base 已經是 blind 分佈（理想上 data=0），但我們也只取 valid 區域
|
||
base_v = base * mask_valid.float()
|
||
base_v = base_v / base_v.sum(dim=-1, keepdim=True).clamp_min(eps)
|
||
|
||
# 2) student: tuned 在 valid 區域重新 normalize
|
||
tuned_v = tuned * mask_valid.float()
|
||
tuned_v_sum = tuned_v.sum(dim=-1, keepdim=True).clamp_min(eps)
|
||
tuned_v = tuned_v / tuned_v_sum
|
||
|
||
# 3) KL(base || tuned)
|
||
kl = (base_v * (torch.log(base_v + eps) - torch.log(tuned_v + eps))).sum(dim=-1).mean()
|
||
|
||
# 4) data mass penalty（不重正規化，直接壓 data 的注意力總量）
|
||
data_mass = (tuned * mask_data.float()).sum(dim=-1).mean()
|
||
|
||
total = total + kl + lambda_data * data_mass
|
||
n += 1
|
||
|
||
return total / max(n, 1)
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# 4. Data Loading
|
||
# -----------------------------------------------------------------------------
|
||
class JsonlMessagesDS(Dataset):
|
||
def __init__(self, jsonl_path: str):
|
||
self.samples = []
|
||
with open(jsonl_path, "r", encoding="utf-8") as f:
|
||
for line in f:
|
||
if line.strip():
|
||
obj = json.loads(line)
|
||
self.samples.append(obj)
|
||
random.shuffle(self.samples)
|
||
print(f"📊 Loaded {len(self.samples)} samples.")
|
||
|
||
def __len__(self): return len(self.samples)
|
||
def __getitem__(self, idx): return self.samples[idx]
|
||
|
||
def collate(batch, tokenizer):
|
||
input_ids, attention_mask, data_mask = apply_chat_tokenize_with_strip_and_mark(
|
||
batch, tokenizer, device="cpu", add_generation_prompt=True,
|
||
encode_kwargs = {"padding_side" :"left" },
|
||
mode="custom_mask == 'inst'",
|
||
custom_mask_identifier={"data": ["<data>", "</data>"], "inst": ["<inst>", "</inst>"]},
|
||
return_tensors="pt"
|
||
)
|
||
|
||
return {"input_ids": input_ids, "attention_mask": attention_mask, "data_mask": data_mask}
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# 5. Training Loop
|
||
# -----------------------------------------------------------------------------
|
||
def save_model(model, tok, out_dir, epoch, batch_idx, data_path):
|
||
"""Safely saves model by temporarily cleaning config, then logs eval results."""
|
||
print("Saving checkpoint...")
|
||
|
||
# 🚑 CRITICAL: Temporarily remove the tensor map from config to prevent DeepCopy error
|
||
temp_map = getattr(model.config, "retrieve_attn_map", None)
|
||
if temp_map is not None:
|
||
del model.config.retrieve_attn_map
|
||
|
||
save_dir = os.path.join(out_dir, f"batch_{epoch}_{batch_idx}")
|
||
model.save_pretrained(save_dir)
|
||
tok.save_pretrained(save_dir)
|
||
|
||
print(f"✅ Saved to {save_dir}")
|
||
|
||
mllu_result = quick_eval_mmlu(model, tok)
|
||
asr_result = quick_eval_asr_util(
|
||
model,
|
||
tok,
|
||
training_data_path=data_path,
|
||
)
|
||
log_path = os.path.join(out_dir, "training_log.csv")
|
||
write_header = not os.path.exists(log_path)
|
||
with open(log_path, "a", newline="") as f:
|
||
writer = csv.writer(f)
|
||
attack_columns = [
|
||
("naive", "a_naive", "v_naive"),
|
||
("ignore", "a_ignore", "v_ignore"),
|
||
("escape_separation", "a_escape", "v_escape"),
|
||
("completion_realcmb", "a_cmb", "v_cmb"),
|
||
("conv_attack", "a_conv", "v_conv"),
|
||
("none", "a_none", "v_none"),
|
||
]
|
||
if write_header:
|
||
header = ["epoch", "batch", "mmlu score", "asr_mode", "asr_status"]
|
||
for _, a_col, v_col in attack_columns:
|
||
header.append(a_col)
|
||
header.append(v_col)
|
||
writer.writerow(header)
|
||
row = [epoch, batch_idx, mllu_result.get("accuracy"), asr_result.get("eval_mode"), asr_result.get("status")]
|
||
metrics = asr_result.get("metrics", {}) if asr_result.get("status") == "ok" else {}
|
||
for attack, _, _ in attack_columns:
|
||
values = metrics.get(attack, {})
|
||
row.append(values.get("asr"))
|
||
row.append(values.get("valid_rate"))
|
||
writer.writerow(row)
|
||
return save_dir
|
||
|
||
def tune_new(model, tok, tok_inf, target_structure, data_path, out_dir, epochs, bs, lr, resume_adapter=None, start_epoch=0, start_batch_idx=0, batch_save_interval=-1):
|
||
# Setup LoRA
|
||
layers = sorted(target_structure.keys())
|
||
targets = get_lora_targets(model, layers)
|
||
|
||
lora_cfg = LoraConfig(r=32, lora_alpha=16, bias="none", target_modules=targets, task_type="CAUSAL_LM")
|
||
model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=False)
|
||
|
||
if resume_adapter:
|
||
model = PeftModel.from_pretrained(model, resume_adapter, is_trainable=True)
|
||
print(f"♻️ Resumed LoRA: {resume_adapter}")
|
||
else:
|
||
model = get_peft_model(model, lora_cfg)
|
||
|
||
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||
total = sum(p.numel() for p in model.parameters())
|
||
ratio = trainable / total * 100
|
||
print(f"🔧 Trainable parameters: {trainable:,} / {total:,} ({ratio:.4f}% of total)")
|
||
|
||
|
||
dl = DataLoader(JsonlMessagesDS(data_path), batch_size=bs, shuffle=True, collate_fn=lambda b: collate(b, tok))
|
||
opt = torch.optim.AdamW(model.parameters(), lr=lr)
|
||
sch = get_linear_schedule_with_warmup(opt, int(0.05 * epochs * len(dl)), epochs * len(dl))
|
||
|
||
current_idx = start_epoch
|
||
|
||
for ep in range(epochs):
|
||
if ep < start_epoch:
|
||
continue
|
||
pbar = tqdm(dl, desc=f"Ep {ep+1}/{epochs}")
|
||
for idx, batch in enumerate(pbar):
|
||
if ep == start_epoch and idx < start_batch_idx:
|
||
continue
|
||
device = next(model.parameters()).device
|
||
MAX_LEN = 50000
|
||
|
||
input_ids = batch["input_ids"][:, :MAX_LEN].to(device)
|
||
attn_mask = batch["attention_mask"][:, :MAX_LEN].to(device) # 原始 Mask (看得到全部)
|
||
data_mask = batch["data_mask"][:, :MAX_LEN].to(device) # Data 標記
|
||
|
||
# -------------------------------------------------------
|
||
# [Step 1] 建構 "Blind Mask" 給 Base Model
|
||
# -------------------------------------------------------
|
||
# data_mask 為 1 的地方是要隱藏的。
|
||
# attn_mask 為 1 是可見，0 是 Padding。
|
||
# 我們要讓 Base Model 在 data_mask 為 1 的地方也變成不可見 (0)。
|
||
# 邏輯：blind_mask = attn_mask AND (NOT data_mask)
|
||
|
||
blind_attn_mask = attn_mask * (~data_mask).long()
|
||
# --- A. Base Model Pass (Reference: Blind) ---
|
||
model.eval()
|
||
base_map_container = {l: {h: None for h in h_list} for l, h_list in target_structure.items()}
|
||
model.config.retrieve_attn_map = base_map_container
|
||
|
||
with torch.no_grad(), model.disable_adapter():
|
||
# 關鍵：這裡傳入 blind_attn_mask
|
||
# Base Model 的 Flash Attn 會收到這個 Mask，
|
||
# 導致它計算出的 Attention Weight 在 Data 區域直接被 Mask 成 -inf (softmax後為0)，
|
||
# 剩餘的 Valid 區域會重新歸一化 (Sum=1)。
|
||
model(input_ids=input_ids, attention_mask=blind_attn_mask, output_attentions=False, use_cache=False)
|
||
base_attns_map = base_map_container
|
||
|
||
# --- B. Tuned Model Pass (Training: See Everything) ---
|
||
model.train()
|
||
tuned_map_container = {l: {h: None for h in h_list} for l, h_list in target_structure.items()}
|
||
model.config.retrieve_attn_map = tuned_map_container
|
||
|
||
# Tuned Model 仍然傳入原始 attn_mask (它看得到 Data，但我們要訓練它忽略)
|
||
model(input_ids=input_ids, attention_mask=attn_mask, output_attentions=False, use_cache=False)
|
||
tuned_attns_map = model.config.retrieve_attn_map
|
||
|
||
# --- C. Loss & Step ---
|
||
# 計算 Loss
|
||
loss = head_attention_loss(tuned_attns_map, base_attns_map, data_mask)
|
||
|
||
loss.backward()
|
||
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
|
||
opt.step()
|
||
sch.step()
|
||
opt.zero_grad()
|
||
|
||
del model.config.retrieve_attn_map
|
||
pbar.set_postfix(loss=f"{loss.item():.4f}")
|
||
if batch_save_interval > 0 and (idx + 1) % batch_save_interval == 0:
|
||
save_model(model, tok, out_dir, current_idx, idx, data_path)
|
||
save_model(model, tok, out_dir, current_idx, idx, data_path)
|
||
current_idx += 1
|
||
|
||
# -----------------------------------------------------------------------------
|
||
# 6. Main
|
||
# -----------------------------------------------------------------------------
|
||
def main():
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--model_path", required=True)
|
||
parser.add_argument("--data_path", default="../3-1_model_training_preprocess/inj-likechen/crafted_instruction_data_tri_injection_qa.jsonl")
|
||
parser.add_argument("--head_path", default="../2-2_head_identification/head_scoring/llama31-8b_injsq_dev/heads_sorted/user_prop_inst_0.1.json")
|
||
|
||
parser.add_argument("--output_dir", default="outputs_lora")
|
||
parser.add_argument("--lora_path", default=None)
|
||
parser.add_argument("--epochs", type=int, default=3)
|
||
parser.add_argument("--batch_size", type=int, default=4)
|
||
parser.add_argument("--batch-save-interval", type=int, default=-1)
|
||
parser.add_argument("--lr", type=float, default=1e-4)
|
||
parser.add_argument("--topk", type=str, default="18.75p")
|
||
args = parser.parse_args()
|
||
|
||
resume_adapter, start_epoch, start_batch_idx = discover_existing_adapter(args.output_dir)
|
||
if args.lora_path:
|
||
resume_adapter, start_epoch, start_batch_idx = args.lora_path, 0, 0
|
||
|
||
model, tok, tok_inf = load_model(args.model_path)
|
||
target_structure = get_target_structure(args.head_path, args.topk)
|
||
os.makedirs(args.output_dir, exist_ok=True)
|
||
tune_new(
|
||
model, tok, tok_inf, target_structure,
|
||
args.data_path, args.output_dir,
|
||
args.epochs, args.batch_size, args.lr,
|
||
resume_adapter=resume_adapter, start_epoch=start_epoch, start_batch_idx=start_batch_idx,
|
||
batch_save_interval=args.batch_save_interval
|
||
)
|
||
|
||
if __name__ == "__main__":
|
||
main()
|