Premchan369/Q-TensorFormer

1"""2Q-TensorFormer: Complete Model Architecture.3 4Unified Information-to-Resource Allocation Architecture.5Integrates:6  - Nested TT-FFN (ranks 1, 2, 4, 8 with instantaneous slicing)7  - Token Information State (8D feature tracking)8  - Information-Value Resource Allocator (marginal utility optimization)9  - Hardware-Aware Cost Profiling10  - Adaptive KV Cache with multi-precision & eviction11  - Deployment Presets:12      QTF_FULL, QTF_BALANCED, QTF_LATENCY, QTF_MEMORY,13      QTF_ENERGY, QTF_EDGE, QTF_CLASSICAL_ONLY14"""15 16import torch17import torch.nn as nn18import torch.nn.functional as F19import math20from typing import Optional, Dict, List, Tuple, Union21 22from .blocks import HybridBlock23from .config import ModelConfig24from .kv_cache import AdaptiveKVCache25from .resource_allocator import AllocationBudget26from .hardware_cost_model import HardwareCostModel27 28 29class PositionalEncoding(nn.Module):30    def __init__(self, d_model: int, max_len: int = 512, dropout: float = 0.1):31        super().__init__()32        self.dropout = nn.Dropout(dropout)33        pe = torch.zeros(max_len, d_model)34        pos = torch.arange(0, max_len).float().unsqueeze(1)35        div = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))36        pe[:, 0::2] = torch.sin(pos * div)37        pe[:, 1::2] = torch.cos(pos * div)38        self.register_buffer("pe", pe.unsqueeze(0))39 40    def forward(self, x: torch.Tensor, start_pos: int = 0) -> torch.Tensor:41        seq_len = x.size(1)42        return self.dropout(x + self.pe[:, start_pos:start_pos + seq_len, :])43 44 45class QTensorFormer(nn.Module):46    """47    Q-TensorFormer: Closed-Loop Information-to-Resource Adaptive Transformer.48    """49 50    DEPLOYMENT_PRESETS = [51        "QTF_FULL",52        "QTF_BALANCED",53        "QTF_LATENCY",54        "QTF_MEMORY",55        "QTF_ENERGY",56        "QTF_EDGE",57        "QTF_CLASSICAL_ONLY",58    ]59 60    def __init__(self, config: ModelConfig, preset: str = "QTF_BALANCED"):61        super().__init__()62        self.config = config63        self.current_preset = preset.upper()64 65        self.embedding = nn.Embedding(config.vocab_size, config.d_model)66        self.pos_encoding = PositionalEncoding(config.d_model, config.max_seq_len, config.dropout)67 68        # Transformer blocks69        backend_type = getattr(config, "backend_type", "classical_surrogate")70        self.blocks = nn.ModuleList([71            HybridBlock(72                d_model=config.d_model,73                n_heads=config.n_heads,74                n_kv_heads=getattr(config, "n_kv_heads", None),75                ff_multiplier=config.ff_multiplier,76                tt_rank=config.tt_rank,77                tt_min_rank=config.tt_min_rank,78                use_quantum=config.use_quantum,79                n_qubits=config.n_qubits,80                backend_type=backend_type,81                dropout=config.dropout,82                max_seq_len=config.max_seq_len,83                hysteresis_tau=getattr(config, "hysteresis_tau", 0.15),84                default_preset=self.current_preset.replace("QTF_", "").lower(),85                enable_early_exit=getattr(config, "enable_early_exit", False),86                early_exit_threshold=getattr(config, "early_exit_threshold", 0.20),87                vocab_size=config.vocab_size,88            )89            for _ in range(config.n_layers)90        ])91 92        self.ln_f = nn.LayerNorm(config.d_model)93        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)94        self.lm_head.weight = self.embedding.weight  # Weight tying95 96        # Hardware cost model97        self.hardware_model = HardwareCostModel()98 99        self._init_weights()100 101    def _init_weights(self):102        for name, p in self.named_parameters():103            if "weight" in name and p.dim() >= 2:104                nn.init.xavier_uniform_(p)105            elif "bias" in name:106                nn.init.zeros_(p)107 108    def set_preset(self, preset_name: str):109        """Switch deployment preset."""110        preset = preset_name.upper()111        if preset not in self.DEPLOYMENT_PRESETS:112            raise ValueError(f"Unknown preset {preset}. Choose from {self.DEPLOYMENT_PRESETS}")113        self.current_preset = preset114 115    def forward(116        self,117        input_ids: torch.Tensor,118        attention_mask: Optional[torch.Tensor] = None,119        kv_caches: Optional[List[AdaptiveKVCache]] = None,120        budget: Optional[AllocationBudget] = None,121        return_stats: bool = False,122        preset: Optional[str] = None,123    ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[Dict]]]:124        """125        Forward pass with dynamic information-to-resource allocation.126        """127        if preset is not None:128            self.set_preset(preset)129 130        B, T = input_ids.shape131        start_pos = kv_caches[0].seq_len if (kv_caches is not None and kv_caches[0].seq_len > 0) else 0132 133        x = self.embedding(input_ids)134        x = self.pos_encoding(x, start_pos=start_pos)135 136        preset_str = self.current_preset.replace("QTF_", "").lower()137        force_classical = (self.current_preset in ["QTF_CLASSICAL_ONLY", "QTF_EDGE"] or not self.config.use_quantum)138 139        all_stats = []140        for i, block in enumerate(self.blocks):141            layer_kv = kv_caches[i] if kv_caches is not None else None142            x, stats = block(143                x,144                mask=attention_mask,145                kv_cache=layer_kv,146                budget=budget,147                preset=preset_str,148                force_classical=force_classical,149            )150            all_stats.append(stats)151            if stats.get("early_exit_triggered", False) and "early_exit_logits" in stats:152                logits = stats["early_exit_logits"]153                if return_stats:154                    return logits, all_stats155                return logits156 157        x = self.ln_f(x)158        logits = self.lm_head(x)159 160        if return_stats:161            return logits, all_stats162        return logits163 164    @torch.no_grad()165    def generate(166        self,167        input_ids: torch.Tensor,168        max_new_tokens: int = 20,169        temperature: float = 1.0,170        top_k: int = 50,171        use_cache: bool = True,172    ) -> torch.Tensor:173        """Autoregressive generation with KV Cache."""174        self.eval()175        B = input_ids.shape[0]176 177        kv_caches = None178        if use_cache:179            kv_caches = [180                AdaptiveKVCache(max_capacity=self.config.max_seq_len)181                for _ in range(self.config.n_layers)182            ]183 184        # Prefill phase185        curr_input = input_ids186        for _ in range(max_new_tokens):187            if kv_caches is not None and kv_caches[0].seq_len > 0:188                step_input = curr_input[:, -1:]189            else:190                step_input = curr_input191 192            logits = self(step_input, kv_caches=kv_caches)193            next_logits = logits[:, -1, :] / max(1e-5, temperature)194 195            if top_k > 0:196                v, _ = torch.topk(next_logits, min(top_k, next_logits.size(-1)))197                next_logits[next_logits < v[:, [-1]]] = float("-inf")198 199            probs = F.softmax(next_logits, dim=-1)200            next_token = torch.multinomial(probs, num_samples=1)201            curr_input = torch.cat([curr_input, next_token], dim=-1)202 203        return curr_input204 205    def reset_schedulers(self):206        for block in self.blocks:207            block.reset_scheduler()208 209    @property210    def total_params(self) -> int:211        return sum(p.numel() for p in self.parameters())212 213    @property214    def active_params(self) -> int:215        tt_active = sum(b.tt_ffn.active_params for b in self.blocks)216        base = self.total_params - sum(b.tt_ffn.total_params for b in self.blocks)217        return base + tt_active218 219    @property220    def compression_ratio(self) -> float:221        dense_per_block = 2 * self.config.d_model * self.config.d_model * self.config.ff_multiplier222        base = self.total_params - sum(b.tt_ffn.total_params for b in self.blocks)223        active_tt = sum(b.tt_ffn.active_params for b in self.blocks)224        return (base + dense_per_block * self.config.n_layers) / max(base + active_tt, 1)225 226    def flops_estimate(self, batch_size: int = 1, seq_len: int = 32) -> Dict:227        total = 0228        breakdown = {}229        for i, block in enumerate(self.blocks):230            b = block.flops_estimate(batch_size, seq_len)231            total += b["total"]232            breakdown[f"block_{i}"] = b233        return {"total": total, "breakdown": breakdown}234 235 236class DenseBaseline(nn.Module):237    """238    Standard dense transformer baseline with identical hyperparameters.239    """240 241    def __init__(self, config: ModelConfig):242        super().__init__()243        self.config = config244 245        self.embedding = nn.Embedding(config.vocab_size, config.d_model)246        self.pos_encoding = PositionalEncoding(config.d_model, config.max_seq_len, config.dropout)247 248        self.blocks = nn.ModuleList([249            nn.ModuleDict({250                "ln1": nn.LayerNorm(config.d_model),251                "attn": nn.MultiheadAttention(config.d_model, config.n_heads, dropout=config.dropout, batch_first=True),252                "ln2": nn.LayerNorm(config.d_model),253                "ffn": nn.Sequential(254                    nn.Linear(config.d_model, config.d_model * config.ff_multiplier),255                    nn.GELU(),256                    nn.Linear(config.d_model * config.ff_multiplier, config.d_model),257                ),258                "dropout": nn.Dropout(config.dropout),259            })260            for _ in range(config.n_layers)261        ])262 263        self.ln_f = nn.LayerNorm(config.d_model)264        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)265        self.lm_head.weight = self.embedding.weight266        self._init_weights()267 268    def _init_weights(self):269        for name, p in self.named_parameters():270            if "weight" in name and p.dim() >= 2:271                nn.init.xavier_uniform_(p)272            elif "bias" in name:273                nn.init.zeros_(p)274 275    def forward(self, input_ids, attention_mask=None, return_stats=False):276        x = self.embedding(input_ids)277        x = self.pos_encoding(x)278 279        for block in self.blocks:280            normed = block["ln1"](x)281            attn_out, _ = block["attn"](normed, normed, normed, need_weights=False)282            x = x + block["dropout"](attn_out)283            ffn_out = block["ffn"](block["ln2"](x))284            x = x + block["dropout"](ffn_out)285 286        x = self.ln_f(x)287        logits = self.lm_head(x)288        if return_stats:289            return logits, []290        return logits291 292    @property293    def total_params(self) -> int:294        return sum(p.numel() for p in self.parameters())295 296 297def create_model(config: ModelConfig, model_type: str = "qtensor", preset: str = "QTF_BALANCED"):298    """Factory helper to build models."""299    if model_type in ["qtensor", "hybrid"]:300        return QTensorFormer(config, preset=preset)301    elif model_type in ["dense", "baseline"]:302        return DenseBaseline(config)303    elif model_type == "tensor_only":304        cfg = ModelConfig(**config.__dict__)305        cfg.use_quantum = False306        return QTensorFormer(cfg, preset="QTF_CLASSICAL_ONLY")307    else:308        raise ValueError(f"Unknown model_type: {model_type}")309