Premchan369/Q-TensorFormer

1"""2Tensor-Train decomposed linear layers with Nested-Core Dynamic Execution.3 4Architecture:5  Replaces dense linear matrix W with d TT-cores:6    W_{i1...ik, o1...ok} = G^(1) · G^(2) ··· G^(d)7 8Key Innovations:9  1. Nested-Core Dynamic Execution:10     Allocates cores up to max_rank (default 8). At runtime, dynamically11     slices cores to active rank r in {1, 2, 4, 8} with ZERO runtime SVD overhead!12  2. Memory Traffic Instrumentation:13     Tracks actual bytes read (weights + inputs) and bytes written (activations)14     per forward pass for true hardware bandwidth optimization.15  3. Batched Vectorized Contraction:16     Hardware-efficient einsum/bmm contraction path.17"""18 19import torch20import torch.nn as nn21import torch.nn.functional as F22import math23from typing import Tuple, Optional, Dict, List24 25 26def factorize_dim(dim: int, max_factors: int = 4) -> Tuple[int, ...]:27    """28    Factorize a dimension into factors >= 2 to avoid trivial padding cores.29    """30    if dim <= 1:31        return (1,)32    factors = []33    remaining = dim34    for p in [2, 2, 3, 2, 5, 2, 3, 7]:35        while remaining % p == 0 and len(factors) < max_factors - 1:36            factors.append(p)37            remaining //= p38        if remaining == 1:39            break40    if remaining > 1 and len(factors) < max_factors:41        factors.append(remaining)42    while len(factors) < 2:43        val = factors[0] if factors else dim44        root = int(math.isqrt(val))45        for d in range(root, 1, -1):46            if val % d == 0:47                factors = [d, val // d]48                break49        else:50            factors = [1, val]51    return tuple(factors[:max_factors])52 53 54class TTLinear(nn.Module):55    """56    Tensor-Train decomposed linear layer with nested rank slicing.57    """58 59    def __init__(60        self,61        in_features: int,62        out_features: int,63        max_rank: int = 8,64        bias: bool = True,65        rank: Optional[int] = None,66        **kwargs,67    ):68        super().__init__()69        if rank is not None:70            effective_max_rank = max(max_rank, rank, 8)71            init_active = rank72        else:73            effective_max_rank = max(max_rank, 8)74            init_active = max_rank75 76        self.in_features = in_features77        self.out_features = out_features78        self.max_rank = effective_max_rank79        self.active_rank = init_active80 81        # Factorize input and output dimensions82        in_factors = list(factorize_dim(in_features))83        out_factors = list(factorize_dim(out_features))84        self.ndim = max(len(in_factors), len(out_factors))85 86        while len(in_factors) < self.ndim:87            in_factors.append(1)88        while len(out_factors) < self.ndim:89            out_factors.append(1)90 91        self.in_shape = tuple(in_factors)92        self.out_shape = tuple(out_factors)93 94        # Allocate nested TT-cores: G^(k) of shape (r_k, out_k, in_k, r_{k+1})95        self.cores = nn.ParameterList()96        for k in range(self.ndim):97            r_left = 1 if k == 0 else max_rank98            r_right = 1 if k == self.ndim - 1 else max_rank99            core = torch.empty(r_left, out_factors[k], in_factors[k], r_right)100            fan = max(1, r_left * in_factors[k] + r_right * out_factors[k])101            bound = math.sqrt(6.0 / fan)102            nn.init.uniform_(core, -bound, bound)103            self.cores.append(core)104 105        self.bias = nn.Parameter(torch.zeros(out_features)) if bias else None106 107        # Statistics108        self.dense_params = in_features * out_features109        self._last_bytes_read = 0110        self._last_bytes_written = 0111 112    @property113    def rank(self) -> int:114        return self.active_rank115 116    @rank.setter117    def rank(self, value: int):118        self.set_rank(value)119 120    def set_rank(self, rank: int):121        """Set active rank for nested core slicing via zero-copy tensor stride extraction."""122        self.active_rank = max(1, min(self.max_rank, int(rank)))123 124    def _slice_cores_for_rank(self, r: int) -> List[torch.Tensor]:125        """Slice nested cores according to specified rank."""126        sliced_cores = []127        for k, core in enumerate(self.cores):128            r_left = 1 if k == 0 else min(r, core.shape[0])129            r_right = 1 if k == self.ndim - 1 else min(r, core.shape[3])130            sliced_cores.append(core[:r_left, :, :, :r_right])131        return sliced_cores132 133    def get_active_cores(self) -> List[torch.Tensor]:134        """Slice nested cores according to current active rank."""135        return self._slice_cores_for_rank(self.active_rank)136 137    def measure_slicing_overhead_us(self, n_runs: int = 1000) -> float:138        """Empirically measure microsecond overhead of core slice extraction."""139        import time140        t0 = time.perf_counter()141        for _ in range(n_runs):142            _ = self._slice_cores_for_rank(4)143        return ((time.perf_counter() - t0) / n_runs) * 1e6144 145    @property146    def active_param_count(self) -> int:147        """Parameters actively used under current rank."""148        cores = self.get_active_cores()149        total = sum(c.numel() for c in cores)150        if self.bias is not None:151            total += self.bias.numel()152        return total153 154    @property155    def compression_ratio(self) -> float:156        """Effective compression ratio vs dense linear layer."""157        return self.dense_params / max(1, self.active_param_count)158 159    def _contract_with_cores(self, x_flat: torch.Tensor, cores: List[torch.Tensor]) -> torch.Tensor:160        """Contract 2D input (N, in_features) through given sliced cores."""161        N = x_flat.shape[0]162        state = x_flat.reshape(N, *self.in_shape)163        for k in range(self.ndim):164            core = cores[k]165            r_k, o_k, i_k, r_kp1 = core.shape166            if k == 0:167                rest = math.prod(self.in_shape[1:]) if self.ndim > 1 else 1168                s = state.reshape(N, i_k, rest)169                cm = core.squeeze(0).permute(1, 0, 2).reshape(i_k, o_k * r_kp1)170                s = torch.bmm(s.transpose(1, 2), cm.unsqueeze(0).expand(N, -1, -1))171                s = s.reshape(N, rest, o_k, r_kp1).permute(0, 3, 2, 1)172                state = s.reshape(N, r_kp1, -1)173            elif k == self.ndim - 1:174                prev_os = math.prod(self.out_shape[:k]) if k > 0 else 1175                s = state.reshape(N, r_k, prev_os, i_k)176                cm = core.squeeze(-1)177                s = torch.einsum("brpi,roi->bpo", s, cm)178                state = s.reshape(N, prev_os * o_k)179            else:180                prev_os = math.prod(self.out_shape[:k]) if k > 0 else 1181                rest_in = math.prod(self.in_shape[k + 1:])182                s = state.reshape(N, r_k, prev_os, i_k, rest_in)183                s = torch.einsum("brpix,roiq->bpoqx", s, core)184                s = s.permute(0, 3, 1, 2, 4)185                state = s.reshape(N, r_kp1, prev_os * o_k * rest_in)186        out = state.reshape(N, self.out_features)187        if self.bias is not None:188            out = out + self.bias189        return out190 191    def to_matrix(self, rank: Optional[int] = None) -> torch.Tensor:192        """Reconstruct the effective dense weight matrix (out_features, in_features) for rank r."""193        r = self.active_rank if rank is None else rank194        cores = self._slice_cores_for_rank(r)195        device = self.cores[0].device196        dtype = self.cores[0].dtype197        eye = torch.eye(self.in_features, device=device, dtype=dtype)198        bias_backup = self.bias199        self.bias = None200        try:201            wt = self._contract_with_cores(eye, cores)202        finally:203            self.bias = bias_backup204        return wt.t()205 206 207    def forward(self, x: torch.Tensor, token_ranks: Optional[torch.Tensor] = None) -> torch.Tensor:208        """209        Forward pass with nested core contraction. Supports token-grouped dynamic rank execution.210        """211        batch_shape = x.shape[:-1]212        B = math.prod(batch_shape) if batch_shape else 1213        x_flat = x.reshape(B, self.in_features)214 215        if token_ranks is not None and token_ranks.numel() == B:216            flat_ranks = token_ranks.reshape(-1)217            out = torch.empty(B, self.out_features, device=x.device, dtype=x.dtype)218            total_read = 0219            for r in flat_ranks.unique():220                r_int = int(r.item())221                idx = (flat_ranks == r).nonzero(as_tuple=True)[0]222                cores_r = self._slice_cores_for_rank(r_int)223                out[idx] = self._contract_with_cores(x_flat[idx], cores_r)224                total_read += sum(c.numel() * c.element_size() for c in cores_r) + x_flat[idx].numel() * x_flat.element_size()225            self._last_bytes_read = total_read226            self._last_bytes_written = out.numel() * out.element_size()227            return out.reshape(*batch_shape, self.out_features)228        else:229            cores = self.get_active_cores()230            self._last_bytes_read = sum(c.numel() * c.element_size() for c in cores) + x_flat.numel() * x_flat.element_size()231            out = self._contract_with_cores(x_flat, cores)232            self._last_bytes_written = out.numel() * out.element_size()233            return out.reshape(*batch_shape, self.out_features)234 235    def flops(self, batch_size: int = 1, seq_len: int = 1) -> int:236        """Compute active FLOPs for current rank."""237        r = self.active_rank238        avg_dim = (sum(self.in_shape) + sum(self.out_shape)) / (2 * self.ndim)239        return int(2 * r**2 * self.ndim * avg_dim * batch_size * seq_len)240 241    def get_memory_traffic(self) -> Dict[str, int]:242        """Return bytes read and written in last forward pass."""243        return {244            "bytes_read": self._last_bytes_read,245            "bytes_written": self._last_bytes_written,246            "total_bytes": self._last_bytes_read + self._last_bytes_written,247        }248 249 250class TTFeedForward(nn.Module):251    """252    Tensor-Train Feed-Forward Network with nested dynamic ranks.253    """254 255    def __init__(256        self,257        hidden_dim: int,258        ff_multiplier: int = 4,259        rank: int = 8,260        activation=F.gelu,261    ):262        super().__init__()263        self.hidden_dim = hidden_dim264        self.expanded_dim = hidden_dim * ff_multiplier265        self.max_rank = rank266 267        self.up_proj = TTLinear(hidden_dim, self.expanded_dim, max_rank=rank, bias=True)268        self.down_proj = TTLinear(self.expanded_dim, hidden_dim, max_rank=rank, bias=True)269        self.activation = activation270 271    def forward(self, x: torch.Tensor, token_ranks: Optional[torch.Tensor] = None) -> torch.Tensor:272        return self.down_proj(self.activation(self.up_proj(x, token_ranks=token_ranks)), token_ranks=token_ranks)273 274    def set_rank(self, rank: int):275        self.up_proj.set_rank(rank)276        self.down_proj.set_rank(rank)277 278    @property279    def rank(self) -> int:280        return self.up_proj.active_rank281 282    @property283    def total_params(self) -> int:284        return sum(p.numel() for p in self.parameters())285 286    @property287    def active_params(self) -> int:288        return self.up_proj.active_param_count + self.down_proj.active_param_count289 290    def flops(self, batch_size: int = 1, seq_len: int = 1) -> int:291        return self.up_proj.flops(batch_size, seq_len) + self.down_proj.flops(batch_size, seq_len)292 293    def get_memory_traffic(self) -> Dict[str, int]:294        up = self.up_proj.get_memory_traffic()295        down = self.down_proj.get_memory_traffic()296        return {297            "bytes_read": up["bytes_read"] + down["bytes_read"],298            "bytes_written": up["bytes_written"] + down["bytes_written"],299            "total_bytes": up["total_bytes"] + down["total_bytes"],300        }301