Activation Capture refactor (#46)

* Refactor activation capture and fix issues with generate_dataset.

Signed-off-by: Kira Selby <kaselby@uwaterloo.ca>

* Add simple script to plot sparsities

Signed-off-by: Kira Selby <kaselby@uwaterloo.ca>

* Activation capture code for phi3

Signed-off-by: Kira Selby <kaselby@uwaterloo.ca>

* Delete measure contextual sparsity and replace with updated version.

Signed-off-by: Kira Selby <kaselby@uwaterloo.ca>

---------

Signed-off-by: Kira Selby <kaselby@uwaterloo.ca>

Author: kaselby

generate_dataset.py

@@ -49,10 +49,10 @@
 from datasets import load_dataset
 from torch.utils.data import DataLoader as TorchDataLoader
 from tqdm import tqdm
-from src.activation_capture import ActivationCapture
+from src.activation_capture import ACTIVATION_CAPTURE, ActivationCapture
 import csv
 import glob
-from src.predictor_trainer import get_sample_by_index
+from src.trainer import get_sample_by_index
 # Setup logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -149,7 +149,7 @@ def process_batch(
                 hidden_states_dict[layer_idx].append(hidden_state)
 
                 # Get last token's MLP activations
-                mlp_activation = capture.get_mlp_activations(layer_idx)
+                mlp_activation = capture.get_gate_activations(layer_idx)
                 if mlp_activation is not None:
                     mlp_act = mlp_activation[batch_idx,-1,:].cpu().numpy().astype(np.float32)
                     mlp_activations_dict[layer_idx].append(mlp_act)
@@ -196,14 +196,15 @@ def generate_dataset(
 
     model.eval()
 
+    # Setup activation capture
+    capture_cls = ACTIVATION_CAPTURE[model.config.model_type]
+    capture = capture_cls()
+    capture.register_hooks(model)
+
     # Get model dimensions
     hidden_dim = model.config.hidden_size
     intermediate_dim = model.config.intermediate_size
-    num_layers = len(model.model.layers)
-    
-    # Setup activation capture
-    capture = ActivationCapture()
-    capture.register_hooks(model)
+    num_layers = len(capture.get_layers(model))
 
     # Load dataset
     logger.info(f"Loading dataset: {dataset_name}")

measure_contextual_sparsity.py

@@ -1,60 +1,67 @@
 import torch.nn.functional as F
+from abc import ABC, abstractmethod
 
 
-class ActivationCapture:
+class ActivationCapture(ABC):
     """Helper class to capture activations from model layers."""
+    has_gate_proj: bool
+    has_up_proj: bool
 
     def __init__(self):
         self.hidden_states = {}
         self.mlp_activations = {}
         self.handles = []
-        
+
+    @abstractmethod
+    def _register_gate_hook(self, layer_idx, layer):
+        pass
+
+    @abstractmethod
+    def _register_up_hook(self, layer_idx, layer):
+        pass
+
+    @abstractmethod
+    def get_layers(self, model):
+        pass
+
+    def _register_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].detach()
+            return output
+        handle = layer.register_forward_hook(
+            hook,
+            with_kwargs=True
+        )
+        return handle
+
     def register_hooks(self, model):
         """Register forward hooks to capture activations."""
         # Clear any existing hooks
         self.remove_hooks()
 
         # Hook into each transformer layer
-        for i, layer in enumerate(model.model.layers):
+        for i, layer in enumerate(self.get_layers(model)):
 
             # Capture hidden states before MLP
-            handle = layer.register_forward_hook(
-                self._create_hidden_state_hook(i),
-                with_kwargs=True
-            )
-            self.handles.append(handle)
+            handle = self._register_hidden_state_hook(i, layer)
+            if handle is not None:
+                self.handles.append(handle)
 
             # Capture MLP gate activations (after activation function)
-            if hasattr(layer.mlp, 'gate_proj'):
-                handle = layer.mlp.gate_proj.register_forward_hook(
-                    self._create_mlp_hook(i, 'gate')
-                )
-                self.handles.append(handle)
+            if self.has_gate_proj:
+                handle = self._register_gate_hook(i, layer)
+                if handle is not None:
+                    self.handles.append(handle)
 
             # Also capture up_proj activations
-            if hasattr(layer.mlp, 'up_proj'):
-                handle = layer.mlp.up_proj.register_forward_hook(
-                    self._create_mlp_hook(i, 'up')
-                )
-                self.handles.append(handle)
-    
-    def _create_hidden_state_hook(self, layer_idx):
-        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].detach()
-            return output
-        return hook
-    
-    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.detach()
-            return output
-        return hook
-    
+            if self.has_up_proj:
+                handle = self._register_up_hook(i, layer)
+                if handle is not None:
+                    self.handles.append(handle)
+
     def remove_hooks(self):
         """Remove all registered hooks."""
         for handle in self.handles:
@@ -65,7 +72,46 @@ def clear_captures(self):
         """Clear captured activations."""
         self.hidden_states = {}
         self.mlp_activations = {}
