| """Frozen-Pythia attachment for the split canonical router/translator.""" |
|
|
| from __future__ import annotations |
|
|
| import torch |
| from torch import Tensor, nn |
| from pathlib import Path |
|
|
| from pcm.planner.canonical import CanonicalPStore |
| from pcm.planner.split_translator import ( |
| ByteEntityEncoder, |
| CanonicalPRouter, |
| CanonicalRouterIndex, |
| RouteResult, |
| FactorizedCanonicalQuery, |
| SplitPTranslatePackage, |
| config_checksum, |
| ) |
|
|
|
|
| def pythia_model_identifier(base_model: nn.Module) -> str: |
| configured = str(getattr(base_model.config, "_name_or_path", "")).strip() |
| if configured and Path(configured).is_absolute(): |
| configured = Path(configured).name |
| return configured or str(getattr(base_model.config, "model_type", "gpt_neox")) |
|
|
|
|
| class PythiaSplitTranslatedModel(nn.Module): |
| def __init__( |
| self, |
| base_model: nn.Module, |
| package: SplitPTranslatePackage, |
| router: CanonicalPRouter, |
| entity_encoder: ByteEntityEncoder, |
| ) -> None: |
| super().__init__() |
| if not hasattr(base_model, "gpt_neox"): |
| raise TypeError("base model must expose GPT-NeoX transformer layers") |
| package.validate_compatibility( |
| model_id=pythia_model_identifier(base_model), |
| model_hidden_width=int(base_model.config.hidden_size), |
| attachment_layers=package.config.attachment_layers, |
| model_config_sha256=config_checksum(base_model.config), |
| ) |
| self.base_model = base_model |
| self.package = package |
| self.router = router |
| self.entity_encoder = entity_encoder |
| for parameter in base_model.parameters(): |
| parameter.requires_grad_(False) |
| layers = base_model.gpt_neox.layers |
| if any(index < 0 or index >= len(layers) for index in package.config.attachment_layers): |
| raise IndexError("split translator attachment layer is outside Pythia depth") |
| self._store: CanonicalPStore | None = None |
| self._index: CanonicalRouterIndex | None = None |
| self._oracle_indices: Tensor | None = None |
| self._query_entity_anchor: Tensor | None = None |
| self._gate_enabled = True |
| self._injection_enabled = True |
| self._collect = False |
| self._gate_telemetry: list[Tensor] = [] |
| self._route_telemetry: list[RouteResult] = [] |
| self._query_telemetry: list[FactorizedCanonicalQuery] = [] |
| self._handles = [ |
| layers[index].register_forward_hook(self._hook) |
| for index in package.config.attachment_layers |
| ] |
| self.base_model.eval() |
|
|
| def _oracle_route(self, query, hidden: Tensor) -> RouteResult: |
| assert self._index is not None and self._oracle_indices is not None |
| scores, features = self.router.all_scores(query, self._index) |
| batch, sequence = hidden.shape[:2] |
| indices = self._oracle_indices.to(hidden.device).view(batch, 1, 1).expand(batch, sequence, 1) |
| selected_scores = scores.gather(-1, indices) |
| selected_features = features.gather( |
| -2, indices.unsqueeze(-1).expand(batch, sequence, 1, 4) |
| ) |
| return RouteResult( |
| indices=indices, |
| scores=selected_scores, |
| weights=torch.ones_like(selected_scores), |
| features=selected_features, |
| accepted=torch.ones((batch, sequence), dtype=torch.bool, device=hidden.device), |
| has_valid=True, |
| ) |
|
|
| def _hook(self, _module, _inputs, hidden: Tensor): |
| if self._store is None or self._store.cache.occupied == 0: |
| return hidden |
| assert self._index is not None |
| entity_anchor = None |
| if self._query_entity_anchor is not None: |
| entity_anchor = self._query_entity_anchor[:, None, :].expand( |
| hidden.shape[0], hidden.shape[1], -1 |
| ) |
| query = self.package.query_projector(hidden, entity_anchor=entity_anchor) |
| if self._collect: |
| self._query_telemetry.append(FactorizedCanonicalQuery( |
| entity=query.entity.detach(), |
| relation_logits=query.relation_logits.detach(), |
| metadata_logits=query.metadata_logits.detach(), |
| )) |
| route = ( |
| self._oracle_route(query, hidden) |
| if self._oracle_indices is not None |
| else self.router.route(query, self._index, top_k=self.package.config.top_k) |
| ) |
| if self._collect: |
| self._route_telemetry.append(RouteResult( |
| indices=route.indices.detach(), scores=route.scores.detach(), |
| weights=route.weights.detach(), features=route.features.detach(), |
| accepted=route.accepted.detach(), |
| has_valid=route.has_valid, |
| )) |
| if not self._injection_enabled or not route.has_valid: |
| return hidden |
| canonical = self._store.canonical_values.to( |
| device=hidden.device, dtype=route.weights.dtype |
| ) |
| selected = canonical[route.indices] |
| pooled = torch.einsum("...k,...kd->...d", route.weights, selected) |
| translated = self.package.value_translator(pooled) |
| route_features = torch.einsum( |
| "...k,...kf->...f", route.weights, route.features |
| ) |
| gate = ( |
| self.package.gate(hidden, translated, route_features) |
| if self._gate_enabled |
| else torch.ones(hidden.shape[:-1], device=hidden.device, dtype=translated.dtype) |
| ) |
| gate = gate * route.accepted.to(gate.dtype) |
| if self._collect: |
| self._gate_telemetry.append(gate.detach()) |
| return hidden + (gate.unsqueeze(-1) * translated).to(hidden.dtype) |
|
|
| def train(self, mode: bool = True): |
| super().train(mode) |
| self.base_model.eval() |
| self.package.train(mode) |
| self.router.train(mode) |
| return self |
|
|
| def forward( |
| self, |
| *args, |
| p_store: CanonicalPStore | None = None, |
| query_entity_surfaces: list[str] | tuple[str, ...] | None = None, |
| oracle_indices: Tensor | None = None, |
| gate_enabled: bool = True, |
| injection_enabled: bool = True, |
| collect_telemetry: bool = False, |
| **kwargs, |
| ): |
| if self._store is not None: |
| raise RuntimeError("PythiaSplitTranslatedModel is not reentrant") |
| self._store = p_store |
| self._oracle_indices = oracle_indices |
| if query_entity_surfaces is not None: |
| if "input_ids" in kwargs and len(query_entity_surfaces) != kwargs["input_ids"].shape[0]: |
| raise ValueError("query entity surface count must match the input batch") |
| self._query_entity_anchor = self.entity_encoder( |
| list(query_entity_surfaces) |
| ).to(next(self.package.parameters()).device) |
| else: |
| self._query_entity_anchor = None |
| self._gate_enabled = gate_enabled |
| self._injection_enabled = injection_enabled |
| self._collect = collect_telemetry |
| self._gate_telemetry.clear() |
| self._route_telemetry.clear() |
| self._query_telemetry.clear() |
| if p_store is not None and p_store.cache.occupied: |
| self._index = self.router.build_index( |
| p_store, self.entity_encoder, device=next(self.package.parameters()).device |
| ) |
| try: |
| return self.base_model(*args, **kwargs) |
| finally: |
| self._store = None |
| self._index = None |
| self._oracle_indices = None |
| self._query_entity_anchor = None |
| self._collect = False |
|
|
| @property |
| def gate_telemetry(self): |
| return tuple(self._gate_telemetry) |
|
|
| @property |
| def route_telemetry(self): |
| return tuple(self._route_telemetry) |
|
|
| @property |
| def query_telemetry(self): |
| return tuple(self._query_telemetry) |
|
|
| def close(self): |
| for handle in self._handles: |
| handle.remove() |
| self._handles.clear() |
|
|
