first commit
This commit is contained in:
HenryChou020514
334
Codes/2-2_head_identification/lib/head_mask_inference.py
Normal file
334
Codes/2-2_head_identification/lib/head_mask_inference.py
Normal file
@ -0,0 +1,334 @@
|
||||
import math
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
|
||||
def select_heads(head_list: List[str], topk: int) -> List[str]:
|
||||
if topk <= 0:
|
||||
return []
|
||||
return head_list[: min(topk, len(head_list))]
|
||||
|
||||
|
||||
def build_head_map(head_names: List[str]) -> dict:
|
||||
head_map = {}
|
||||
for head_name in head_names:
|
||||
if not head_name.startswith("L"):
|
||||
raise ValueError(f"Invalid head name: {head_name}")
|
||||
try:
|
||||
layer_part, head_part = head_name.split("_", 1)
|
||||
layer_idx = int(layer_part[1:])
|
||||
head_idx = int(head_part[1:])
|
||||
except Exception as exc:
|
||||
raise ValueError(f"Invalid head name: {head_name}") from exc
|
||||
head_map.setdefault(layer_idx, []).append(head_idx)
|
||||
return head_map
|
||||
|
||||
|
||||
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)
|
||||
if tok.pad_token_id is None:
|
||||
tok.pad_token = tok.eos_token
|
||||
tok.pad_token_id = tok.eos_token_id
|
||||
tok.padding_side = "left"
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_path,
|
||||
config=cfg,
|
||||
device_map="auto",
|
||||
torch_dtype=torch.bfloat16,
|
||||
trust_remote_code=True,
|
||||
attn_implementation="eager",
|
||||
)
|
||||
model.eval()
|
||||
return model, tok
|
||||
|
||||
|
||||
def mask_attention(attn_head: torch.Tensor, data_positions: List[int]) -> torch.Tensor:
|
||||
if not data_positions:
|
||||
return attn_head
|
||||
masked = attn_head.clone()
|
||||
masked[:, data_positions] = 0.0
|
||||
return masked
|
||||
|
||||
|
||||
def build_head_summaries(
|
||||
attentions,
|
||||
selected_heads: List[str],
|
||||
data_positions_batch: List[List[int]],
|
||||
n_layers: int,
|
||||
n_heads: int,
|
||||
):
|
||||
summaries = []
|
||||
for batch_idx, data_positions in enumerate(data_positions_batch):
|
||||
head_info = {}
|
||||
for head_name in selected_heads:
|
||||
if not head_name.startswith("L"):
|
||||
raise ValueError(f"Invalid head name: {head_name}")
|
||||
try:
|
||||
layer_part, head_part = head_name.split("_", 1)
|
||||
layer_idx = int(layer_part[1:])
|
||||
head_idx = int(head_part[1:])
|
||||
except Exception as exc:
|
||||
raise ValueError(f"Invalid head name: {head_name}") from exc
|
||||
if layer_idx < 0 or layer_idx >= n_layers or head_idx < 0 or head_idx >= n_heads:
|
||||
raise ValueError(f"Head out of range for model: {head_name}")
|
||||
attn_head = attentions[layer_idx][batch_idx, head_idx]
|
||||
masked = mask_attention(attn_head, data_positions)
|
||||
pre_sum = attn_head[:, data_positions].sum().float().item() if data_positions else 0.0
|
||||
post_sum = masked[:, data_positions].sum().float().item() if data_positions else 0.0
|
||||
head_info[head_name] = {
|
||||
"pre_data_attention_sum": pre_sum,
|
||||
"post_data_attention_sum": post_sum,
|
||||
}
|
||||
summaries.append(head_info)
|
||||
return summaries
|
||||
|
||||
|
||||
def total_attn_to_data_batch(attn, data_positions_by_sample: List[List[int]]) -> List[float]:
|
||||
totals = []
|
||||
for b, data_positions in enumerate(data_positions_by_sample):
|
||||
if not data_positions:
|
||||
totals.append(0.0)
|
||||
continue
|
||||
total = 0.0
|
||||
for layer_attn in attn:
|
||||
total += layer_attn[b, :, -1, data_positions].sum().float().item()
|
||||
totals.append(total)
|
||||
return totals
|
||||
|
||||
|
||||
def debug_print_attention_totals(
|
||||
data_indices,
|
||||
unmasked_attn,
|
||||
masked_attn,
|
||||
data_positions_batch: List[List[int]],
|
||||
debug: bool = False,
|
||||
):
|
||||
if not debug:
|
||||
return
|
||||
unmasked_totals = total_attn_to_data_batch(unmasked_attn, data_positions_batch)
|
||||
masked_totals = total_attn_to_data_batch(masked_attn, data_positions_batch)
|
||||
for idx, (u, m) in enumerate(zip(unmasked_totals, masked_totals)):
|
||||
sample_id = data_indices[idx] if idx < len(data_indices) else idx
|
||||
print(f"[DEBUG] idx={sample_id} unmasked_attn_sum={u:.6f}")
|
||||
print(f"[DEBUG] idx={sample_id} masked_attn_sum={m:.6f}")
|
||||
|
||||
|
||||
def debug_print_head_summaries(
|
||||
data_indices,
|
||||
head_summaries: List[dict],
|
||||
debug: bool = False,
|
||||
):
|
||||
if not debug:
|
||||
return
|
||||
for idx, head_info in enumerate(head_summaries):
|
||||
sample_id = data_indices[idx] if idx < len(data_indices) else idx
|
||||
print(f"[DEBUG] idx={sample_id} head_summaries={len(head_info)}")
|
||||
|
||||
|
||||
def _make_head_mask_hook(
|
||||
layer_idx: int,
|
||||
head_indices: List[int],
|
||||
data_positions_by_sample: List[List[int]],
|
||||
num_heads: int,
|
||||
debug: bool = False,
|
||||
):
|
||||
head_indices = list(sorted(set(head_indices)))
|
||||
seen = {"printed": False}
|
||||
|
||||
def _hook(_module, args, kwargs):
|
||||
if not data_positions_by_sample:
|
||||
return None
|
||||
|
||||
attention_mask = None
|
||||
if kwargs is not None:
|
||||
attention_mask = kwargs.get("attention_mask", None)
|
||||
if attention_mask is None and len(args) >= 2:
|
||||
attention_mask = args[1]
|
||||
|
||||
if attention_mask is None:
|
||||
return None
|
||||
|
||||
bsz, mask_heads, q_len, k_len = attention_mask.shape
|
||||
|
||||
if mask_heads == 1 and num_heads > 1:
|
||||
attention_mask = attention_mask.expand(bsz, num_heads, q_len, k_len).clone()
|
||||
|
||||
if attention_mask.dtype == torch.bool:
|
||||
val_to_fill = True
|
||||
else:
|
||||
val_to_fill = torch.finfo(attention_mask.dtype).min
|
||||
|
||||
for b in range(bsz):
|
||||
masked_positions = [p for p in data_positions_by_sample[b] if p < k_len]
|
||||
if not masked_positions:
|
||||
continue
|
||||
masked_pos_tensor = torch.tensor(masked_positions, device=attention_mask.device)
|
||||
if debug and layer_idx == 0 and not seen["printed"]:
|
||||
target_head = head_indices[0] if head_indices else 0
|
||||
sample_pos = masked_positions[:3]
|
||||
if sample_pos:
|
||||
sample_vals = attention_mask[b, target_head, -1, sample_pos].detach().cpu().tolist()
|
||||
print(f"[DEBUG] L{layer_idx}: head={target_head} sample_before={sample_vals}")
|
||||
|
||||
for h in head_indices:
|
||||
attention_mask[b, h].index_fill_(1, masked_pos_tensor, val_to_fill)
|
||||
|
||||
if debug and layer_idx == 0 and not seen["printed"]:
|
||||
target_head = head_indices[0] if head_indices else 0
|
||||
sample_pos = masked_positions[:3]
|
||||
if sample_pos:
|
||||
sample_vals = attention_mask[b, target_head, -1, sample_pos].detach().cpu().tolist()
|
||||
print(f"[DEBUG] L{layer_idx}: head={target_head} sample_after={sample_vals}")
|
||||
seen["printed"] = True
|
||||
|
||||
if kwargs is not None and "attention_mask" in kwargs:
|
||||
kwargs["attention_mask"] = attention_mask
|
||||
return args, kwargs
|
||||
|
||||
new_args = list(args)
|
||||
if len(new_args) >= 2:
|
||||
new_args[1] = attention_mask
|
||||
return tuple(new_args), kwargs
|
||||
|
||||
return None
|
||||
|
||||
return _hook
|
||||
|
||||
|
||||
def _install_mask_hooks(
|
||||
model,
|
||||
head_map: dict,
|
||||
data_positions_by_sample: List[List[int]],
|
||||
num_heads: int,
|
||||
debug: bool = False,
|
||||
):
|
||||
handles = []
|
||||
layers = getattr(getattr(model, "model", None), "layers", None)
|
||||
if layers is None:
|
||||
raise ValueError("Unsupported model layout: missing model.model.layers")
|
||||
for layer_idx, head_indices in head_map.items():
|
||||
if layer_idx < 0 or layer_idx >= len(layers):
|
||||
raise ValueError(f"Layer index out of range: L{layer_idx}")
|
||||
attn = getattr(layers[layer_idx], "self_attn", None)
|
||||
if attn is None:
|
||||
raise ValueError(f"Layer L{layer_idx} missing self_attn")
|
||||
hook = _make_head_mask_hook(layer_idx, head_indices, data_positions_by_sample, num_heads, debug=debug)
|
||||
handles.append(attn.register_forward_pre_hook(hook, with_kwargs=True))
|
||||
return handles
|
||||
|
||||
|
||||
class MaskedCausalLM:
|
||||
def __init__(self, model, head_map: dict, num_heads: int, debug: bool = False):
|
||||
self._model = model
|
||||
self._head_map = head_map
|
||||
self._num_heads = num_heads
|
||||
self._debug = debug
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self._model, name)
|
||||
|
||||
def _validate_data_positions(self, data_positions_batch, batch_size):
|
||||
if data_positions_batch is None:
|
||||
raise ValueError("data_positions_batch is required for masked inference.")
|
||||
if batch_size is not None and len(data_positions_batch) != batch_size:
|
||||
raise ValueError(
|
||||
f"data_positions_batch size {len(data_positions_batch)} does not match batch size {batch_size}."
|
||||
)
|
||||
|
||||
def _with_masking(self, data_positions_batch, fn):
|
||||
if not self._head_map:
|
||||
return fn()
|
||||
handles = _install_mask_hooks(
|
||||
self._model,
|
||||
self._head_map,
|
||||
data_positions_batch,
|
||||
self._num_heads,
|
||||
debug=self._debug,
|
||||
)
|
||||
try:
|
||||
return fn()
|
||||
finally:
|
||||
for handle in handles:
|
||||
handle.remove()
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
data_positions_batch = kwargs.pop("data_positions_batch", None)
|
||||
batch_size = None
|
||||
input_ids = kwargs.get("input_ids", None)
|
||||
if input_ids is None and args:
|
||||
input_ids = args[0]
|
||||
if input_ids is not None and hasattr(input_ids, "shape"):
|
||||
batch_size = input_ids.shape[0]
|
||||
self._validate_data_positions(data_positions_batch, batch_size)
|
||||
return self._with_masking(data_positions_batch, lambda: self._model(*args, **kwargs))
|
||||
|
||||
def generate(self, *args, **kwargs):
|
||||
data_positions_batch = kwargs.pop("data_positions_batch", None)
|
||||
batch_size = None
|
||||
input_ids = kwargs.get("input_ids", None)
|
||||
if input_ids is None and args:
|
||||
input_ids = args[0]
|
||||
if input_ids is not None and hasattr(input_ids, "shape"):
|
||||
batch_size = input_ids.shape[0]
|
||||
self._validate_data_positions(data_positions_batch, batch_size)
|
||||
return self._with_masking(data_positions_batch, lambda: self._model.generate(*args, **kwargs))
|
||||
|
||||
|
||||
def build_masked_model(model, head_list: List[str], topk: str, debug: bool = False):
|
||||
if not isinstance(head_list, list) or not head_list:
|
||||
raise ValueError("head_list must be a non-empty list like ['L1H5', 'L15H23'].")
|
||||
if len(head_list) <= 0:
|
||||
topk_count = 0
|
||||
elif topk is None:
|
||||
topk_count = len(head_list)
|
||||
else:
|
||||
topk_str = str(topk).strip().lower()
|
||||
if topk_str.endswith("p"):
|
||||
pct = float(topk_str[:-1])
|
||||
if pct <= 0:
|
||||
topk_count = 0
|
||||
else:
|
||||
topk_count = max(1, int(math.ceil(len(head_list) * pct / 100.0)))
|
||||
else:
|
||||
topk_count = max(0, int(topk_str))
|
||||
selected_heads = select_heads(head_list, topk_count)
|
||||
num_heads = getattr(model.config, "num_attention_heads", None)
|
||||
if num_heads is None:
|
||||
raise ValueError("Model config missing num_attention_heads.")
|
||||
head_map = build_head_map(selected_heads)
|
||||
masked_model = MaskedCausalLM(model, head_map, num_heads, debug=debug)
|
||||
return masked_model, selected_heads
|
||||
|
||||
def _generate_batch(model, tok, input_ids_batch, attention_mask_batch, max_new_tokens, data_positions_batch=None):
|
||||
if not input_ids_batch:
|
||||
return []
|
||||
input_ids_tensor = torch.tensor(input_ids_batch, dtype=torch.long, device=model.device)
|
||||
attention_mask_tensor = torch.tensor(attention_mask_batch, dtype=torch.long, device=model.device)
|
||||
if data_positions_batch is None or (type(model) != MaskedCausalLM):
|
||||
out = model.generate(
|
||||
input_ids=input_ids_tensor,
|
||||
attention_mask=attention_mask_tensor,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
eos_token_id=tok.eos_token_id,
|
||||
pad_token_id=tok.pad_token_id,
|
||||
)
|
||||
else:
|
||||
out = model.generate(
|
||||
input_ids=input_ids_tensor,
|
||||
attention_mask=attention_mask_tensor,
|
||||
data_positions_batch=data_positions_batch,
|
||||
max_new_tokens=max_new_tokens,
|
||||
do_sample=False,
|
||||
eos_token_id=tok.eos_token_id,
|
||||
pad_token_id=tok.pad_token_id,
|
||||
)
|
||||
prompt_len = len(input_ids_batch[0])
|
||||
outputs = []
|
||||
for row in out:
|
||||
gen_ids = row.tolist()
|
||||
outputs.append(tok.decode(gen_ids[prompt_len:], skip_special_tokens=True))
|
||||
return outputs
|
||||
Reference in New Issue
Block a user