""" Projector nn module for the unified memory """ import torch import torch.nn as nn from torch import Tensor from transformers import Cache, DynamicCache from typing import Optional, Tuple, Literal, Union import copy import math from rosetta.utils.registry import register_model, get_projector_class, PROJECTOR_REGISTRY, capture_init_args, save_object, load_object class Projector(nn.Module): """Base projector class for unified memory""" def forward(self, source_kv: Tuple[Tensor, Tensor], target_kv: Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tensor]: """ Project and combine the source key-value tensors to the target key-value tensors Args: source_kv: Tuple of (key, value) tensors, each (..., D_s) where ... are arbitrary leading dimensions target_kv: Tuple of (key, value) tensors, each (..., D_t) where ... are arbitrary leading dimensions Returns: Tuple of (key, value) tensors, each (..., D_t) with same leading dimensions as input """ raise NotImplementedError("Subclasses must implement forward method") def cache_project(self, source_kv_cache: Cache, target_kv_cache: Cache) -> Cache: """ Project the source kv cache to the target kv cache """ if not isinstance(source_kv_cache, DynamicCache) or not isinstance(target_kv_cache, DynamicCache): raise ValueError("Only DynamicCache is supported") projected_cache = DynamicCache() # Process each layer for layer_idx in range(len(source_kv_cache.key_cache)): source_key = source_kv_cache.key_cache[layer_idx] # (B, H, N, D_s) source_value = source_kv_cache.value_cache[layer_idx] # (B, H, N, D_s) # Get corresponding target tensors (for reference/combination) if layer_idx < len(target_kv_cache.key_cache): target_key = target_kv_cache.key_cache[layer_idx] # (B, H, N, D_t) target_value = target_kv_cache.value_cache[layer_idx] # (B, H, N, D_t) else: # If target cache doesn't have this layer, create dummy tensors B, H, N, D_s = source_key.shape D_t = source_key.shape[-1] # Assume same dimension for simplicity target_key = torch.zeros(B, H, N, D_t, device=source_key.device, dtype=source_key.dtype) target_value = torch.zeros(B, H, N, D_t, device=source_value.device, dtype=source_value.dtype) # Reshape for forward pass: DynamicCache format (B, H, N, D) -> projector format (B, N, H, D) source_key_reshaped = source_key.transpose(1, 2) source_value_reshaped = source_value.transpose(1, 2) target_key_reshaped = target_key.transpose(1, 2) target_value_reshaped = target_value.transpose(1, 2) # Project using forward method with tuple input/output source_kv = (source_key_reshaped, source_value_reshaped) target_kv = (target_key_reshaped, target_value_reshaped) projected_key, projected_value = self.forward(source_kv, target_kv) # Reshape back: projector format (B, N, H, D) -> DynamicCache format (B, H, N, D) projected_key = projected_key.transpose(1, 2) projected_value = projected_value.transpose(1, 2) # Update cache projected_cache.update(projected_key, projected_value, layer_idx) return projected_cache @register_model @capture_init_args class TrivialProjector(Projector): """ Trivial projector that directly outputs the target key-value pairs without any modification. This is useful as a baseline or when you want to effectively disable projection. """ def __init__(self, **kwargs): """ Initialize the trivial projector. Args: source_dim: Source dimension (ignored, kept for compatibility) target_dim: Target dimension (ignored, kept for compatibility) **kwargs: Additional arguments (ignored, kept for compatibility) """ super().__init__() def forward(self, source_kv: Tuple[Tensor, Tensor], target_kv: Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tensor]: """ Return the target key-value pairs unchanged, ignoring the source. Args: source_kv: Tuple of (key, value) tensors (ignored) target_kv: Tuple of (key, value) tensors to return unchanged Returns: The target key-value pairs unchanged """ return target_kv @register_model @capture_init_args class ReplaceProjector(Projector): """ Replacement projector that projects source key-value tensors to target dimension using MLP, then replace target tensors using learnable weights. """ def __init__( self, source_dim: int, target_dim: int, hidden_dim: int = 512, num_layers: int = 2, dropout: float = 0.1, activation: str = "gelu", use_layer_norm: bool = True, init_weight: float = 0.1, anneal_steps: int = 1360, initial_temperature: float = 1.0, final_temperature: float = 0.01, scalar_temperature: float = 0.005, # shared_key_projection: nn.Module = None, # shared_value_projection: nn.Module = None, dtype: torch.dtype = torch.float32 ): super().__init__() self.source_dim = source_dim self.target_dim = target_dim self.hidden_dim = hidden_dim self.num_layers = num_layers # Activation function if activation.lower() == "gelu": self.activation = nn.GELU() elif activation.lower() == "relu": self.activation = nn.ReLU() elif activation.lower() == "silu": self.activation = nn.SiLU() else: raise ValueError(f"Unsupported activation: {activation}") # Build separate MLP layers for key and value projection self.key_projection = self._build_mlp(source_dim, hidden_dim, target_dim, num_layers, use_layer_norm, dropout, dtype) self.value_projection = self._build_mlp(source_dim, hidden_dim, target_dim, num_layers, use_layer_norm, dropout, dtype) def _build_mlp(self, source_dim: int, hidden_dim: int, target_dim: int, num_layers: int, use_layer_norm: bool, dropout: float, dtype: torch.dtype) -> nn.Sequential: """Build a single MLP projection module""" layers = [] # Input projection layers.append(nn.Linear(source_dim, hidden_dim, dtype=dtype)) if use_layer_norm: layers.append(nn.LayerNorm(hidden_dim, dtype=dtype)) layers.append(copy.deepcopy(self.activation)) layers.append(nn.Dropout(dropout)) # Hidden layers for _ in range(num_layers - 2): layers.append(nn.Linear(hidden_dim, hidden_dim, dtype=dtype)) if use_layer_norm: layers.append(nn.LayerNorm(hidden_dim, dtype=dtype)) layers.append(copy.deepcopy(self.activation)) layers.append(nn.Dropout(dropout)) # Output projection if num_layers > 1: layers.append(nn.Linear(hidden_dim, target_dim, dtype=dtype)) else: # Single layer case layers = [nn.Linear(source_dim, target_dim, dtype=dtype)] return nn.Sequential(*layers) def forward(self, source_kv: Tuple[Tensor, Tensor], target_kv: Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tensor]: """ Project source key-value tensors to target dimension and add to target tensors with learnable weights Args: source_kv: Tuple of (key, value) tensors, each (..., D_s) where ... are arbitrary leading dimensions target_kv: Tuple of (key, value) tensors, each (..., D_t) where ... are arbitrary leading dimensions Returns: Tuple of (key, value) tensors, each (..., D_t) with same leading dimensions as input """ source_key, source_value = source_kv target_key, target_value = target_kv # Get shapes - assuming format is (B, H, N, D) where H is num_heads, N is seq_len, D is head_dim source_shape = source_key.shape # (B, H_s, N, D_s) target_shape = target_key.shape # (B, H_t, N, D_t) # Extract dimensions batch_size, source_num_heads, seq_len, source_head_dim = source_shape _, target_num_heads, _, target_head_dim = target_shape # Reshape source: merge num_heads and head_dim for projection # (B, H_s, N, D_s) -> (B, N, H_s * D_s) source_key_reshaped = source_key.transpose(1, 2) # (B, N, H_s, D_s) source_value_reshaped = source_value.transpose(1, 2) # (B, N, H_s, D_s) source_key_flat = source_key_reshaped.contiguous().view(batch_size, seq_len, source_num_heads * source_head_dim) source_value_flat = source_value_reshaped.contiguous().view(batch_size, seq_len, source_num_heads * source_head_dim) # Project source tensors from (H_s * D_s) to (H_t * D_t) projected_key_flat = self.key_projection(source_key_flat) # (B, N, H_t * D_t) projected_value_flat = self.value_projection(source_value_flat) # (B, N, H_t * D_t) # Reshape projected tensors back to target format # (B, N, H_t * D_t) -> (B, N, H_t, D_t) -> (B, H_t, N, D_t) projected_key_reshaped = projected_key_flat.view(batch_size, seq_len, target_num_heads, target_head_dim) projected_value_reshaped = projected_value_flat.view(batch_size, seq_len, target_num_heads, target_head_dim) projected_key = projected_key_reshaped.transpose(1, 2) # (B, H_t, N, D_t) projected_value = projected_value_reshaped.transpose(1, 2) # (B, H_t, N, D_t) return (projected_key, projected_value) class ModernMLP(nn.Module): """ Modern MLP with residual connections, layer normalization, and configurable architecture. """ def __init__( self, input_dim: int, output_dim: int, hidden_dim: int = 512, num_layers: int = 2, activation: str = "gelu", use_layer_norm: bool = True, use_residual: bool = True, dropout: float = 0.1, use_swiglu: bool = False, dtype: torch.dtype = torch.float32 ): super().__init__() self.input_dim = input_dim self.output_dim = output_dim self.hidden_dim = hidden_dim self.num_layers = num_layers self.use_residual = use_residual and (input_dim == output_dim) self.use_swiglu = use_swiglu # Activation function if activation.lower() == "gelu": self.activation = nn.GELU() elif activation.lower() == "relu": self.activation = nn.ReLU() elif activation.lower() == "silu": self.activation = nn.SiLU() else: raise ValueError(f"Unsupported activation: {activation}") # Build layers self.layers = nn.ModuleList() for i in range(num_layers): layer_input_dim = input_dim if i == 0 else hidden_dim layer_output_dim = output_dim if i == num_layers - 1 else hidden_dim if self.use_swiglu and i < num_layers - 1: # Don't use SwiGLU on output layer layer = SwiGLUBlock(layer_input_dim, layer_output_dim, dtype=dtype) else: layer = nn.Linear(layer_input_dim, layer_output_dim, dtype=dtype) self.layers.append(layer) # Add layer norm after each layer except the last one if use_layer_norm and i < num_layers - 1: self.layers.append(nn.LayerNorm(layer_output_dim, dtype=dtype)) # Add activation after each layer except the last one if i < num_layers - 1 and not self.use_swiglu: self.layers.append(copy.deepcopy(self.activation)) # Add dropout after activation if dropout > 0 and i < num_layers - 1: self.layers.append(nn.Dropout(dropout)) # Residual projection if dimensions don't match if self.use_residual and input_dim != output_dim: self.residual_proj = nn.Linear(input_dim, output_dim, dtype=dtype) else: self.residual_proj = None def forward(self, x: Tensor) -> Tensor: """Forward pass with optional residual connection.""" residual = x for layer in self.layers: x = layer(x) # Add residual connection if self.use_residual: if self.residual_proj is not None: residual = self.residual_proj(residual) x = x + residual return x class SwiGLUBlock(nn.Module): """SwiGLU activation block for modern transformer architectures.""" def __init__(self, input_dim: int, output_dim: int, dtype: torch.dtype = torch.float32): super().__init__() self.gate_proj = nn.Linear(input_dim, output_dim, dtype=dtype) self.up_proj = nn.Linear(input_dim, output_dim, dtype=dtype) self.activation = nn.SiLU() def forward(self, x: Tensor) -> Tensor: gate = self.activation(self.gate_proj(x)) up = self.up_proj(x) return gate * up @register_model @capture_init_args class AllInOneProjector(Projector): """ Unified projector that consolidates all projection functionalities with modern patterns. Features: 1. Gate logit granularity: scalar, token-wise, head-wise, or value-wise 2. (DISABLED) Selector logit granularity: scalar, token-wise, head-wise, or value-wise 3. Key/Value weight granularity: scalar, token-wise, head-wise, or value-wise 4. Input-dependent gates and weights via MLP or parameters 5. Optional concatenation with combiner networks 6. Modern MLP architecture with residual connections and SwiGLU 7. Configurable target preservation: choose between traditional blending or simplified projection 8. Optional adding of target (self) signal to outputs via add_self Target Preservation Modes: - preserve_target_weight=True (default): output = (1-weight)*target + gate*selector*weight*projected - preserve_target_weight=False: output = target + gate*selector*weight*projected (no weight coefficient on target) Note: Selector functionality has been disabled/commented out. """ def __init__( self, source_dim: int, target_dim: int, source_num_heads: int = 1, target_num_heads: int = 1, hidden_dim: int = 512, num_layers: int = 2, dropout: float = 0.1, activation: str = "gelu", use_layer_norm: bool = True, use_residual: bool = True, use_swiglu: bool = False, # Gate configuration gate_granularity: Literal["scalar", "token", "head", "head_merged", "value"] = "scalar", gate_depends_on_input: bool = False, gate_input_features: Optional[str] = "target_key", # "target_key", "target_value", "both", "target_projected_key", "target_projected_value", "target_projected_both" gate_init_value: float = 0.0, # Weight configuration weight_granularity: Literal["scalar", "token", "head", "head_merged", "value"] = "scalar", weight_depends_on_input: bool = False, weight_input_features: Optional[str] = "target_key", # "target_key", "target_value", "both", "target_projected_key", "target_projected_value", "target_projected_both" weight_init_value: float = 0.0, # Target preservation configuration preserve_target_weight: bool = True, # If False, target won't be multiplied by (1 - normalized_weight) add_self: bool = True, # If False, target (self) won't be added to outputs # Concat configuration use_concat: bool = False, # combiner_hidden_dim: int = 128, weight_hidden_dim: int = 1024, # Temperature and gumbel use_gumbel: bool = True, initial_temperature: float = 1.0, final_temperature: float = 0.01, anneal_steps: int = 1360, scalar_temperature: float = 0.005, # Sequence length configuration max_sequence_length: int = 8192, # Maximum sequence length for token-level parameters pos_emb: bool = False, dtype: torch.dtype = torch.float32 ): super().__init__() self.source_dim = source_dim self.target_dim = target_dim self.source_num_heads = source_num_heads self.target_num_heads = target_num_heads self.hidden_dim = hidden_dim self.weight_hidden_dim = weight_hidden_dim self.max_sequence_length = max_sequence_length # Configuration self.gate_granularity = gate_granularity self.gate_depends_on_input = gate_depends_on_input self.gate_input_features = gate_input_features self.weight_granularity = weight_granularity self.weight_depends_on_input = weight_depends_on_input self.weight_input_features = weight_input_features self.preserve_target_weight = preserve_target_weight self.add_self = add_self self.use_concat = use_concat self.use_gumbel = use_gumbel self.scalar_temperature = scalar_temperature # Temperature annealing for gate only (removed selector temperature) self.register_buffer("gate_temperature", torch.tensor(initial_temperature, dtype=dtype)) self.initial_temperature = initial_temperature self.final_temperature = final_temperature self.anneal_steps = anneal_steps # Build projection networks self.key_projection = self._build_projection_mlp( source_dim * source_num_heads, target_dim * target_num_heads, hidden_dim, num_layers, activation, use_layer_norm, use_residual, dropout, use_swiglu, dtype ) self.value_projection = self._build_projection_mlp( source_dim * source_num_heads, target_dim * target_num_heads, hidden_dim, num_layers, activation, use_layer_norm, use_residual, dropout, use_swiglu, dtype ) # Build gate components self._build_gate_components(dtype) # Build weight components self._build_weight_components(weight_init_value, dtype) # Build concat components if needed if self.use_concat: in_dim = target_dim * target_num_heads * 2 out_dim = target_dim * target_num_heads self.key_combiner = nn.Linear(in_dim, out_dim, dtype=dtype) self.value_combiner = nn.Linear(in_dim, out_dim, dtype=dtype) def _build_projection_mlp( self, input_dim: int, output_dim: int, hidden_dim: int, num_layers: int, activation: str, use_layer_norm: bool, use_residual: bool, dropout: float, use_swiglu: bool, dtype: torch.dtype ) -> ModernMLP: """Build modern MLP for projection.""" return ModernMLP( input_dim=input_dim, output_dim=output_dim, hidden_dim=hidden_dim, num_layers=num_layers, activation=activation, use_layer_norm=use_layer_norm, use_residual=use_residual, dropout=dropout, use_swiglu=use_swiglu, dtype=dtype ) def _build_gate_components(self, dtype: torch.dtype): """Build gate logit components based on configuration.""" if not self.gate_depends_on_input: # Parameter-based gate gate_shape = self._get_parameter_shape(self.gate_granularity) self.gate_logit = nn.Parameter(torch.zeros(gate_shape, dtype=dtype)) else: # Input-dependent gate via MLP input_dim = self._get_gate_input_dim() output_dim = self._get_gate_output_dim() self.gate_generator = ModernMLP( input_dim=input_dim, output_dim=output_dim, hidden_dim=self.hidden_dim, num_layers=2, activation="gelu", use_layer_norm=True, use_residual=False, dropout=0.1, dtype=dtype ) def _build_weight_components(self, weight_init_value: float, dtype: torch.dtype): """Build weight components based on configuration.""" if not self.weight_depends_on_input: # Parameter-based weights weight_shape = self._get_parameter_shape(self.weight_granularity) self.key_weight = nn.Parameter(torch.full(weight_shape, weight_init_value, dtype=dtype)) self.value_weight = nn.Parameter(torch.full(weight_shape, weight_init_value, dtype=dtype)) else: # Input-dependent weights via MLP input_dim = self._get_weight_input_dim() output_dim = self._get_weight_output_dim() # Shared hidden layer for efficiency self.weight_hidden = ModernMLP( input_dim=input_dim, output_dim=self.weight_hidden_dim, hidden_dim=self.weight_hidden_dim, num_layers=2, activation="gelu", use_layer_norm=True, use_residual=False, dropout=0.1, dtype=dtype ) # Separate heads for key and value weights self.key_weight_head = nn.Linear(self.weight_hidden_dim, output_dim, dtype=dtype) self.value_weight_head = nn.Linear(self.weight_hidden_dim, output_dim, dtype=dtype) def _get_parameter_shape(self, granularity: str) -> tuple: """Get parameter shape based on granularity.""" if granularity == "scalar": return () # Scalar elif granularity == "token": return (self.max_sequence_length,) # Token-level parameters with max sequence length elif granularity == "head": return (self.max_sequence_length, self.target_num_heads) # Token and head level parameters elif granularity == "head_merged": return (self.max_sequence_length, self.target_num_heads) # Token and head level parameters elif granularity == "value": return (self.max_sequence_length, self.target_num_heads, self.target_dim) # Token, head and value level parameters else: raise ValueError(f"Invalid granularity: {granularity}") def _get_gate_input_dim(self) -> int: """Get input dimension for gate generator.""" base_dim = 0 if self.gate_input_features == "target_key": base_dim = self.target_dim elif self.gate_input_features == "target_value": base_dim = self.target_dim elif self.gate_input_features == "both": base_dim = self.target_dim * 2 elif self.gate_input_features == "target_projected_key": base_dim = self.target_dim * 2 # target_key + projected_key elif self.gate_input_features == "target_projected_value": base_dim = self.target_dim * 2 # target_value + projected_value elif self.gate_input_features == "target_projected_both": base_dim = self.target_dim * 4 # target_key + target_value + projected_key + projected_value else: raise ValueError(f"Invalid gate input features: {self.gate_input_features}") # Adjust for granularity processing strategy if self.gate_granularity == "scalar": # Scalar: process aggregated features across all heads return base_dim # Use pooled features elif self.gate_granularity == "token": # Token: process merged head dimensions return base_dim * self.target_num_heads # Flatten (H, D) to (H*D) elif self.gate_granularity == "head_merged": # Head-merged: similar to token granularity, merge H and D return base_dim * self.target_num_heads # (B, N, H*D) elif self.gate_granularity == "head": # Head-local: per head processing, do not merge heads return base_dim # (B, H, N, D) else: # value # Value: process per-head features return base_dim # Keep per-head processing (B, H, N, D) def _get_gate_output_dim(self) -> int: """Get output dimension for gate generator.""" if self.gate_granularity == "scalar": return 1 elif self.gate_granularity == "token": return 1 # Per token elif self.gate_granularity == "head_merged": # Per token per head after merge: output one value per head return self.target_num_heads elif self.gate_granularity == "head": # Per token per head: scalar per head return 1 elif self.gate_granularity == "value": return self.target_dim # Per token per head per value (but processed per-head, so output D per head) else: raise ValueError(f"Invalid gate granularity: {self.gate_granularity}") def _get_weight_input_dim(self) -> int: """Get input dimension for weight generator.""" base_dim = 0 if self.weight_input_features == "target_key": base_dim = self.target_dim elif self.weight_input_features == "target_value": base_dim = self.target_dim elif self.weight_input_features == "both": base_dim = self.target_dim * 2 elif self.weight_input_features == "target_projected_key": base_dim = self.target_dim * 2 # target_key + projected_key elif self.weight_input_features == "target_projected_value": base_dim = self.target_dim * 2 # target_value + projected_value elif self.weight_input_features == "target_projected_both": base_dim = self.target_dim * 4 # target_key + target_value + projected_key + projected_value else: raise ValueError(f"Invalid weight input features: {self.weight_input_features}") # Adjust for granularity processing strategy if self.weight_granularity == "scalar": # Scalar: process aggregated features across all heads return base_dim # Use pooled features elif self.weight_granularity == "token": # Token: process merged head dimensions return base_dim * self.target_num_heads # Flatten (H, D) to (H*D) elif self.weight_granularity == "head_merged": # Head-merged: similar to token granularity, merge H and D return base_dim * self.target_num_heads # (B, N, H*D) elif self.weight_granularity == "head": # Head-local: per head processing, do not merge heads return base_dim # (B, H, N, D) else: # value # Value: process per-head features return base_dim # Keep per-head processing (B, H, N, D) def _get_weight_output_dim(self) -> int: """Get output dimension for weight generator.""" if self.weight_granularity == "scalar": return 1 elif self.weight_granularity == "token": return 1 # Per token elif self.weight_granularity == "head_merged": # Per token per head after merge: output one value per head return self.target_num_heads elif self.weight_granularity == "head": # Per token per head: scalar per head return 1 elif self.weight_granularity == "value": return self.target_dim # Per token per head per value (but processed per-head, so output D per head) else: raise ValueError(f"Invalid weight granularity: {self.weight_granularity}") def _generate_gates(self, target_key: Tensor, target_value: Tensor, projected_key: Tensor = None, projected_value: Tensor = None) -> Tensor: """Generate gate logits based on configuration.""" if not self.gate_depends_on_input: # Use parameter-based gate return self.gate_logit else: # Generate input-dependent gate # First, prepare the base input features if self.gate_input_features == "target_key": base_input = target_key elif self.gate_input_features == "target_value": base_input = target_value elif self.gate_input_features == "both": base_input = torch.cat([target_key, target_value], dim=-1) elif self.gate_input_features == "target_projected_key": if projected_key is None: raise ValueError("projected_key is required for target_projected_key input features") base_input = torch.cat([target_key, projected_key], dim=-1) elif self.gate_input_features == "target_projected_value": if projected_value is None: raise ValueError("projected_value is required for target_projected_value input features") base_input = torch.cat([target_value, projected_value], dim=-1) elif self.gate_input_features == "target_projected_both": if projected_key is None or projected_value is None: raise ValueError("Both projected_key and projected_value are required for target_projected_both input features") base_input = torch.cat([target_key, target_value, projected_key, projected_value], dim=-1) # Now process based on granularity # base_input shape: (B, H, N, D_input) B, H, N, D_input = base_input.shape if self.gate_granularity == "scalar": # For scalar granularity, aggregate all dimensions: (B, H, N, D_input) -> (B, D_input) gate_input = base_input.mean(dim=(1, 2)) # Average over heads and tokens elif self.gate_granularity == "token": # For token granularity, merge H and D_input dimensions: (B, H, N, D_input) -> (B, N, H*D_input) gate_input = base_input.transpose(1, 2).contiguous().view(B, N, H * D_input) elif self.gate_granularity == "head_merged": # For head granularity, merge H and D like token: (B, H, N, D_in) -> (B, N, H*D_in) gate_input = base_input.transpose(1, 2).contiguous().view(B, N, H * D_input) elif self.gate_granularity == "head": # For head granularity, keep per-head processing: (B, H, N, D_input) gate_input = base_input elif self.gate_granularity == "value": # For value granularity, keep per-head processing: (B, H, N, D_input) gate_input = base_input return self.gate_generator(gate_input) def _generate_weights(self, target_key: Tensor, target_value: Tensor, projected_key: Tensor = None, projected_value: Tensor = None) -> Tuple[Tensor, Tensor]: """Generate weights based on configuration.""" if not self.weight_depends_on_input: # Use parameter-based weights return self.key_weight, self.value_weight else: # Generate input-dependent weights # First, prepare the base input features if self.weight_input_features == "target_key": base_input = target_key elif self.weight_input_features == "target_value": base_input = target_value elif self.weight_input_features == "both": base_input = torch.cat([target_key, target_value], dim=-1) elif self.weight_input_features == "target_projected_key": if projected_key is None: raise ValueError("projected_key is required for target_projected_key input features") base_input = torch.cat([target_key, projected_key], dim=-1) elif self.weight_input_features == "target_projected_value": if projected_value is None: raise ValueError("projected_value is required for target_projected_value input features") base_input = torch.cat([target_value, projected_value], dim=-1) elif self.weight_input_features == "target_projected_both": if projected_key is None or projected_value is None: raise ValueError("Both projected_key and projected_value are required for target_projected_both input features") base_input = torch.cat([target_key, target_value, projected_key, projected_value], dim=-1) # Now process based on granularity # base_input shape: (B, H, N, D_input) B, H, N, D_input = base_input.shape if self.weight_granularity == "scalar": # For scalar granularity, aggregate all dimensions: (B, H, N, D_input) -> (B, D_input) weight_input = base_input.mean(dim=(1, 2)) # Average over heads and tokens elif self.weight_granularity == "token": # For token granularity, merge H and D_input dimensions: (B, H, N, D_input) -> (B, N, H*D_input) weight_input = base_input.transpose(1, 2).contiguous().view(B, N, H * D_input) elif self.weight_granularity == "head_merged": # For head granularity, merge H and D like token: (B, H, N, D_in) -> (B, N, H*D_in) weight_input = base_input.transpose(1, 2).contiguous().view(B, N, H * D_input) elif self.weight_granularity == "head": # For head granularity, keep per-head processing: (B, H, N, D_input) weight_input = base_input elif self.weight_granularity == "value": # For value granularity, keep per-head processing: (B, H, N, D_input) weight_input = base_input weight_hidden = self.weight_hidden(weight_input) key_weight = self.key_weight_head(weight_hidden) value_weight = self.value_weight_head(weight_hidden) return key_weight, value_weight def _apply_gumbel_sigmoid(self, gate_logit: Tensor) -> Tensor: """Apply Gumbel sigmoid trick for training.""" if self.training and self.use_gumbel: gumbel_noise = self._sample_gumbel(gate_logit.shape, gate_logit.device, gate_logit.dtype) return torch.sigmoid((gate_logit + gumbel_noise) / self.gate_temperature) else: return (gate_logit > 0).float() @staticmethod def _sample_gumbel(shape: tuple, device: torch.device, dtype: torch.dtype, eps: float = 1e-20) -> Tensor: """Sample from Gumbel distribution.""" u = torch.rand(shape, device=device, dtype=dtype) return -torch.log(-torch.log(u + eps) + eps) def _reshape_for_granularity(self, tensor: Tensor, granularity: str, target_shape: tuple) -> Tensor: """Reshape tensor to match target shape based on granularity.""" B, H, N, D = target_shape if granularity == "scalar": # Scalar -> (B, H, N, D) return tensor.view(1, 1, 1, 1).expand(B, H, N, D) elif granularity == "token": # (max_seq_len,) -> (B, H, N, D) - slice to actual sequence length token_params = tensor[:N] # Take first N tokens return token_params.view(1, 1, N, 1).expand(B, H, N, D) elif granularity == "head": # (max_seq_len, H) -> (B, H, N, D) - slice to actual sequence length, each token each head independent head_params = tensor[:N, :] # Take first N tokens, all heads: (N, H) return head_params.view(1, N, H, 1).transpose(1, 2).expand(B, H, N, D) # (1, N, H, 1) -> (1, H, N, 1) -> (B, H, N, D) elif granularity == "head_merged": raise NotImplementedError elif granularity == "value": # (max_seq_len, H, D) -> (B, H, N, D) - slice to actual sequence length, each token each head each value independent value_params = tensor[:N, :, :] # Take first N tokens: (N, H, D) return value_params.view(1, N, H, D).transpose(1, 2).expand(B, H, N, D) # (1, N, H, D) -> (1, H, N, D) -> (B, H, N, D) else: raise ValueError(f"Invalid granularity: {granularity}") def update_temperature(self, step: int): """Update temperature using exponential annealing schedule for gate only.""" # Update gate temperature gate_ratio = min(step / self.anneal_steps, 1.0) gate_temp = self.initial_temperature * (self.final_temperature / self.initial_temperature) ** gate_ratio self.gate_temperature.fill_(gate_temp) def forward(self, source_kv: Tuple[Tensor, Tensor], target_kv: Tuple[Tensor, Tensor], position_ids: Optional[Tensor] = None, max_pos: Optional[Tensor] = None) -> Tuple[Tensor, Tensor]: """ Forward pass with unified projection logic. Args: source_kv: Tuple of (key, value) tensors, each (B, H_s, N, D_s) target_kv: Tuple of (key, value) tensors, each (B, H_t, N, D_t) position_ids: Position ids tensor (B, N), optional, required if pos_emb=True Returns: Tuple of (key, value) tensors, each (B, H_t, N, D_t) """ source_key, source_value = source_kv target_key, target_value = target_kv # Get shapes B, H_s, N, D_s = source_key.shape _, H_t, _, D_t = target_key.shape # Reshape for projection: (B, H, N, D) -> (B, N, H*D) source_key_flat = source_key.transpose(1, 2).contiguous().view(B, N, H_s * D_s) source_value_flat = source_value.transpose(1, 2).contiguous().view(B, N, H_s * D_s) # Project source to target dimension projected_key_flat = self.key_projection(source_key_flat) # (B, N, H_t * D_t) projected_value_flat = self.value_projection(source_value_flat) # (B, N, H_t * D_t) # Handle concatenation if enabled if self.use_concat: target_key_flat = target_key.transpose(1, 2).contiguous().view(B, N, H_t * D_t) target_value_flat = target_value.transpose(1, 2).contiguous().view(B, N, H_t * D_t) # Concatenate and combine combined_key = torch.cat([projected_key_flat, target_key_flat], dim=-1) combined_value = torch.cat([projected_value_flat, target_value_flat], dim=-1) final_projected_key_flat = self.key_combiner(combined_key) final_projected_value_flat = self.value_combiner(combined_value) else: final_projected_key_flat = projected_key_flat final_projected_value_flat = projected_value_flat # Reshape back: (B, N, H_t * D_t) -> (B, H_t, N, D_t) projected_key = final_projected_key_flat.view(B, N, H_t, D_t).transpose(1, 2) projected_value = final_projected_value_flat.view(B, N, H_t, D_t).transpose(1, 2) # Generate gates, selectors and weights (may need projected tensors for input features) needs_projected_for_gate = self.gate_depends_on_input and self.gate_input_features in [ "target_projected_key", "target_projected_value", "target_projected_both" ] needs_projected_for_weight = self.weight_depends_on_input and self.weight_input_features in [ "target_projected_key", "target_projected_value", "target_projected_both" ] if needs_projected_for_gate or needs_projected_for_weight: gate_logit = self._generate_gates(target_key, target_value, projected_key, projected_value) key_weight, value_weight = self._generate_weights(target_key, target_value, projected_key, projected_value) else: gate_logit = self._generate_gates(target_key, target_value) key_weight, value_weight = self._generate_weights(target_key, target_value) # Reshape gates and weights to match target shape target_shape = (B, H_t, N, D_t) if self.gate_depends_on_input: # Reshape based on gate granularity - all preserve token dimension N if self.gate_granularity == "scalar": # For scalar, gate_logit is already (B, 1) from MLP, just expand gate_logit = gate_logit.view(B, 1, 1, 1).expand(target_shape) elif self.gate_granularity == "token": gate_logit = gate_logit.unsqueeze(1).unsqueeze(-1).expand(target_shape) # (B, N, 1) -> (B, H, N, D) elif self.gate_granularity == "head_merged": # (B, N, H) -> (B, H, N, D) - per token per head, broadcast over D gate_logit = gate_logit.permute(0, 2, 1).unsqueeze(-1).expand(B, H_t, N, D_t) elif self.gate_granularity == "head": # (B, H, N, 1) -> (B, H, N, D) - per token per head scalar, broadcast over D gate_logit = gate_logit.expand(B, H_t, N, D_t) elif self.gate_granularity == "value": # (B, H, N, D) -> (B, H, N, D) - each token each head each value has one value pass # Already in correct shape else: gate_logit = self._reshape_for_granularity(gate_logit, self.gate_granularity, target_shape) if self.weight_depends_on_input: # Reshape weights based on granularity - all preserve token dimension N if self.weight_granularity == "scalar": # For scalar, weights are already (B, 1) from MLP, just expand key_weight = key_weight.view(B, 1, 1, 1).expand(target_shape) value_weight = value_weight.view(B, 1, 1, 1).expand(target_shape) elif self.weight_granularity == "token": key_weight = key_weight.unsqueeze(1).expand(target_shape) # (B, N, 1) -> (B, H, N, D) value_weight = value_weight.unsqueeze(1).expand(target_shape) elif self.weight_granularity == "head_merged": # (B, N, H) -> (B, H, N, D) - per token per head, broadcast over D key_weight = key_weight.permute(0, 2, 1).unsqueeze(-1).expand(B, H_t, N, D_t) value_weight = value_weight.permute(0, 2, 1).unsqueeze(-1).expand(B, H_t, N, D_t) elif self.weight_granularity == "head": # (B, H, N, 1) -> (B, H, N, D) - per token per head scalar, broadcast over D key_weight = key_weight.expand(B, H_t, N, D_t) value_weight = value_weight.expand(B, H_t, N, D_t) elif self.weight_granularity == "value": # (B, H, N, D) -> (B, H, N, D) - each token each head each value has one value pass # Already in correct shape else: key_weight = self._reshape_for_granularity(key_weight, self.weight_granularity, target_shape) value_weight = self._reshape_for_granularity(value_weight, self.weight_granularity, target_shape) # Apply gating and selection gate = self._apply_gumbel_sigmoid(gate_logit) # Normalize weights using dynamic temperature normalized_key_weight = torch.sigmoid(key_weight / self.scalar_temperature) normalized_value_weight = torch.sigmoid(value_weight / self.scalar_temperature) # Final combination # Compute projected contribution (always present) projected_key_term = gate * normalized_key_weight * projected_key projected_value_term = gate * normalized_value_weight * projected_value # Compute target (self) contribution depending on flags if self.add_self: if self.preserve_target_weight: target_key_term = (1 - normalized_key_weight) * target_key target_value_term = (1 - normalized_value_weight) * target_value else: target_key_term = target_key target_value_term = target_value else: target_key_term = torch.zeros_like(target_key) target_value_term = torch.zeros_like(target_value) # Final outputs output_key = target_key_term + projected_key_term output_value = target_value_term + projected_value_term return (output_key, output_value) class QwenStyleLayer(nn.Module): """ One Qwen3-style MLP sublayer: y = x + Dropout( down( SiLU(gate(LN(x))) * up(LN(x)) ) ) - Pre-norm with RMSNorm - Bias-free linears """ def __init__(self, hidden_size: int, intermediate_size: int, dropout: float = 0.0, dtype: torch.dtype = torch.float32): super().__init__() self.norm = nn.RMSNorm(hidden_size, eps=1e-6, dtype=dtype) self.gate = nn.Linear(hidden_size, intermediate_size, bias=False, dtype=dtype) self.up = nn.Linear(hidden_size, intermediate_size, bias=False, dtype=dtype) self.down = nn.Linear(intermediate_size, hidden_size, bias=False, dtype=dtype) self.act = nn.SiLU() self.drop = nn.Dropout(dropout) if dropout > 0 else nn.Identity() def forward(self, x: Tensor) -> Tensor: h = self.norm(x) h = self.act(self.gate(h)) * self.up(h) # SwiGLU h = self.down(h) h = self.drop(h) return x + h class StandardFFNLayer(nn.Module): """ Pre-norm RMSNorm, classic MLP: y = x + Dropout( W2( Act( W1( RMSNorm(x) ) ) ) ) - No SwiGLU: single hidden nonlinearity (GELU/ReLU/SiLU) - Bias-free linears (common in modern LLM FFNs) """ def __init__( self, hidden_size: int, intermediate_size: int, dropout: float = 0.0, dtype: torch.dtype = torch.float32, activation: str = "gelu", ): super().__init__() self.norm = nn.RMSNorm(hidden_size, eps=1e-6, dtype=dtype) self.w1 = nn.Linear(hidden_size, intermediate_size, bias=False, dtype=dtype) self.w2 = nn.Linear(intermediate_size, hidden_size, bias=False, dtype=dtype) self.drop = nn.Dropout(dropout) if dropout > 0 else nn.Identity() act = activation.lower() if act == "gelu": self.act = nn.GELU() elif act == "relu": self.act = nn.ReLU() elif act == "silu": self.act = nn.SiLU() else: raise ValueError(f"Unsupported activation: {activation}") def forward(self, x: Tensor) -> Tensor: h = self.norm(x) h = self.act(self.w1(h)) h = self.w2(h) h = self.drop(h) return x + h class RegularMLP(nn.Module): """ Qwen3-style stacked MLP operating at a fixed hidden size. - No input/output projections; caller is responsible for projections. - num_layers repeats of Qwen-style FFN sublayer (pre-RMSNorm, SwiGLU, bias-free) """ def __init__( self, hidden_dim: int = 1024, intermediate_dim: int = 3072, num_layers: int = 3, dropout: float = 0.1, dtype: torch.dtype = torch.float32, ): super().__init__() assert num_layers >= 1, "num_layers must be >= 1" self.blocks = nn.ModuleList([ StandardFFNLayer(hidden_size=hidden_dim, intermediate_size=intermediate_dim, dropout=dropout, dtype=dtype) for _ in range(num_layers) ]) def forward(self, x: Tensor) -> Tensor: for blk in self.blocks: x = blk(x) return x @register_model @capture_init_args class C2CProjector(Projector): """ Concise projector specialized to a fixed C2C configuration using StandardMLP. - Projections: StandardMLP (pre-RMSNorm, SwiGLU, residual per sublayer) - Concat: enabled, followed by linear combiner to target size - Gate: scalar parameter with Gumbel-sigmoid during training - Weights: input-dependent, head_merged granularity using target and projected key - Target preservation: add_self=True, preserve_target_weight=False - Temperatures: annealed gate temperature (1.0 -> 0.001 over 1929 steps), scalar_temperature=1.0 """ def __init__( self, source_dim: int, target_dim: int, source_num_heads: int = 1, target_num_heads: int = 1, intermediate_dim: int = 1024, hidden_dim: int = 1024, num_layers: int = 3, dropout: float = 0.1, initial_temperature: float = 1.0, final_temperature: float = 0.001, anneal_steps: int = 1929, dtype: torch.dtype = torch.float32, ): super().__init__() assert num_layers >= 3, "num_layers must be >= 3" # Dimensions self.source_dim = source_dim self.target_dim = target_dim self.source_num_heads = source_num_heads self.target_num_heads = target_num_heads # Sizes in_dim = source_dim * source_num_heads out_dim = target_dim * target_num_heads # 1) concat(source_X, target_X) then project to hidden_dim self.key_in = nn.Linear(in_dim + out_dim, hidden_dim, bias=True, dtype=dtype) self.value_in = nn.Linear(in_dim + out_dim, hidden_dim, bias=True, dtype=dtype) # 2) one-layer common embedding MLP to get intermediate representation (at hidden_dim) self.key_mlp1 = RegularMLP(hidden_dim=hidden_dim, intermediate_dim=intermediate_dim, num_layers=1, dropout=dropout, dtype=dtype) self.value_mlp1 = RegularMLP(hidden_dim=hidden_dim, intermediate_dim=intermediate_dim, num_layers=1, dropout=dropout, dtype=dtype) # 3a) intermediate representation → (L-2)-layer MLP for weights → project to head dim self.key_scalar_mlp2 = RegularMLP(hidden_dim=hidden_dim, intermediate_dim=hidden_dim, num_layers=1, dropout=dropout, dtype=dtype) self.value_scalar_mlp2 = RegularMLP(hidden_dim=hidden_dim, intermediate_dim=hidden_dim, num_layers=1, dropout=dropout, dtype=dtype) self.key_scalar_head = nn.Linear(hidden_dim, target_num_heads, dtype=dtype) self.value_scalar_head = nn.Linear(hidden_dim, target_num_heads, dtype=dtype) # 3b) intermediate representation → (L-2)-layer MLP for projected_X → finally project hidden_dim → out_dim self.key_proj_mlp2 = RegularMLP(hidden_dim=hidden_dim, intermediate_dim=intermediate_dim, num_layers=num_layers-2, dropout=dropout, dtype=dtype) self.value_proj_mlp2 = RegularMLP(hidden_dim=hidden_dim, intermediate_dim=intermediate_dim, num_layers=num_layers-2, dropout=dropout, dtype=dtype) self.key_proj_out = nn.Linear(hidden_dim, out_dim, bias=True, dtype=dtype) self.value_proj_out = nn.Linear(hidden_dim, out_dim, bias=True, dtype=dtype) # Scalar key/value gate parameters and temperature schedule self.key_gate_logit = nn.Parameter(torch.tensor(0.0, dtype=dtype)) self.value_gate_logit = nn.Parameter(torch.tensor(0.0, dtype=dtype)) self.use_gumbel = True self.register_buffer("gate_temperature", torch.tensor(initial_temperature, dtype=dtype)) self.initial_temperature = initial_temperature self.final_temperature = final_temperature self.anneal_steps = anneal_steps # Temperature for weight normalization self.scalar_temperature = 1.0 def update_temperature(self, step: int): ratio = min(step / self.anneal_steps, 1.0) temp = self.initial_temperature * (self.final_temperature / self.initial_temperature) ** ratio self.gate_temperature.fill_(temp) def forward( self, source_kv: Tuple[Tensor, Tensor], target_kv: Tuple[Tensor, Tensor], position_ids: Optional[Tensor] = None, max_pos: Optional[Tensor] = None, ) -> Tuple[Tensor, Tensor]: source_key, source_value = source_kv target_key, target_value = target_kv B, Hs, N, Ds = source_key.shape _, Ht, _, Dt = target_key.shape # Flatten heads source_key_flat = source_key.transpose(1, 2).contiguous().view(B, N, Hs * Ds) source_value_flat = source_value.transpose(1, 2).contiguous().view(B, N, Hs * Ds) target_key_flat = target_key.transpose(1, 2).contiguous().view(B, N, Ht * Dt) target_value_flat = target_value.transpose(1, 2).contiguous().view(B, N, Ht * Dt) # 1) concat source and target features along channel key_cat = torch.cat([source_key_flat, target_key_flat], dim=-1) value_cat = torch.cat([source_value_flat, target_value_flat], dim=-1) # 2) project to hidden dim key_hidden = self.key_in(key_cat) value_hidden = self.value_in(value_cat) # 3) one-layer common embedding MLP to get intermediate representation (at hidden_dim) key_hidden = self.key_mlp1(key_hidden) value_hidden = self.value_mlp1(value_hidden) # 4b) intermediate representation -> projected feature path key_proj_hidden = self.key_proj_out(self.key_proj_mlp2(key_hidden)) # (B, N, Ht * Dt) value_proj_hidden = self.value_proj_out(self.value_proj_mlp2(value_hidden)) # (B, N, Ht * Dt) projected_key = key_proj_hidden.view(B, N, Ht, Dt).transpose(1, 2) # (B, Ht, N, Dt) projected_value = value_proj_hidden.view(B, N, Ht, Dt).transpose(1, 2) # (B, Ht, N, Dt) # 4a) intermediate representation -> scalar path key_scalar = self.key_scalar_head(self.key_scalar_mlp2(key_hidden)) # (B, N, Ht) value_scalar = self.value_scalar_head(self.value_scalar_mlp2(value_hidden)) # (B, N, Ht) key_scalar = key_scalar.permute(0, 2, 1).unsqueeze(-1) # (B, Ht, N, 1) value_scalar = value_scalar.permute(0, 2, 1).unsqueeze(-1) # (B, Ht, N, 1) # Key/value gates: element-wise Gumbel noise with scalar logits (broadcast over channels) key_gate_logit = self.key_gate_logit.view(1, 1, 1, 1) value_gate_logit = self.value_gate_logit.view(1, 1, 1, 1) if self.training and self.use_gumbel: u1 = torch.rand(B, Ht, N, 1, device=key_gate_logit.device, dtype=key_gate_logit.dtype) u2 = torch.rand(B, Ht, N, 1, device=value_gate_logit.device, dtype=value_gate_logit.dtype) g1 = -torch.log(-torch.log(u1 + 1e-20) + 1e-20) g2 = -torch.log(-torch.log(u2 + 1e-20) + 1e-20) key_gate = torch.sigmoid((key_gate_logit + g1) / self.gate_temperature) value_gate = torch.sigmoid((value_gate_logit + g2) / self.gate_temperature) else: key_gate = (key_gate_logit > 0).float() value_gate = (value_gate_logit > 0).float() # Normalize scalars (scalar_temperature=1.0) norm_key_scalar = torch.sigmoid(key_scalar) norm_value_scalar = torch.sigmoid(value_scalar) # Combine (preserve_target_weight=False, add_self=True) output_key = target_key + key_gate * norm_key_scalar * projected_key output_value = target_value + value_gate * norm_value_scalar * projected_value # Expose capture attributes for downstream analysis scripts try: # Store normalized scalars (detach to avoid autograd, keep device-agnostic via CPU) self.last_norm_key_scalar = norm_key_scalar.detach().cpu() self.last_norm_value_scalar = norm_value_scalar.detach().cpu() # Store gate logits as python floats (parameters are scalar) self.last_key_gate_logit = float(self.key_gate_logit.detach().cpu().item()) self.last_value_gate_logit = float(self.value_gate_logit.detach().cpu().item()) except Exception: # Best-effort capture; never break forward path pass return output_key, output_value def save_projector(obj: Projector, file_path: str) -> None: save_object(obj, file_path) def load_projector(file_path: str, override_args: Optional[dict] = None) -> Projector: return load_object(file_path, get_projector_class, override_args) def create_projector(projector_type: str, **kwargs) -> Projector: """ Factory function to create a projector based on type. Args: projector_type: String indicating the type of projector **kwargs: Additional arguments to pass to the projector constructor Returns: An instance of the appropriate projector """ # Prefer using the unified registry getter (handles case-insensitive keys) try: cls = get_projector_class(projector_type) except ValueError as e: raise e return cls(**kwargs)