-    
+
+    @abstractmethod
+    def get_mlp_activations(self, layer_idx):
+        """Get combined MLP activations for a layer."""
+        pass
+
+    @abstractmethod
+    def get_gate_activations(self, layer_idx):
+        """Get combined MLP activations for a layer."""
+        return 
+
+
+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, model):
+        return model.model.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.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_mlp_activations(self, layer_idx):
         """Get combined MLP activations for a layer."""
         gate_key = f"{layer_idx}_gate"
@@ -80,4 +126,18 @@ def get_mlp_activations(self, layer_idx):
             gated_act = F.silu(gate_act) * up_act
             return gated_act
 
-        return None
+        return None
+    
+    def get_gate_activations(self, layer_idx):
+        """Get combined MLP activations for a layer."""
+        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
+    
+
+ACTIVATION_CAPTURE = {}
+
+def register_activation_capture(model_name, activation_capture):
+    ACTIVATION_CAPTURE[model_name] = activation_capture

src/activation_capture.py

@@ -7,4 +7,7 @@
 AutoConfig.register("llama-skip", LlamaSkipConnectionConfig)
 AutoModelForCausalLM.register(LlamaSkipConnectionConfig, LlamaSkipConnectionForCausalLM)
 
+from src.activation_capture import register_activation_capture, ActivationCaptureDefault
+register_activation_capture('llama', ActivationCaptureDefault)
+
 __all__ = [configuration_llama_skip, modelling_llama_skip]

src/models/llama/__init__.py

@@ -7,4 +7,7 @@
 AutoConfig.register("mistral-skip", MistralSkipConnectionConfig)
 AutoModelForCausalLM.register(MistralSkipConnectionConfig, MistralSkipConnectionForCausalLM)
 
+from src.activation_capture import register_activation_capture, ActivationCaptureDefault
+register_activation_capture('mistral', ActivationCaptureDefault)
+
 __all__ = [configuration_mistral_skip, modelling_mistral_skip]

src/models/mistral/__init__.py

@@ -7,4 +7,8 @@
 AutoConfig.register("phi3-skip", Phi3SkipConnectionConfig)
 AutoModelForCausalLM.register(Phi3SkipConnectionConfig, Phi3SkipConnectionForCausalLM)
 
+from .activation_capture_phi import ActivationCapturePhi3
+from src.activation_capture import register_activation_capture
+register_activation_capture('phi3', ActivationCapturePhi3)
+
 __all__ = [configuration_phi_skip, modelling_phi_skip]

src/models/phi3/__init__.py

@@ -0,0 +1,58 @@
+from src.activation_capture import ActivationCapture
+import torch.nn.functional as F
+
+
+
+class ActivationCapturePhi3(ActivationCapture):
+    """Helper class to capture activations from model layers."""
+    has_gate_proj: bool = True
+    has_up_proj: bool = True
+
+    def get_layers(self, model):
+        return model.model.layers
+
+    def _register_gate_hook(self, layer_idx, layer):
+        def hook(module, input, output):
+            key1 = f"{layer_idx}_{'gate'}"
+            key2 = f"{layer_idx}_{'up'}"
+            # Just detach, don't clone or move to CPU yet
+            gate_outputs, up_outputs = output.chunk(2, dim=1)
+            self.mlp_activations[key1] = gate_outputs.detach()
+            self.mlp_activations[key2] = up_outputs.detach()
+            return output
+        handle = layer.mlp.gate_up_proj.register_forward_hook(hook)
+        return handle
+
+    def _register_up_hook(self, layer_idx, layer):
+        def hook(module, input, output):
+            key = f"{layer_idx}_{'up'}"
+            # Just detach, don't clone or move to CPU yet
+            up_outputs = output.chunk(2, dim=1)[1]
+            self.mlp_activations[key] = up_outputs.detach()
+            return output
+        handle = layer.mlp.gate_up_proj.register_forward_hook(hook)
+        return handle
+    
+    def get_gate_activations(self, layer_idx):
+        """Get combined MLP activations for a layer."""
+        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_mlp_activations(self, layer_idx):
+        """Get combined MLP activations for a layer."""
+        gate_key = f"{layer_idx}_gate"
+        up_key = f"{layer_idx}_up"
+        
+        if gate_key in self.mlp_activations and up_key in self.mlp_activations:
+            # Compute gated activations: gate(x) * up(x)
+            gate_act = self.mlp_activations[gate_key]
+            up_act = self.mlp_activations[up_key]
+            
+            # Apply SwiGLU activation: silu(gate) * up
+            gated_act = F.silu(gate_act) * up_act
+            return gated_act
+        
+        return None

src/models/phi3/activation_capture_phi.py

@@ -7,4 +7,7 @@
 AutoConfig.register("qwen2-skip", Qwen2SkipConnectionConfig)
 AutoModelForCausalLM.register(Qwen2SkipConnectionConfig, Qwen2SkipConnectionForCausalLM)
 
+from src.activation_capture import register_activation_capture, ActivationCaptureDefault
+register_activation_capture('qwen2', ActivationCaptureDefault)
+
 __all__ = [configuration_qwen_skip, modelling_qwen_skip]