56 lines
2.1 KiB
Python
56 lines
2.1 KiB
Python
class AttentionManager:
|
|
"""
|
|
Captures attention weights via hooks to avoid OOM.
|
|
Slices specific token positions immediately and discards the full matrix.
|
|
"""
|
|
def __init__(self, model):
|
|
self.model = model
|
|
self.attentions = {} # Store {layer_idx: attention_tensor}
|
|
self.hooks = []
|
|
self._register_hooks()
|
|
|
|
def _register_hooks(self):
|
|
# Locate the actual decoder layers.
|
|
# For Llama/Qwen + PEFT, it is usually model.base_model.model.layers or model.model.layers
|
|
if hasattr(self.model, "base_model"):
|
|
layers = self.model.base_model.model.layers
|
|
else:
|
|
layers = self.model.model.layers
|
|
|
|
for i, layer in enumerate(layers):
|
|
self.hooks.append(layer.register_forward_hook(self._make_hook(i)))
|
|
|
|
def _make_hook(self, idx):
|
|
def hook(module, args, output):
|
|
# output signature for LlamaDecoderLayer: (hidden_states, self_attn_weights, present_key_value)
|
|
# We want output[1] (self_attn_weights)
|
|
|
|
# Note: output is a tuple, so we must return a new tuple
|
|
if len(output) > 1 and output[1] is not None:
|
|
full_attn = output[1] # Shape: [bs, heads, seq_len, seq_len]
|
|
|
|
# --- CRITICAL OPTIMIZATION ---
|
|
# Slice ONLY the last token query, preserving gradients if needed.
|
|
# Shape becomes: [bs, heads, 1, seq_len]
|
|
# This is tiny compared to the full matrix.
|
|
print(f"hook len {len(output)}")
|
|
self.attentions[idx] = full_attn[..., -1, :]
|
|
|
|
# Replace the full attention in the output with None.
|
|
# This frees the GBs of memory immediately.
|
|
new_output = list(output)
|
|
new_output[1] = None
|
|
return tuple(new_output)
|
|
|
|
return output
|
|
return hook
|
|
|
|
def get_results(self):
|
|
return self.attentions
|
|
|
|
def clear(self):
|
|
self.attentions = {}
|
|
|
|
def remove_hooks(self):
|
|
for h in self.hooks:
|
|
h.remove() |