| """Learned dense-to-local+algebra overlap blocks with physical export.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
| from strata.modeling.algebra import HardConcreteGate |
| from strata.modeling.compose.lm_adapter import AlgebraGraphReadAdapter, SparseAlgebraGraphRead |
| from strata.modeling.ph_pat.config import PHPATConfig |
| from strata.modeling.ph_pat.replacement_block import ( |
| ChunkedLocalCausalAttention, |
| DenseGlobalBlock, |
| DenseGlobalCausalAttention, |
| ) |
|
|
|
|
| class AlgebraOverlapBlock(nn.Module): |
| """Exact dense endpoint and trainable local+graph endpoint for one block.""" |
|
|
| def __init__( |
| self, |
| dense_block: DenseGlobalBlock, |
| config: PHPATConfig, |
| *, |
| gamma_max: float = 0.02, |
| ) -> None: |
| super().__init__() |
| self.attn_norm = dense_block.attn_norm |
| self.global_attention = dense_block.attention |
| self.local_attention = ChunkedLocalCausalAttention(config) |
| self.local_attention.load_state_dict(self.global_attention.state_dict(), strict=True) |
| self.ffn_norm = dense_block.ffn_norm |
| self.ffn = dense_block.ffn |
| self.dropout = dense_block.dropout |
| self.graph_adapter = AlgebraGraphReadAdapter(config.d_model, config.d_model, gamma_max=gamma_max) |
| self.global_gate = HardConcreteGate(1, initial_probability=0.99) |
| reference = self.global_attention.qkv.weight |
| self.local_attention.to(device=reference.device, dtype=reference.dtype) |
| self.graph_adapter.to(device=reference.device, dtype=reference.dtype) |
| self.global_gate.to(device=reference.device) |
|
|
| def forward( |
| self, |
| hidden: torch.Tensor, |
| attention_mask: torch.Tensor, |
| *, |
| graph_values: torch.Tensor | None = None, |
| graph_reliability: torch.Tensor | None = None, |
| sparse_graph_read: SparseAlgebraGraphRead | None = None, |
| graph_enabled: bool = False, |
| gate_mode: str = "dense", |
| gate_override: torch.Tensor | float | None = None, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| normalized = self.attn_norm(hidden) |
| if gate_override is not None: |
| gate = torch.as_tensor(gate_override, device=hidden.device, dtype=hidden.dtype) |
| if gate.numel() != 1 or bool((gate < 0).any() or (gate > 1).any()): |
| raise ValueError("gate_override must be one scalar in [0, 1]") |
| gate = gate.reshape(()) |
| if float(gate.detach()) == 1.0: |
| attention = self.global_attention(normalized, attention_mask) |
| elif float(gate.detach()) == 0.0: |
| attention = self.local_attention(normalized, attention_mask) |
| else: |
| local = self.local_attention(normalized, attention_mask) |
| global_value = self.global_attention(normalized, attention_mask) |
| attention = local + gate * (global_value - local) |
| elif gate_mode == "dense": |
| attention = self.global_attention(normalized, attention_mask) |
| gate = hidden.new_tensor(1.0) |
| elif gate_mode == "local": |
| attention = self.local_attention(normalized, attention_mask) |
| gate = hidden.new_tensor(0.0) |
| elif gate_mode in ("sample", "expected"): |
| local = self.local_attention(normalized, attention_mask) |
| global_value = self.global_attention(normalized, attention_mask) |
| gate = self.global_gate(sample=gate_mode == "sample")[0].to(hidden.dtype) |
| attention = local + gate * (global_value - local) |
| else: |
| raise ValueError("gate_mode must be dense, local, sample, or expected") |
| hidden = hidden + self.dropout(attention) |
| if graph_enabled: |
| if sparse_graph_read is not None: |
| hidden, _update, _relative = self.graph_adapter.forward_sparse( |
| hidden, sparse_graph_read |
| ) |
| else: |
| if graph_values is None or graph_reliability is None: |
| raise ValueError("enabled graph reads require values and reliability") |
| hidden, _update, _relative = self.graph_adapter( |
| hidden, graph_values, graph_reliability |
| ) |
| hidden = hidden + self.dropout(self.ffn(self.ffn_norm(hidden))) |
| return hidden, gate |
|
|
| def export_local(self) -> "ExportedAlgebraBlock": |
| return ExportedAlgebraBlock(self) |
|
|
| def restore_global(self, config: PHPATConfig) -> DenseGlobalBlock: |
| """Restore the original dense block without retaining local/chart modules.""" |
|
|
| restored = DenseGlobalBlock(config) |
| restored.attn_norm = self.attn_norm |
| restored.attention = self.global_attention |
| restored.ffn_norm = self.ffn_norm |
| restored.ffn = self.ffn |
| restored.dropout = self.dropout |
| return restored |
|
|
|
|
| class ExportedAlgebraBlock(nn.Module): |
| """Physical local+graph block; no dense Q/K/V module remains.""" |
|
|
| def __init__(self, source: AlgebraOverlapBlock) -> None: |
| super().__init__() |
| self.attn_norm = source.attn_norm |
| self.local_attention = source.local_attention |
| self.ffn_norm = source.ffn_norm |
| self.ffn = source.ffn |
| self.dropout = source.dropout |
| self.graph_adapter = source.graph_adapter |
|
|
| def forward( |
| self, |
| hidden: torch.Tensor, |
| attention_mask: torch.Tensor, |
| *, |
| graph_values: torch.Tensor | None = None, |
| graph_reliability: torch.Tensor | None = None, |
| sparse_graph_read: SparseAlgebraGraphRead | None = None, |
| graph_enabled: bool = False, |
| ) -> torch.Tensor: |
| hidden = hidden + self.dropout(self.local_attention(self.attn_norm(hidden), attention_mask)) |
| if graph_enabled: |
| if self.graph_adapter is not None: |
| if sparse_graph_read is not None: |
| hidden, _update, _relative = self.graph_adapter.forward_sparse( |
| hidden, sparse_graph_read |
| ) |
| else: |
| if graph_values is None or graph_reliability is None: |
| raise ValueError("enabled graph reads require values and reliability") |
| hidden, _update, _relative = self.graph_adapter( |
| hidden, graph_values, graph_reliability |
| ) |
| return hidden + self.dropout(self.ffn(self.ffn_norm(hidden))) |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class PrunableLMOutput: |
| logits: torch.Tensor |
| loss: torch.Tensor | None |
| hidden_states: torch.Tensor |
| global_gates: torch.Tensor |
|
|
|
|
| class PrunableAlgebraLM(nn.Module): |
| """Install overlap blocks without altering unselected dense blocks.""" |
|
|
| def __init__( |
| self, |
| base_model: nn.Module, |
| config: PHPATConfig, |
| replacement_layers: tuple[int, ...], |
| *, |
| gamma_max: float = 0.02, |
| ) -> None: |
| super().__init__() |
| self.base_model = base_model |
| self.config = config |
| self.replacement_layers = tuple(sorted(int(index) for index in replacement_layers)) |
| if len(set(self.replacement_layers)) != len(self.replacement_layers): |
| raise ValueError("replacement layers must be unique") |
| for index in self.replacement_layers: |
| block = self.base_model.blocks[index] |
| if not isinstance(block, DenseGlobalBlock): |
| raise TypeError(f"layer {index} is not a dense global block") |
| self.base_model.blocks[index] = AlgebraOverlapBlock(block, config, gamma_max=gamma_max) |
|
|
| def install_overlap_layers( |
| self, |
| layer_indices: tuple[int, ...], |
| *, |
| gamma_max: float = 0.02, |
| ) -> tuple[int, ...]: |
| """Install candidate overlap paths while preserving prior physical exports.""" |
|
|
| requested = tuple(sorted(set(int(index) for index in layer_indices))) |
| installed = [] |
| for index in requested: |
| block = self.base_model.blocks[index] |
| if isinstance(block, (AlgebraOverlapBlock, ExportedAlgebraBlock)): |
| continue |
| if not isinstance(block, DenseGlobalBlock): |
| raise TypeError(f"layer {index} is not a dense global block") |
| self.base_model.blocks[index] = AlgebraOverlapBlock( |
| block, |
| self.config, |
| gamma_max=gamma_max, |
| ) |
| installed.append(index) |
| self.replacement_layers = tuple(sorted(set(self.replacement_layers).union(installed))) |
| return tuple(installed) |
|
|
| def restore_global_layers(self, layer_indices: tuple[int, ...]) -> tuple[int, ...]: |
| """Remove unselected overlap paths and retain their exact dense branches.""" |
|
|
| restored = [] |
| for index in sorted(set(int(value) for value in layer_indices)): |
| block = self.base_model.blocks[index] |
| if isinstance(block, ExportedAlgebraBlock): |
| raise ValueError(f"cannot restore physically exported layer {index}") |
| if not isinstance(block, AlgebraOverlapBlock): |
| continue |
| self.base_model.blocks[index] = block.restore_global(self.config) |
| restored.append(index) |
| if restored: |
| removed = set(restored) |
| self.replacement_layers = tuple( |
| index for index in self.replacement_layers if index not in removed |
| ) |
| return tuple(restored) |
|
|
| def export_layers(self, layer_indices: tuple[int, ...]) -> tuple[int, ...]: |
| """Physically export an explicit set of overlap branches to local blocks.""" |
|
|
| exported = [] |
| for index in sorted(set(int(value) for value in layer_indices)): |
| block = self.base_model.blocks[index] |
| if not isinstance(block, AlgebraOverlapBlock): |
| raise TypeError(f"layer {index} is not an overlap candidate") |
| self.base_model.blocks[index] = block.export_local() |
| exported.append(index) |
| return tuple(exported) |
|
|
| def forward_hidden( |
| self, |
| input_ids: torch.Tensor, |
| *, |
| attention_mask: torch.Tensor | None = None, |
| graph_values: torch.Tensor | None = None, |
| graph_reliability: torch.Tensor | None = None, |
| sparse_graph_read: SparseAlgebraGraphRead | None = None, |
| graph_enabled: bool = False, |
| graph_layer_indices: tuple[int, ...] | None = None, |
| gate_mode: str = "dense", |
| global_mask: torch.Tensor | None = None, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| if attention_mask is None: |
| attention_mask = torch.ones_like(input_ids, dtype=torch.bool) |
| else: |
| attention_mask = attention_mask.to(torch.bool) |
| positions = torch.arange(input_ids.shape[1], device=input_ids.device) |
| hidden = self.base_model.dropout( |
| self.base_model.token_embeddings(input_ids) |
| + self.base_model.position_embeddings(positions).unsqueeze(0) |
| ) |
| gates = [] |
| replacement_index = 0 |
| active_graph_layers = None if graph_layer_indices is None else set(graph_layer_indices) |
| for layer_index, block in enumerate(self.base_model.blocks): |
| if isinstance(block, AlgebraOverlapBlock): |
| override = None |
| if global_mask is not None: |
| if global_mask.shape != (len(self.replacement_layers),): |
| raise ValueError("global_mask must have one value per replacement layer") |
| override = global_mask[replacement_index] |
| hidden, gate = block( |
| hidden, |
| attention_mask, |
| graph_values=graph_values, |
| graph_reliability=graph_reliability, |
| sparse_graph_read=sparse_graph_read, |
| graph_enabled=graph_enabled and ( |
| active_graph_layers is None or layer_index in active_graph_layers |
| ), |
| gate_mode=gate_mode, |
| gate_override=override, |
| ) |
| gates.append(gate.reshape(1)) |
| replacement_index += 1 |
| elif isinstance(block, ExportedAlgebraBlock): |
| hidden = block( |
| hidden, |
| attention_mask, |
| graph_values=graph_values, |
| graph_reliability=graph_reliability, |
| sparse_graph_read=sparse_graph_read, |
| graph_enabled=graph_enabled and ( |
| active_graph_layers is None or layer_index in active_graph_layers |
| ), |
| ) |
| gates.append(hidden.new_zeros(1)) |
| replacement_index += 1 |
| else: |
| hidden, _read = block(hidden, attention_mask) |
| return hidden, attention_mask, torch.cat(gates) if gates else hidden.new_empty(0) |
|
|
| def forward( |
| self, |
| input_ids: torch.Tensor, |
| *, |
| attention_mask: torch.Tensor | None = None, |
| labels: torch.Tensor | None = None, |
| graph_values: torch.Tensor | None = None, |
| graph_reliability: torch.Tensor | None = None, |
| sparse_graph_read: SparseAlgebraGraphRead | None = None, |
| graph_enabled: bool = False, |
| graph_layer_indices: tuple[int, ...] | None = None, |
| gate_mode: str = "dense", |
| global_mask: torch.Tensor | None = None, |
| ) -> PrunableLMOutput: |
| hidden, _mask, gates = self.forward_hidden( |
| input_ids, |
| attention_mask=attention_mask, |
| graph_values=graph_values, |
| graph_reliability=graph_reliability, |
| sparse_graph_read=sparse_graph_read, |
| graph_enabled=graph_enabled, |
| graph_layer_indices=graph_layer_indices, |
| gate_mode=gate_mode, |
| global_mask=global_mask, |
| ) |
| hidden = self.base_model.final_norm(hidden) |
| logits = self.base_model.lm_head(hidden) |
| loss = None |
| if labels is not None: |
| loss = F.cross_entropy( |
| logits[:, :-1].contiguous().view(-1, logits.shape[-1]), |
| labels[:, 1:].contiguous().view(-1), |
| ignore_index=-100, |
| ) |
| return PrunableLMOutput(logits, loss, hidden, gates) |
|
|
| def expected_global_branches(self) -> torch.Tensor: |
| values = [ |
| block.global_gate.expected_l0()[0] |
| for block in self.base_model.blocks |
| if isinstance(block, AlgebraOverlapBlock) |
| ] |
| return torch.stack(values) if values else next(self.parameters()).new_empty(0) |
|
|
| def deterministic_global_branches(self) -> torch.Tensor: |
| values = [ |
| block.global_gate(sample=False)[0] |
| for block in self.base_model.blocks |
| if isinstance(block, AlgebraOverlapBlock) |
| ] |
| return torch.stack(values) if values else next(self.parameters()).new_empty(0) |
|
|
| def export_zero_gates(self, *, threshold: float = 0.5) -> tuple[int, ...]: |
| exported = [] |
| for index, block in enumerate(self.base_model.blocks): |
| if isinstance(block, AlgebraOverlapBlock) and not bool( |
| block.global_gate.export_mask(threshold=threshold)[0] |
| ): |
| self.base_model.blocks[index] = block.export_local() |
| exported.append(index) |
| return tuple(exported) |
|
|
| def dense_modules_in_replacement_layers(self) -> tuple[int, ...]: |
| found = [] |
| for index in self.replacement_layers: |
| if any(isinstance(module, DenseGlobalCausalAttention) for module in self.base_model.blocks[index].modules()): |
| found.append(index) |
| return tuple(found) |
|
|
| def physically_removed_layers(self) -> tuple[int, ...]: |
| return tuple( |
| index |
| for index, block in enumerate(self.base_model.blocks) |
| if isinstance(block, ExportedAlgebraBlock) |
| ) |
|
|
| def active_graph_layers(self, *, tolerance: float = 0.0) -> tuple[int, ...]: |
| """Return physical graph blocks whose BRR gate is not identically zero.""" |
|
|
| return tuple( |
| index |
| for index, block in enumerate(self.base_model.blocks) |
| if isinstance(block, ExportedAlgebraBlock) |
| and block.graph_adapter is not None |
| and bool(block.graph_adapter.brr.gate_parameter.detach().abs().max() > tolerance) |
| ) |
|
|
| def strip_inactive_graph_adapters(self, *, tolerance: float = 0.0) -> tuple[int, ...]: |
| """Physically remove adapters whose BRR contribution is exactly zero.""" |
|
|
| stripped = [] |
| for index, block in enumerate(self.base_model.blocks): |
| if not isinstance(block, ExportedAlgebraBlock) or block.graph_adapter is None: |
| continue |
| gate = block.graph_adapter.brr.gate_parameter.detach().abs().max() |
| if bool(gate <= tolerance): |
| block.graph_adapter = None |
| stripped.append(index) |
| return tuple(stripped) |
|
|
|
|
| __all__ = [ |
| "AlgebraOverlapBlock", |
| "ExportedAlgebraBlock", |
| "PrunableAlgebraLM", |
| "PrunableLMOutput", |
| ] |
|
|
