Add cortex/core.py — CortexModule, CortexSurgeon, CortexConfig

cortex/core.py

@@ -0,0 +1,493 @@
+"""
+Core abstractions for Cortex layer surgery.
+
+CortexModule: Protocol for all plug-in modules.
+CortexSurgeon: Orchestrator that injects/removes modules into a pretrained model.
+CortexConfig: Configuration for which modules go where.
+"""
+
+import torch
+import torch.nn as nn
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Tuple, Any, Callable, Union
+from enum import Enum
+import copy
+import logging
+
+logger = logging.getLogger("cortex")
+
+
+class InjectionPoint(Enum):
+    """Where in the transformer block a module can be injected."""
+    PRE_ATTENTION = "pre_attention"        # Before self-attention
+    POST_ATTENTION = "post_attention"      # After self-attention, before FFN
+    PRE_FFN = "pre_ffn"                    # Before FFN (same as post_attention for standard transformers)
+    POST_FFN = "post_ffn"                  # After FFN, after residual add
+    RESIDUAL_STREAM = "residual_stream"    # Direct access to the residual stream (wraps entire block)
+    FULL_BLOCK_WRAP = "full_block_wrap"    # Wraps the entire transformer block
+
+
+class CortexModule(ABC, nn.Module):
+    """
+    Base class for all Cortex plug-in modules.
+    
+    Each module must declare:
+      - injection_point: Where in the transformer block it operates
+      - target_layers: Which layers to inject into (list of ints or "all")
+    
+    And implement:
+      - forward(): Takes hidden states + optional kwargs, returns modified hidden states
+      - on_inject(): Called once when the module is injected (for initialization from activations)
+      - on_remove(): Called when the module is removed
+    """
+    
+    def __init__(self, injection_point: InjectionPoint, target_layers: Union[List[int], str] = "all"):
+        super().__init__()
+        self.injection_point = injection_point
+        self.target_layers = target_layers
+        self._original_modules: Dict[str, nn.Module] = {}
+        self._hooks: List[torch.utils.hooks.RemovableHook] = []
+        self._active = True
+    
+    @abstractmethod
+    def forward(
+        self, 
+        hidden_states: torch.Tensor,
+        layer_idx: int,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Any] = None,
+        **kwargs
+    ) -> torch.Tensor:
+        """
+        Process hidden states.
+        
+        Args:
+            hidden_states: [batch, seq_len, hidden_dim]
+            layer_idx: Which transformer layer this is being called from
+            attention_mask: Optional attention mask
+            position_ids: Optional position ids
+            past_key_values: Optional KV cache
+            **kwargs: Additional model-specific arguments
+            
+        Returns:
+            Modified hidden_states [batch, seq_len, hidden_dim]
+        """
+        ...
+    
+    def on_inject(self, model: nn.Module, layer_idx: int, hidden_dim: int):
+        """Called once when the module is injected. Use for calibration-based initialization."""
+        pass
+    
+    def on_remove(self, model: nn.Module, layer_idx: int):
+        """Called when the module is removed. Clean up hooks, restore originals."""
+        for hook in self._hooks:
+            hook.remove()
+        self._hooks.clear()
+    
+    def enable(self):
+        self._active = True
+    
+    def disable(self):
+        self._active = False
+    
+    @property
+    def is_active(self):
+        return self._active
+    
+    def extra_repr(self):
+        return f"injection={self.injection_point.value}, layers={self.target_layers}, active={self._active}"
+
+
+@dataclass
+class CortexConfig:
+    """Configuration for a Cortex surgery session."""
+    modules: Dict[str, Dict] = field(default_factory=dict)
+    # Each entry: {"module_class": "MemoryBank", "params": {...}, "target_layers": [0,1,2,...]}
+    freeze_base_model: bool = True
+    gradient_checkpointing: bool = False
+    
+    def add_module(self, name: str, module_class: str, target_layers: Union[List[int], str] = "all", **params):
+        self.modules[name] = {
+            "module_class": module_class,
+            "params": params,
+            "target_layers": target_layers,
+        }
+        return self
+
+
+class TransformerBlockWrapper(nn.Module):
+    """
+    Wraps a transformer block to intercept activations at specific injection points.
+    Multiple CortexModules can be attached at different injection points.
+    """
+    
+    def __init__(self, original_block: nn.Module, layer_idx: int, model_type: str = "llama"):
+        super().__init__()
+        self.original_block = original_block
+        self.layer_idx = layer_idx
+        self.model_type = model_type
+        
+        # Modules organized by injection point
+        self.pre_attention_modules: nn.ModuleList = nn.ModuleList()
+        self.post_attention_modules: nn.ModuleList = nn.ModuleList()
+        self.pre_ffn_modules: nn.ModuleList = nn.ModuleList()
+        self.post_ffn_modules: nn.ModuleList = nn.ModuleList()
+        self.residual_modules: nn.ModuleList = nn.ModuleList()
+    
+    def add_module_at(self, module: CortexModule):
+        """Add a CortexModule at its declared injection point."""
+        point = module.injection_point
+        if point == InjectionPoint.PRE_ATTENTION:
+            self.pre_attention_modules.append(module)
+        elif point == InjectionPoint.POST_ATTENTION:
+            self.post_attention_modules.append(module)
+        elif point == InjectionPoint.PRE_FFN:
+            self.pre_ffn_modules.append(module)
+        elif point == InjectionPoint.POST_FFN:
+            self.post_ffn_modules.append(module)
+        elif point in (InjectionPoint.RESIDUAL_STREAM, InjectionPoint.FULL_BLOCK_WRAP):
+            self.residual_modules.append(module)
+    
+    def forward(self, hidden_states, **kwargs):
+        """
+        Execute the transformer block with all injected modules.
+        
+        Supports LLaMA/Mistral/Qwen/Gemma-style blocks where:
+        - self_attn is the attention sublayer
+        - mlp is the FFN sublayer
+        - input_layernorm / post_attention_layernorm are the norms
+        """
+        # === RESIDUAL STREAM MODULES (operate on the whole block) ===
+        for mod in self.residual_modules:
+            if mod.is_active:
+                hidden_states = mod(hidden_states, layer_idx=self.layer_idx, **kwargs)
+        
+        # === PRE-ATTENTION ===
+        for mod in self.pre_attention_modules:
+            if mod.is_active:
+                hidden_states = mod(hidden_states, layer_idx=self.layer_idx, **kwargs)
+        
+        # === ORIGINAL BLOCK FORWARD ===
+        outputs = self.original_block(hidden_states, **kwargs)
+        
+        # outputs is typically (hidden_states, ...) — extract hidden states
+        if isinstance(outputs, tuple):
+            hidden_states_out = outputs[0]
+            rest = outputs[1:]
+        else:
+            hidden_states_out = outputs
+            rest = ()
+        
+        # === POST-FFN ===
+        for mod in self.post_ffn_modules:
+            if mod.is_active:
+                hidden_states_out = mod(hidden_states_out, layer_idx=self.layer_idx, **kwargs)
+        
+        if rest:
+            return (hidden_states_out,) + rest
+        return hidden_states_out
+
+
+class CortexSurgeon:
+    """
+    Main orchestrator for performing layer surgery on pretrained models.
+    
+    Usage:
+        model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
+        surgeon = CortexSurgeon(model)
+        
+        # Add modules
+        surgeon.add_module("memory", MemoryBank(hidden_dim=2048, num_slots=64), target_layers=[8,9,10,11])
+        surgeon.add_module("halluc_gate", HallucinationGate(hidden_dim=2048), target_layers="all")
+        
+        # Perform surgery
+        surgeon.operate()
+        
+        # Model now has enhanced capabilities
+        outputs = model(**inputs)
+    """
+    
+    def __init__(self, model: nn.Module, model_type: str = "auto"):
+        self.model = model
+        self.model_type = self._detect_model_type(model) if model_type == "auto" else model_type
+        self.modules: Dict[str, Tuple[CortexModule, Union[List[int], str]]] = {}
+        self._wrappers: Dict[int, TransformerBlockWrapper] = {}
+        self._hooks: List[torch.utils.hooks.RemovableHook] = []
+        self._operated = False
+        
+        # Detect model architecture
+        self.hidden_dim = self._get_hidden_dim()
+        self.num_layers = self._get_num_layers()
+        self.layers_accessor = self._get_layers_accessor()
+        
+        logger.info(f"CortexSurgeon initialized: model_type={self.model_type}, "
+                     f"hidden_dim={self.hidden_dim}, num_layers={self.num_layers}")
+    
+    def _detect_model_type(self, model: nn.Module) -> str:
+        """Auto-detect the model architecture."""
+        class_name = model.__class__.__name__.lower()
+        if "llama" in class_name:
+            return "llama"
+        elif "mistral" in class_name:
+            return "mistral"
+        elif "qwen" in class_name:
+            return "qwen2"
+        elif "gemma" in class_name:
+            return "gemma"
+        elif "gpt2" in class_name or "gptneo" in class_name:
+            return "gpt2"
+        elif "phi" in class_name:
+            return "phi"
+        else:
+            if hasattr(model, "config"):
+                mt = getattr(model.config, "model_type", "").lower()
+                if mt:
+                    return mt
+            return "llama"
+    
+    def _get_hidden_dim(self) -> int:
+        """Get the hidden dimension of the model."""
+        if hasattr(self.model, "config"):
+            return self.model.config.hidden_size
+        raise ValueError("Cannot detect hidden_dim — model has no config")
+    
+    def _get_num_layers(self) -> int:
+        """Get the number of transformer layers."""
+        if hasattr(self.model, "config"):
+            return self.model.config.num_hidden_layers
+        raise ValueError("Cannot detect num_layers — model has no config")
+    
+    def _get_layers_accessor(self) -> str:
+        """Get the attribute path to the list of transformer layers."""
+        for path in ["model.layers", "transformer.h", "gpt_neox.layers", "model.decoder.layers"]:
+            obj = self.model
+            try:
+                for attr in path.split("."):
+                    obj = getattr(obj, attr)
+                if hasattr(obj, "__len__") and len(obj) > 0:
+                    return path
+            except AttributeError:
+                continue
+        raise ValueError("Cannot find transformer layers in model")
+    
+    def _get_layers(self) -> nn.ModuleList:
+        """Access the transformer layers."""
+        obj = self.model
+        for attr in self.layers_accessor.split("."):
+            obj = getattr(obj, attr)
+        return obj
+    
+    def _resolve_target_layers(self, target: Union[List[int], str]) -> List[int]:
+        """Resolve 'all' to a list of layer indices, or pass through a list."""
+        if target == "all":
+            return list(range(self.num_layers))
+        elif isinstance(target, str) and target == "middle":
+            start = self.num_layers // 3
+            end = 2 * self.num_layers // 3
+            return list(range(start, end))
+        elif isinstance(target, str) and target == "deep":
+            start = 2 * self.num_layers // 3
+            return list(range(start, self.num_layers))
+        return list(target)
+    
+    def add_module(
+        self, 
+        name: str, 
+        module: CortexModule, 
+        target_layers: Optional[Union[List[int], str]] = None
+    ):
+        """
+        Register a CortexModule for injection.
+        
+        Args:
+            name: Unique name for this module instance
+            module: The CortexModule to inject
+            target_layers: Override the module's target_layers if provided
+        """
+        if target_layers is not None:
+            module.target_layers = target_layers
+        self.modules[name] = module
+        logger.info(f"Registered module '{name}': {module.__class__.__name__} "
+                     f"at {module.injection_point.value} for layers {module.target_layers}")
+    
+    def operate(self, freeze_base: bool = True):
+        """
+        Perform the surgery — inject all registered modules into the model.
+        
+        Args:
+            freeze_base: If True, freeze all original model parameters (only cortex modules train)
+        """
+        if self._operated:
+            raise RuntimeError("Surgery already performed. Create a new CortexSurgeon to modify further.")
+        
+        layers = self._get_layers()
+        
+        if freeze_base:
+            for param in self.model.parameters():
+                param.requires_grad = False
+            logger.info("Froze all base model parameters")
+        
+        for name, module in self.modules.items():
+            target_layer_idxs = self._resolve_target_layers(module.target_layers)
+            
+            for layer_idx in target_layer_idxs:
+                if layer_idx >= len(layers):
+                    logger.warning(f"Layer {layer_idx} out of range (model has {len(layers)} layers), skipping")
+                    continue
+                
+                layer = layers[layer_idx]
+                module.on_inject(self.model, layer_idx, self.hidden_dim)
+                hook = self._create_hook(module, layer, layer_idx)
+                if hook is not None:
+                    self._hooks.append(hook)
+            
+            device = next(self.model.parameters()).device
+            module.to(device)
+            
+            for param in module.parameters():
+                param.requires_grad = True
+            
+            logger.info(f"Injected '{name}' into layers {target_layer_idxs}")
+        
+        self._operated = True
+        
+        total_params = sum(p.numel() for p in self.model.parameters())
+        cortex_params = sum(
+            sum(p.numel() for p in m.parameters()) 
+            for m in self.modules.values()
+        )
+        logger.info(f"Surgery complete: {cortex_params:,} Cortex params / {total_params:,} total "
+                     f"({cortex_params/total_params*100:.2f}% overhead)")
+        
+        return self
+    
+    def _create_hook(self, module: CortexModule, layer: nn.Module, layer_idx: int):
+        """Create and register the appropriate hook for a module's injection point."""
+        point = module.injection_point
+        
+        if point == InjectionPoint.POST_FFN:
+            def post_ffn_hook(mod, inp, output, _module=module, _layer_idx=layer_idx):
+                if not _module.is_active:
+                    return output
+                if isinstance(output, tuple):
+                    hidden = output[0]
+                    hidden = _module(hidden, layer_idx=_layer_idx)
+                    return (hidden,) + output[1:]
+                return _module(output, layer_idx=_layer_idx)
+            return layer.register_forward_hook(post_ffn_hook)
+        
+        elif point == InjectionPoint.PRE_ATTENTION:
+            def pre_attn_hook(mod, inp, _module=module, _layer_idx=layer_idx):
+                if not _module.is_active:
+                    return inp
+                if isinstance(inp, tuple):
+                    hidden = inp[0]
+                    hidden = _module(hidden, layer_idx=_layer_idx)
+                    return (hidden,) + inp[1:]
+                return (_module(inp, layer_idx=_layer_idx),)
+            return layer.register_forward_pre_hook(pre_attn_hook)
+        
+        elif point == InjectionPoint.POST_ATTENTION:
+            attn = self._get_attention_sublayer(layer)
+            if attn is not None:
+                def post_attn_hook(mod, inp, output, _module=module, _layer_idx=layer_idx):
+                    if not _module.is_active:
+                        return output
+                    if isinstance(output, tuple):
+                        hidden = output[0]
+                        hidden = _module(hidden, layer_idx=_layer_idx)
+                        return (hidden,) + output[1:]
+                    return _module(output, layer_idx=_layer_idx)
+                return attn.register_forward_hook(post_attn_hook)
+        
+        elif point == InjectionPoint.PRE_FFN:
+            ffn = self._get_ffn_sublayer(layer)
+            if ffn is not None:
+                def pre_ffn_hook(mod, inp, _module=module, _layer_idx=layer_idx):
+                    if not _module.is_active:
+                        return inp
+                    if isinstance(inp, tuple):
+                        hidden = inp[0]
+                        hidden = _module(hidden, layer_idx=_layer_idx)
+                        return (hidden,) + inp[1:]
+                    return (_module(inp, layer_idx=_layer_idx),)
+                return ffn.register_forward_pre_hook(pre_ffn_hook)
+        
+        elif point == InjectionPoint.RESIDUAL_STREAM:
+            def residual_hook(mod, inp, output, _module=module, _layer_idx=layer_idx):
+                if not _module.is_active:
+                    return output
+                if isinstance(output, tuple):
+                    hidden = output[0]
+                    hidden = _module(hidden, layer_idx=_layer_idx)
+                    return (hidden,) + output[1:]
+                return _module(output, layer_idx=_layer_idx)
+            return layer.register_forward_hook(residual_hook)
+        
+        return None
+    
+    def _get_attention_sublayer(self, layer: nn.Module) -> Optional[nn.Module]:
+        """Get the attention sublayer of a transformer block."""
+        for name in ["self_attn", "attn", "attention"]:
+            if hasattr(layer, name):
+                return getattr(layer, name)
+        return None
+    
+    def _get_ffn_sublayer(self, layer: nn.Module) -> Optional[nn.Module]:
+        """Get the FFN sublayer of a transformer block."""
+        for name in ["mlp", "feed_forward", "ffn"]:
+            if hasattr(layer, name):
+                return getattr(layer, name)
+        return None
+    
+    def remove_all(self):
+        """Remove all injected modules and restore the original model."""
+        for hook in self._hooks:
+            hook.remove()
+        self._hooks.clear()
+        
+        for name, module in self.modules.items():
+            target_layers = self._resolve_target_layers(module.target_layers)
+            for layer_idx in target_layers:
+                module.on_remove(self.model, layer_idx)
+        
+        self.modules.clear()
+        self._operated = False
+        logger.info("All Cortex modules removed, model restored")
+    
+    def get_trainable_parameters(self):
+        """Yield only the Cortex module parameters (for optimizer)."""
+        for module in self.modules.values():
+            yield from module.parameters()
+    
+    def get_parameter_report(self) -> Dict:
+        """Get a detailed report of parameter counts per module."""
+        report = {}
+        for name, module in self.modules.items():
+            params = sum(p.numel() for p in module.parameters())
+            trainable = sum(p.numel() for p in module.parameters() if p.requires_grad)
+            report[name] = {"total": params, "trainable": trainable}
+        return report
+    
+    def save_cortex_modules(self, path: str):
+        """Save only the Cortex module weights."""
+        state = {}
+        for name, module in self.modules.items():
+            state[name] = {
+                "class": module.__class__.__name__,
+                "state_dict": module.state_dict(),
+                "injection_point": module.injection_point.value,
+                "target_layers": module.target_layers,
+            }
+        torch.save(state, path)
+        logger.info(f"Saved Cortex modules to {path}")
+    
+    def load_cortex_modules(self, path: str):
+        """Load Cortex module weights."""
+        state = torch.load(path, map_location="cpu")
+        for name, module in self.modules.items():
+            if name in state:
+                module.load_state_dict(state[name]["state_dict"])
+                logger.info(f"Loaded weights for '{name}'")