Fixes to activation capture, gemma and CETT

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

Author: kaselby

src/activation_capture.py

@@ -11,7 +11,7 @@ class Hook(Enum):
 
 class ActivationCapture():
     """Helper class to capture activations from model layers."""
-    hooks_available: List[Hook]
+    hooks_available: List[Hook] = [Hook.IN, Hook.ACT, Hook.UP, Hook.OUT]
 
     def __init__(self, model):
         self.model = model
@@ -95,4 +95,6 @@ def remove_hooks(self):
 
     def clear_captures(self):
         """Clear captured activations."""
-        self.mlp_activations = {}
+        self.mlp_activations = {
+            hook: {} for hook in self.hooks_available
+        }

src/cett.py

@@ -2,7 +2,7 @@
 import logging
 import os
 import json
-import tqdm
+from tqdm import tqdm
 import argparse
 
 from datasets import load_dataset
@@ -39,46 +39,16 @@ def CETT(threshold):
     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 = defaultdict(list)
-    
-    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)
-
-            for layer,layer_idx in enumerate(model.activation_capture.get_layers()):
-                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] * layer.mlp.down_proj.weight
-                    threshold = calculate_threshold_one_token(neuron_outputs, cett_target=cett_target, n_quantiles=n_quantiles)
-                    thresholds[layer_idx].append(threshold)
-
-    for layer_idx, layer_thresholds in thresholds.items():
-        thresholds[layer_idx] = sum(layer_thresholds) / len(layer_thresholds)
-
-    return thresholds
-
-
-
 def find_thresholds(
-        model_name, 
-        dataset_name, 
-        dataset_config,
-        max_samples, 
-        cett_target, 
-        n_quantiles,
-        save_path,
-        device,
+        model_name: str, 
+        dataset_name: str, 
+        dataset_config: str,
+        max_samples: int, 
+        cett_target: float, 
+        n_quantiles: int,
+        save_path: str,
+        seed: int,
+        device: torch.device,
     ):
 
     # Load tokenizer and model
@@ -96,7 +66,7 @@ def find_thresholds(
         model = model.to(device)
 
     model.eval()
-    model.activation_capture = model.ACTIVATION_CAPTURE(model)
+    model.activation_capture = ActivationCapture(model)
     model.activation_capture.register_hooks(hooks=[Hook.UP])
 
     # Load dataset
@@ -107,7 +77,7 @@ def find_thresholds(
         )
     else:
         dataset = load_dataset(dataset_name, split="train", streaming=True)
-    dataset = dataset.shuffle(buffer_size=10000, seed=42)
+    dataset = dataset.shuffle(buffer_size=10000, seed=seed)
 
     def sample_and_tokenize(examples):
         """Sample text chunks before tokenization for efficiency using vectorized operations."""
@@ -130,13 +100,12 @@ def sample_and_tokenize(examples):
     logger.info(f"Beginning to compute thresholds using {max_samples} samples")
     thresholds = defaultdict(list)
     with torch.no_grad():
-        for batch in tqdm.tqdm(dataloader):
+        for batch in tqdm(dataloader, total=max_samples):
             input_ids = batch["input_ids"].to(device)
-            attention_mask = batch["attention_mask"].to(device)
 
-            _ = model(input_ids=input_ids, attention_mask=attention_mask)
+            _ = model(input_ids.squeeze(0))
 
-            for layer,layer_idx in enumerate(model.activation_capture.get_layers()):
+            for layer_idx, layer in enumerate(model.activation_capture.get_layers()):
                 activations = model.activation_capture.mlp_activations[Hook.UP][layer_idx]
                 activations = activations.view(-1, activations.size(-1))
 
@@ -232,12 +201,13 @@ def parse_args():
 
     find_thresholds(
         model_name=args.model_name, 
-        dataset_name=args.dataset_name, 
+        dataset_name=args.dataset, 
         dataset_config=args.dataset_config,
         max_samples=args.max_samples, 
         cett_target=args.cett_target, 
         n_quantiles=args.n_quantiles,
         save_path=args.save_path,
-        device=device
+        seed=args.seed,
+        device=device,
     )
 

src/models/gemma3n/activation_capture_gemma.py

@@ -27,7 +27,6 @@
 from sparse_transformers import sparse_mlp_forward
 
 from src.models.gemma3n.configuration_gemma_skip import Gemma3nSkipConnectionConfig
-from src.models.gemma3n.activation_capture_gemma import ActivationCaptureGemma3n
 from src.modeling_skip import SkipMLP, SkipDecoderLayer, build_skip_connection_model, build_skip_connection_model_for_causal_lm
 
 logger = logging.get_logger(__name__)
@@ -413,7 +412,6 @@ def project_per_layer_inputs(
 Gemma3nSkipConnectionForCausalLMBase = build_skip_connection_model_for_causal_lm(Gemma3nSkipPreTrainedModel, Gemma3nSkipConnectionModel)
 
 class Gemma3nSkipConnectionForCausalLM(Gemma3nSkipConnectionForCausalLMBase):
-    ACTIVATION_CAPTURE = ActivationCaptureGemma3n
     _tied_weights_keys = ["lm_head.weight"]
     _tp_plan = {"lm_head": "colwise_rep"}
     _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}

src/models/gemma3n/modelling_gemma_skip.py

@@ -25,7 +25,6 @@
 
 from src.models.phi3.configuration_phi_skip import Phi3SkipConnectionConfig
 from src.modeling_skip import SkipMLP, SkipDecoderLayer, FastLoRAProjection, build_skip_connection_model, build_skip_connection_model_for_causal_lm
-from .activation_capture_phi import ActivationCapturePhi3
 logger = logging.get_logger(__name__)
 
 
@@ -339,8 +338,6 @@ def _prepare_4d_causal_attention_mask_with_cache_position(
 Phi3SkipConnectionForCausalLMBase = build_skip_connection_model_for_causal_lm(Phi3SkipPreTrainedModel, Phi3SkipConnectionModel)
 
 class Phi3SkipConnectionForCausalLM(Phi3SkipConnectionForCausalLMBase):
-    ACTIVATION_CAPTURE = ActivationCapturePhi3
-
     _keys_to_ignore_on_load_missing = [
         "model.layers.*.mlp.combined_proj_buffer",
         "model.layers.*.mlp.down_proj_buffer",