first commit

Codes/4_model_evaluation/lib/attention_hook_last_only.py

@@ -0,0 +1,56 @@
+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()