Added basic code for CETT threshold calculation and refactored activation capture.

kaselby · kaselby · commit 24c07b8e2c44 · 2025-07-10T13:19:56.000-04:00
Signed-off-by: Kira Selby &lt;kaselby@uwaterloo.ca&gt;
diff --git a/generate_dataset.py b/generate_dataset.py
@@ -43,7 +43,7 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from transformers.trainer_utils import set_seed
 
-from src.activation_capture import ActivationCaptureTraining
+from src.activation_capture import Hook
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
@@ -120,14 +120,14 @@ def process_batch(
     hidden_states_dict = {}
     mlp_activations_dict = {}
     for layer_idx in range(num_layers):
-        hidden_state = model.activation_capture.get_hidden_states(layer_idx)[0]
+        hidden_state = model.activation_capture.mlp_activations[Hook.IN][layer_idx][0]
         hidden_states_dict[layer_idx] = (
             hidden_state.view(-1, hidden_state.shape[-1])
             .cpu()
             .numpy()
             .astype(np.float32)
         )
-        mlp_activation = model.activation_capture.get_gate_activations(layer_idx)
+        mlp_activation = model.activation_capture.mlp_activations[Hook.ACT][layer_idx]
         mlp_activations_dict[layer_idx] = (
             mlp_activation[0]
             .view(-1, mlp_activation.shape[-1])
@@ -172,8 +172,8 @@ def generate_dataset(
         model = model.to(device)
 
     model.eval()
-    model.activation_capture = ActivationCaptureTraining(model)
-    model.activation_capture.register_hooks()
+    model.activation_capture = model.ACTIVATION_CAPTURE(model)
+    model.activation_capture.register_hooks(hooks=[Hook.IN, Hook.ACT])
 
     # Get model dimensions
     hidden_dim = model.config.hidden_size
diff --git a/measure_contextual_sparsity.py b/measure_contextual_sparsity.py
@@ -13,7 +13,7 @@
 from transformers.trainer_utils import set_seed
 
 import matplotlib.pyplot as plt
-from src.activation_capture import ActivationCaptureDefault
+from src.activation_capture import Hook
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
@@ -28,16 +28,14 @@ def __init__(self, model, tokenizer, device):
         self.tokenizer = tokenizer
         self.device = device
 
-        model.activation_capture = ActivationCaptureDefault(model)
-        model.activation_capture.register_hooks()
+        model.activation_capture = model.ACTIVATION_CAPTURE(model)
+        model.activation_capture.register_hooks(hooks=[Hook.ACT])
         self.num_layers = len(self.model.activation_capture.get_layers())
 
         self.reset_buffers()
 
     def reset_buffers(self):
-        self.mlp_sparsity = {}
-        self.mlp_sparsity["gate"] = defaultdict(list)
-        self.mlp_sparsity["up"] = defaultdict(list)
+        self.mlp_sparsity = defaultdict(list)
         self.num_seqs = 0
 
     def process_batch(self, input_ids: torch.Tensor, attention_mask: torch.Tensor):
@@ -54,26 +52,19 @@ def process_batch(self, input_ids: torch.Tensor, attention_mask: torch.Tensor):
 
         # Compute sparsity
         for layer_idx in range(self.num_layers):
-            sparsity_masks_gate = (
-                self.model.activation_capture.get_gate_activations(layer_idx) <= 0
-            )
-            sparsity_masks_up = (
-                self.model.activation_capture.get_up_activations(layer_idx) <= 0
+            sparsity_masks = (
+                self.model.activation_capture.mlp_activations[Hook.ACT][layer_idx] <= 0
             )
 
             # Naive sparsity computation
             self.mlp_sparsity["gate"][layer_idx].append(
-                sparsity_masks_gate.float().mean().item()
-            )
-            self.mlp_sparsity["up"][layer_idx].append(
-                sparsity_masks_up.float().mean().item()
+                sparsity_masks.float().mean().item()
             )
 
             # Level of sparsity after union over batch dim
             # union_sparsity_mask = sparsity_masks.any(dim=0)
             # self.union_sparsity[batch_size][layer_idx].append(union_sparsity_mask.float().mean().item())
 
-            # TODO: Add HNSW sparsity computation for both attn heads and mlp neurons
             # TODO: Compute union sparsity over multiple different batch sizes
 
         # Clear GPU tensors from capture to free memory
diff --git a/src/activation_capture.py b/src/activation_capture.py
@@ -1,54 +1,91 @@
-from typing_extensions import override
-import torch.nn.functional as F
-from abc import ABC, abstractmethod
 
+from enum import Enum
+from typing import List
 
-class ActivationCapture(ABC):
+class Hook(Enum):
+    IN = "IN"
+    ACT = "ACT"
+    UP = "UP"
+    OUT = "OUT"
+
+
+class ActivationCapture():
     """Helper class to capture activations from model layers."""
-    has_gate_proj: bool
-    has_up_proj: bool
+    hooks_available: List[Hook]
     
     def __init__(self, model):
         self.model = model
-        self.mlp_activations = {}
+        self.mlp_activations = {
+            hook: {} for hook in self.hooks_available
+        }
         self.handles = []
 
-    @abstractmethod
-    def _register_gate_hook(self, layer_idx, layer):
-        pass
+    def _register_in_hook(self, layer_idx, layer):
+        def hook(module, input, output):
+            # Just detach, don't clone or move to CPU yet
+            self.mlp_activations[Hook.IN][layer_idx] = input[0].clone().detach()
+            return output
+        handle = layer.mlp.register_forward_hook(hook)
+        return handle
+
+    def _register_act_hook(self, layer_idx, layer):
+        def hook(module, input, output):
+            # Just detach, don't clone or move to CPU yet
+            self.mlp_activations[Hook.ACT][layer_idx] = input[0].clone().detach()
+            return output
+        handle = layer.mlp.act_fn.register_forward_hook(hook)
+        return handle
 
-    @abstractmethod
     def _register_up_hook(self, layer_idx, layer):
-        pass
+        def hook(module, input, output):
+            # Just detach, don't clone or move to CPU yet
+            self.mlp_activations[Hook.UP][layer_idx] = input[0].clone().detach()
+            return output
+        handle = layer.mlp.down_proj.register_forward_hook(hook)
+        return handle
+    
+    def _register_out_hook(self, layer_idx, layer):
+        def hook(module, input, output):
+            # Just detach, don't clone or move to CPU yet
+            self.mlp_activations[Hook.OUT][layer_idx] = output.clone().detach()
+            return output
+        handle = layer.mlp.register_forward_hook(hook)
+        return handle
 
-    @abstractmethod
     def get_layers(self):
-        pass
-
-
-    @abstractmethod
-    def get_gate_activations(self, layer_idx):
-        """Get combined MLP activations for a layer."""
-        pass
+        return self.model.get_decoder().layers
 
-    def register_hooks(self):
+    def register_hooks(self, hooks=(Hook.ACT, Hook.UP, Hook.OUT)):
         """Register forward hooks to capture activations."""
         # Clear any existing hooks
         self.remove_hooks()
         
         # Hook into each transformer layer
-        for i, layer in enumerate(self.get_layers()):            
-            # Capture MLP gate activations (after activation function)
-            if self.has_gate_proj:
-                handle = self._register_gate_hook(i, layer)
+        for i, layer in enumerate(self.get_layers()):   
+            # Hooks capturing inputs to the MLP layer
+            if Hook.IN in hooks and Hook.IN in self.hooks_available:
+                handle = self._register_in_hook(i, layer)
                 if handle is not None:
                     self.handles.append(handle)
-                        
-            # Also capture up_proj activations
-            if self.has_up_proj:
+
+            # Hooks capturing inputs to the activation function      
+            if Hook.ACT in hooks and Hook.ACT in self.hooks_available:
+                handle = self._register_act_hook(i, layer)
+                if handle is not None:
+                    self.handles.append(handle)
+
+            # Hooks capturing inputs to the down projection
+            if Hook.UP in hooks and Hook.UP in self.hooks_available:
                 handle = self._register_up_hook(i, layer)
                 if handle is not None:
                     self.handles.append(handle)
+
+            # Hooks capturing the final MLP output
+            if Hook.OUT in hooks and Hook.OUT in self.hooks_available:
+                handle = self._register_out_hook(i, layer)
+                if handle is not None:
+                    self.handles.append(handle)
+
     
     def remove_hooks(self):
         """Remove all registered hooks."""
@@ -59,91 +96,3 @@ def remove_hooks(self):
     def clear_captures(self):
         """Clear captured activations."""
         self.mlp_activations = {}
-
-
-
-class ActivationCaptureDefault(ActivationCapture):
-    """Helper class to capture activations from model layers."""
-    has_gate_proj: bool = True
-    has_up_proj: bool = True
-
-    def get_layers(self):
-        return self.model.get_decoder().layers
-
-    def _create_mlp_hook(self, layer_idx, proj_type):
-        def hook(module, input, output):
-            key = f"{layer_idx}_{proj_type}"
-            # Just detach, don't clone or move to CPU yet
-            self.mlp_activations[key] = output.clone().detach()
-            return output
-        return hook
-
-    def _register_gate_hook(self, layer_idx, layer):
-        handle = layer.mlp.gate_proj.register_forward_hook(
-            self._create_mlp_hook(layer_idx, 'gate')
-        )
-        return handle
-
-    def _register_up_hook(self, layer_idx, layer):
-        handle = layer.mlp.up_proj.register_forward_hook(
-            self._create_mlp_hook(layer_idx, 'up')
-        )
-        return handle
-    
-    def get_gate_activations(self, layer_idx):
-        gate_key = f"{layer_idx}_gate"
-        if gate_key in self.mlp_activations:
-            gate_act = self.mlp_activations[gate_key]
-            return F.silu(gate_act)
-        return None
-
-    def get_up_activations(self, layer_idx):
-        up_key = f"{layer_idx}_up"
-        if up_key in self.mlp_activations:
-            up_act = self.mlp_activations[up_key]
-            return up_act
-        return None
-
-class ActivationCaptureTraining(ActivationCaptureDefault):
-    """Additional Hidden State capture for training dataset generation"""
-    def __init__(self, model):
-        super().__init__(model)
-        self.hidden_states = {}
-    
-    def _create_hidden_state_hook(self, layer_idx, layer):
-        def hook(module, args, kwargs, output):
-            # args[0] is the input hidden states to the layer
-            if len(args) > 0:
-                # Just detach, don't clone or move to CPU yet
-                self.hidden_states[layer_idx] = args[0].clone().detach()
-            return output
-        return hook
-    
-    def _register_hidden_state_hook(self, layer_idx, layer):
-        handle = layer.register_forward_hook(
-            self._create_hidden_state_hook(layer_idx, layer),
-            with_kwargs=True
-        )
-        return handle
-
-    @override
-    def clear_captures(self):
-        """Clear captured activations."""
-        super().clear_captures()
-        self.hidden_states = {}
-
-    @override
-    def register_hooks(self):
-        """Register forward hooks to capture activations."""
-        # Clear any existing hooks
-        super().register_hooks()
-        # Hook into each transformer layer
-        for i, layer in enumerate(self.get_layers()):            
-            # Capture hidden states before MLP
-            handle = self._register_hidden_state_hook(i, layer)
-            if handle is not None:
-                self.handles.append(handle)
-    
-    def get_hidden_states(self, layer_idx):
-        """Get hidden states for a layer."""
-        return self.hidden_states[layer_idx]
diff --git a/src/cett.py b/src/cett.py
@@ -0,0 +1,54 @@
+
+
+import torch
+
+from src.activation_capture import ActivationCapture, Hook
+
+def calculate_threshold_one_token(neuron_outputs, cett_target, n_quantiles=1000):
+    norms = neuron_outputs.norm(dim=0)
+    quantiles = norms.quantile(torch.linspace(0,1,n_quantiles))
+    tot_norm = neuron_outputs.sum(dim=1).norm()
+
+    def CETT(threshold):
+        threshold_norm = ((norms < threshold) * neuron_outputs).sum(dim=1).norm()
+        return threshold_norm / tot_norm
+
+    left = 0
+    right = quantiles.size(0)
+    threshold = 0
+    while left < right:
+        mid = (left + right) // 2
+        cett = CETT(quantiles[mid])
+        if cett <= cett_target:
+            left = mid + 1
+            threshold = quantiles[mid]
+        else:
+            right = mid - 1
+    return threshold
+
+
+def find_threshold(model, dataloader, layer_idx, cett_target=0.2, n_quantiles=500):
+    model.activation_capture = model.ACTIVATION_CAPTURE(model)
+    model.activation_capture.register_hooks(hooks=[Hook.UP])
+
+    thresholds = []
+    
+    with torch.no_grad():
+        for batch in dataloader:
+            input_ids = batch["input_ids"]
+            attention_mask = batch["attention_mask"]
+
+            model.activation_capture.clear_captures()
+        
+            _ = model(input_ids=input_ids, attention_mask=attention_mask)
+
+            activations = model.activation_capture.mlp_activations[Hook.UP][layer_idx]
+            activations = activations.view(-1, activations.size(-1))
+
+            for i in range(activations.size(0)):
+                neuron_outputs = activations[i] * model.model.layers[0].mlp.down_proj.weight
+                threshold = calculate_threshold_one_token(neuron_outputs, cett_target=cett_target, n_quantiles=n_quantiles)
+                thresholds.append(threshold)
+
+    return sum(thresholds)/len(thresholds)
+                
diff --git a/src/modeling_skip.py b/src/modeling_skip.py
@@ -20,7 +20,7 @@
 from transformers.utils.import_utils import is_torch_flex_attn_available
 
 from sparse_transformers import WeightCache, sparse_mlp_forward
-from src.activation_capture import ActivationCaptureDefault
+from src.activation_capture import ActivationCapture
 
 if is_torch_flex_attn_available():
     from torch.nn.attention.flex_attention import BlockMask
@@ -352,7 +352,7 @@ def forward(
 
 
 def build_skip_connection_model_for_causal_lm(pretrained_model_class: type[PreTrainedModel], base_model_class: type[PreTrainedModel]):
-    ACTIVATION_CAPTURE = ActivationCaptureDefault
+    ACTIVATION_CAPTURE = ActivationCapture
 
     class SkipConnectionModelForCausalLM(pretrained_model_class, GenerationMixin):
         _tied_weights_keys = ["lm_head.weight"]
diff --git a/src/models/phi3/modelling_phi_skip.py b/src/models/phi3/modelling_phi_skip.py
