Upload q_tensor_former.py

q_tensor_former.py

@@ -0,0 +1,493 @@
+#!/usr/bin/env python3
+"""
+Q-TensorFormer: Quantum-Enhanced Tensor Network LLM Compression Engine
+=======================================================================
+Hybrid quantum-tensor transformer with:
+  - Pure PyTorch Tensor-Train FFN layers (no compiled deps)
+  - PennyLane quantum angle encoding with TorchLayer
+  - Entanglement-guided adaptive rank scheduling
+  - Selective quantum routing (only "hard" tokens)
+  - Full benchmark against identical-architecture baseline
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math, os
+from typing import Optional, Tuple
+from dataclasses import dataclass
+
+import pennylane as qml
+
+print("=" * 65)
+print(" Q-TENSORFORMER: Quantum-Tensor LLM Compressor")
+print("=" * 65)
+print(f" PyTorch {torch.__version__}  |  PennyLane {qml.__version__}")
+print()
+
+# ═════════════════════════════════════════════════════════════════════
+# CONFIG
+# ═════════════════════════════════════════════════════════════════════
+
+@dataclass
+class CFG:
+    d_model: int = 64
+    n_heads: int = 4
+    n_layers: int = 2
+    ff_mult: int = 4
+    max_seq: int = 64
+    vocab: int = 1000
+    tt_rank: int = 8
+    min_rank: int = 2
+    q_qubits: int = 4
+    q_layers: int = 2
+    q_sparsity: float = 0.3
+    dropout: float = 0.1
+    lr: float = 3e-4
+    rank_alpha: float = 2.0
+    rank_smoothing: float = 0.9
+
+# ═════════════════════════════════════════════════════════════════════
+# 1. PURE PYTORCH TENSOR-TRAIN LINEAR LAYER
+# ═════════════════════════════════════════════════════════════════════
+
+def auto_factor(n, max_f=4):
+    if n <= 1: return (1, 1)
+    f, r = [], n
+    for p in [2,2,2,2,2,3,3,5,7]:
+        while r % p == 0 and len(f) < max_f:
+            f.append(p); r //= p
+    if r > 1:
+        if len(f) < max_f: f.append(r)
+        else: f[-1] *= r
+    while len(f) < 2: f.insert(0, 1)
+    return tuple(f[:max_f])
+
+class TTLinear(nn.Module):
+    """Tensor-Train decomposed linear layer. Pure PyTorch, zero compiled deps."""
+    def __init__(self, in_shape, out_shape, rank=8, bias=True):
+        super().__init__()
+        in_shape = tuple(in_shape)
+        out_shape = tuple(out_shape)
+        max_d = max(len(in_shape), len(out_shape))
+        in_shape = (1,) * (max_d - len(in_shape)) + in_shape
+        out_shape = (1,) * (max_d - len(out_shape)) + out_shape
+        assert len(in_shape) == len(out_shape)
+        self.in_shape, self.out_shape = in_shape, out_shape
+        self.rank, self.ndim = rank, len(in_shape)
+        self.in_feat = math.prod(in_shape)
+        self.out_feat = math.prod(out_shape)
+        self.cores = nn.ParameterList()
+        for k in range(self.ndim):
+            rl = 1 if k == 0 else rank
+            rr = 1 if k == self.ndim - 1 else rank
+            c = torch.empty(rl, out_shape[k], in_shape[k], rr)
+            bnd = math.sqrt(6.0 / max(1, rl*in_shape[k] + rr*out_shape[k]))
+            nn.init.uniform_(c, -bnd, bnd)
+            self.cores.append(c)
+        self.bias = nn.Parameter(torch.zeros(self.out_feat)) if bias else None
+        tp = sum(c.numel() for c in self.cores) + (self.bias.numel() if bias else 0)
+        self.compr = (self.in_feat * self.out_feat) / max(tp, 1)
+
+    def forward(self, x):
+        bs = x.shape[:-1]
+        B = math.prod(bs)
+        x = x.reshape(B, self.in_feat)
+        state = x.reshape(B, *self.in_shape)
+
+        for k in range(self.ndim):
+            core = self.cores[k]
+            r_k, o_k, i_k, r_kp1 = core.shape
+
+            if k == 0:
+                rest = math.prod(self.in_shape[1:])
+                s = state.reshape(B, i_k, rest)
+                cm = core.squeeze(0).permute(1, 0, 2).reshape(i_k, o_k * r_kp1)
+                s = torch.bmm(s.transpose(1, 2), cm.unsqueeze(0).expand(B, -1, -1))
+                s = s.reshape(B, rest, o_k, r_kp1).permute(0, 3, 2, 1)
+                state = s.reshape(B, r_kp1, -1)
+
+            elif k == self.ndim - 1:
+                prev_os = math.prod(self.out_shape[:k])
+                s = state.reshape(B, r_k, prev_os, i_k)
+                cm = core.squeeze(-1)
+                s = torch.einsum('brpi,roi->bpo', s, cm)
+                state = s.reshape(B, prev_os * o_k)
+
+            else:
+                prev_os = math.prod(self.out_shape[:k])
+                rest_in = math.prod(self.in_shape[k+1:])
+                s = state.reshape(B, r_k, prev_os * i_k * rest_in)
+                s = s.reshape(B, r_k, prev_os, i_k, rest_in)
+                s = torch.einsum('brpix,roiq->bpoqx', s, core)
+                s = s.permute(0, 3, 1, 2, 4)
+                state = s.reshape(B, r_kp1, prev_os * o_k * rest_in)
+
+        out = state.reshape(B, self.out_feat)
+        if self.bias is not None: out = out + self.bias
+        return out.reshape(*bs, self.out_feat)
+
+    def set_rank(self, nr):
+        for i, c in enumerate(self.cores):
+            s = [slice(None)]*4
+            if i > 0: s[0] = slice(None, nr)
+            if i < self.ndim - 1: s[3] = slice(None, nr)
+            self.cores[i] = nn.Parameter(c[tuple(s)].clone())
+
+# ═════════════════════════════════════════════════════════════════════
+# 2. QUANTUM ANGLE EMBEDDING (PennyLane)
+# ═════════════════════════════════════════════════════════════════════
+
+class QuantumEmbed(nn.Module):
+    """Angle embedding → variational circuit → PauliZ expectations."""
+    def __init__(self, n_q=4, layers=2, n_out=None):
+        super().__init__()
+        self.n_q, self.layers = n_q, layers
+        n_out = n_out or n_q
+        dev = qml.device("default.qubit", wires=n_q)
+
+        @qml.qnode(dev, interface="torch", diff_method="backprop")
+        def circ(inputs, w):
+            for i in range(n_q): qml.RX(inputs[..., i], wires=i)
+            for L in range(layers):
+                for i in range(n_q): qml.RY(w[L, i], wires=i)
+                for i in range(n_q-1): qml.CNOT(wires=[i, i+1])
+                if n_q > 2: qml.CNOT(wires=[n_q-1, 0])
+            return [qml.expval(qml.PauliZ(i)) for i in range(n_out)]
+
+        self.qlayer = qml.qnn.TorchLayer(circ, {"w": (layers, n_q)})
+
+    def forward(self, x): return self.qlayer(x)
+
+# ═════════════════════════════════════════════════════════════════════
+# 3. TT FEED-FORWARD
+# ═════════════════════════════════════════════════════════════════════
+
+class TTFFN(nn.Module):
+    def __init__(self, D, ff_mult=4, rank=8):
+        super().__init__()
+        E = D * ff_mult
+        self.up = TTLinear(auto_factor(D), auto_factor(E), rank, True)
+        self.down = TTLinear(auto_factor(E), auto_factor(D), rank, True)
+    def forward(self, x): return self.down(F.gelu(self.up(x)))
+    def set_rank(self, r): self.up.set_rank(r); self.down.set_rank(r)
+
+# ═════════════════════════════════════════════════════════════════════
+# 4. RANK SCHEDULER
+# ═════════════════════════════════════════════════════════════════════
+
+class RankScheduler(nn.Module):
+    """rank = r_min + alpha * entropy (EMA-smoothed)"""
+    def __init__(self, mn=2, mx=16, a=2.0, sm=0.9):
+        super().__init__()
+        self.mn, self.mx = mn, mx
+        self.alpha = nn.Parameter(torch.tensor(a))
+        self.sm = sm
+        self.register_buffer('ema', torch.tensor(0.5))
+        self.register_buffer('cur', torch.tensor(float(mx)))
+    def forward(self, ent):
+        s = ent.mean().detach() if ent.numel()>1 else ent.detach()
+        self.ema = self.sm*self.ema + (1-self.sm)*s
+        raw = self.mn + self.alpha*self.ema
+        r = int(torch.clamp(raw, self.mn, self.mx).round().item())
+        if self.training: self.cur.fill_(r)
+        return r
+    @property
+    def current(self): return int(self.cur.item())
+
+# ═════════════════════════════════════════════════════════════════════
+# 5. QUANTUM ROUTER
+# ═════════════════════════════════════════════════════════════════════
+
+class QuantumRouter(nn.Module):
+    """Learned gate: routes only hard tokens through quantum circuit."""
+    def __init__(self, D, qmod, thr=0.5):
+        super().__init__()
+        self.qmod = qmod
+        self.thr = thr
+        self.gate = nn.Sequential(
+            nn.Linear(D, D//4), nn.ReLU(), nn.Linear(D//4,1), nn.Sigmoid())
+        self.register_buffer('tot', torch.tensor(0.0))
+        self.register_buffer('qtok', torch.tensor(0.0))
+    def forward(self, x):
+        B,S,D = x.shape
+        g = self.gate(x.reshape(-1,D)).squeeze(-1).reshape(B,S)
+        m = (g > self.thr).float()
+        if self.training:
+            m = m.detach() + g - g.detach()
+        xf = x.reshape(-1,D); mf = m.reshape(-1)
+        sel = xf[mf > 0.5]; out = xf.clone()
+        if sel.shape[0]>0:
+            qo = self.qmod(sel)
+            if qo.shape[-1]!=D:
+                if not hasattr(self,'_proj'):
+                    self._proj = nn.Linear(qo.shape[-1],D).to(x.device)
+                qo = self._proj(qo)
+            out[mf > 0.5] = qo.to(out.dtype)
+        self.tot += B*S; self.qtok += m.sum()
+        return out.reshape(B,S,D), g
+    def sparsity(self):
+        if self.tot>0: return 1.0-(self.qtok/self.tot).item()
+        return 1.0
+
+# ═════════════════════════════════════════════════════════════════════
+# 6. ATTENTION
+# ═════════════════════════════════════════════════════════════════════
+
+class MHA(nn.Module):
+    def __init__(self, D, heads=4, drop=0.1):
+        super().__init__()
+        assert D%heads==0
+        self.h, self.hd = heads, D//heads
+        self.scale = self.hd**-0.5
+        self.qkv = nn.Linear(D, 3*D, bias=False)
+        self.out = nn.Linear(D, D)
+        self.drop = nn.Dropout(drop)
+    def forward(self, x, mask=None):
+        B,S,D = x.shape
+        qkv = self.qkv(x).reshape(B,S,3,self.h,self.hd).permute(2,0,3,1,4)
+        q,k,v = qkv[0], qkv[1], qkv[2]
+        a = (q@k.transpose(-2,-1))*self.scale
+        if mask is not None:
+            a = a.masked_fill(mask[:,None,None,:]==0, float('-inf'))
+        aw = F.softmax(a, dim=-1); aw = self.drop(aw)
+        o = (aw@v).transpose(1,2).reshape(B,S,D)
+        return self.out(o), aw
+
+# ═════════════════════════════════════════════════════════════════════
+# 7. HYBRID BLOCK
+# ═════════════════════════════════════════════════════════════════════
+
+class HybridBlock(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        D = cfg.d_model
+        self.a_norm = nn.LayerNorm(D)
+        self.attn = MHA(D, cfg.n_heads, cfg.dropout)
+        self.f_norm = nn.LayerNorm(D)
+        self.ffn = TTFFN(D, cfg.ff_mult, cfg.tt_rank)
+        self.qrouter = None
+        if cfg.q_qubits:
+            qc = QuantumEmbed(cfg.q_qubits, cfg.q_layers, cfg.q_qubits)
+            qw = nn.Sequential(nn.Linear(D, cfg.q_qubits), qc)
+            self.qrouter = QuantumRouter(D, qw)
+        self.rs = RankScheduler(cfg.min_rank, cfg.tt_rank, cfg.rank_alpha, cfg.rank_smoothing)
+        self.drop = nn.Dropout(cfg.dropout)
+    def forward(self, x, mask=None, adapt=True):
+        ao, aw = self.attn(self.a_norm(x), mask)
+        x = x + self.drop(ao)
+        eps=1e-8
+        ent = -torch.sum(aw*torch.log(aw+eps), dim=-1).mean(dim=-1).mean()
+        tr = self.rs(ent) if adapt else self.rs.mx
+        if adapt: self.ffn.set_rank(tr)
+        n = self.f_norm(x)
+        qs = 1.0
+        if self.qrouter is not None:
+            qo, _ = self.qrouter(n)
+            n = n + self.drop(qo - n.detach() + n)
+            qs = self.qrouter.sparsity()
+        x = x + self.drop(self.ffn(n))
+        return {'out':x, 'aw':aw, 'entropy':ent, 'rank':tr, 'qsparse':qs}
+
+# ═════════════════════════════════════════════════════════════════════
+# 8. Q-TENSORFORMER MODEL
+# ═════════════════════════════════════════════════════════════════════
+
+class QTensorFormer(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.tok = nn.Embedding(cfg.vocab, cfg.d_model)
+        self.pos = nn.Parameter(torch.randn(1, cfg.max_seq, cfg.d_model)*0.02)
+        self.layers = nn.ModuleList([HybridBlock(cfg) for _ in range(cfg.n_layers)])
+        self.norm = nn.LayerNorm(cfg.d_model)
+        self.head = nn.Linear(cfg.d_model, cfg.vocab, bias=False)
+        self.head.weight = self.tok.weight
+        self._init()
+    def _init(self):
+        for p in self.parameters():
+            if p.dim()>=2: nn.init.xavier_uniform_(p)
+    def forward(self, ids, mask=None, adapt=True):
+        B,S = ids.shape
+        x = self.tok(ids) + self.pos[:,:S,:]
+        if mask is not None: mask = mask[:,None,None,:]
+        bos = []
+        for l in self.layers:
+            o = l(x, mask, adapt); x=o['out']; bos.append(o)
+        x = self.norm(x); logits = self.head(x)
+        ent = torch.stack([b['entropy'] for b in bos]).mean()
+        rk = sum(b['rank'] for b in bos)/len(bos)
+        qs = sum(b['qsparse'] for b in bos)/len(bos)
+        return {'logits':logits,'entropy':ent,'rank':rk,'qsparse':qs}
+    def loss(self, ids, mask=None, labels=None):
+        if labels is None: labels=ids.clone()
+        out = self(ids, mask)
+        sl = out['logits'][:,:-1].contiguous()
+        ll = labels[:,1:].contiguous()
+        l = F.cross_entropy(sl.reshape(-1,self.cfg.vocab), ll.reshape(-1), ignore_index=-100)
+        return {'loss':l,'ppl':torch.exp(l),'entropy':out['entropy'],'rank':out['rank'],'qsparse':out['qsparse']}
+    def nparams(self):
+        t = sum(p.numel() for p in self.parameters())
+        tr = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        return {'total':t,'trainable':tr}
+
+# ═════════════════════════════════════════════════════════════════════
+# 9. BASELINE (identical architecture, dense FFN)
+# ═════════════════════════════════════════════════════════════════════
+
+class Baseline(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.tok = nn.Embedding(cfg.vocab, cfg.d_model)
+        self.pos = nn.Parameter(torch.randn(1, cfg.max_seq, cfg.d_model)*0.02)
+        self.drop = nn.Dropout(cfg.dropout)
+        self.layers = nn.ModuleList()
+        for _ in range(cfg.n_layers):
+            self.layers.append(nn.ModuleDict({
+                'a_n': nn.LayerNorm(cfg.d_model),
+                'a': MHA(cfg.d_model, cfg.n_heads, cfg.dropout),
+                'f_n': nn.LayerNorm(cfg.d_model),
+                'ff': nn.Sequential(
+                    nn.Linear(cfg.d_model, cfg.d_model*cfg.ff_mult),
+                    nn.GELU(), nn.Dropout(cfg.dropout),
+                    nn.Linear(cfg.d_model*cfg.ff_mult, cfg.d_model)),
+            }))
+        self.norm = nn.LayerNorm(cfg.d_model)
+        self.head = nn.Linear(cfg.d_model, cfg.vocab, bias=False)
+        self.head.weight = self.tok.weight
+        self._init()
+    def _init(self):
+        for p in self.parameters():
+            if p.dim()>=2: nn.init.xavier_uniform_(p)
+    def forward(self, ids, mask=None):
+        B,S = ids.shape
+        x = self.tok(ids)+self.pos[:,:S,:]; x=self.drop(x)
+        m = mask[:,None,None,:] if mask is not None else None
+        for l in self.layers:
+            ao,_ = l['a'](l['a_n'](x),m); x=x+self.drop(ao)
+            x = x+self.drop(l['ff'](l['f_n'](x)))
+        return {'logits':self.head(self.norm(x))}
+    def loss(self, ids, mask=None, labels=None):
+        if labels is None: labels=ids.clone()
+        out = self(ids, mask)
+        sl = out['logits'][:,:-1].contiguous()
+        ll = labels[:,1:].contiguous()
+        l = F.cross_entropy(sl.reshape(-1,self.cfg.vocab), ll.reshape(-1), ignore_index=-100)
+        return {'loss':l,'ppl':torch.exp(l)}
+    def nparams(self):
+        t = sum(p.numel() for p in self.parameters())
+        tr = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        return {'total':t,'trainable':tr}
+
+# ═════════════════════════════════════════════════════════════════════
+# 10. TRAINING UTILITIES
+# ═════════════════════════════════════════════════════════════════════
+
+def make_data(vocab=1000, seq=64, n=500, bs=16):
+    d = torch.randint(1, vocab, (n, seq))
+    ds = torch.utils.data.TensorDataset(d)
+    return torch.utils.data.DataLoader(ds, batch_size=bs, shuffle=True,
+        collate_fn=lambda batch: {'input_ids': torch.stack([item[0] for item in batch])})
+
+def train_epoch(model, dl, opt, sched, e, tag="M"):
+    model.train(); tl,tp,nb = 0.0,0.0,0; ex={}
+    for b in dl:
+        ids = b['input_ids']; m = b.get('attention_mask')
+        opt.zero_grad()
+        out = model.loss(ids, m); out['loss'].backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        opt.step()
+        if sched: sched.step()
+        tl += out['loss'].item(); tp += out['ppl'].item(); nb += 1
+        for k in ['entropy','rank','qsparse']:
+            if k in out: ex[k]=ex.get(k,0.0)+(out[k].item() if isinstance(out[k],torch.Tensor) else out[k])
+    al,ap = tl/max(nb,1), tp/max(nb,1)
+    s = f"[{tag}] E{e:2d}  loss={al:.4f}  ppl={ap:6.1f}"
+    for k,v in ex.items(): s+=f"  {k}={v/max(nb,1):.3f}"
+    print(s); return al,ap
+
+@torch.no_grad()
+def evaluate(model, dl):
+    model.eval(); tl,tp,nb=0.0,0.0,0
+    for b in dl:
+        ids=b['input_ids']; m=b.get('attention_mask')
+        out=model.loss(ids,m); tl+=out['loss'].item(); tp+=out['ppl'].item(); nb+=1
+    return tl/max(nb,1), tp/max(nb,1)
+
+# ═════════════════════════════════════════════════════════════════════
+# 11. MAIN BENCHMARK
+# ═════════════════════════════════════════════════════════════════════
+
+def main():
+    torch.manual_seed(42)
+    cfg = CFG(d_model=64, n_layers=2, n_heads=4, tt_rank=8,
+              q_qubits=4, q_sparsity=0.3, vocab=1000, max_seq=64)
+
+    print(f"Config: d={cfg.d_model} layers={cfg.n_layers} heads={cfg.n_heads} rank={cfg.tt_rank}")
+    print(f"Quantum: qubits={cfg.q_qubits} sparsity={cfg.q_sparsity}")
+    print(f"Tensor FFN: ON\n")
+
+    qt = QTensorFormer(cfg)
+    bl = Baseline(cfg)
+
+    pq = qt.nparams(); pb = bl.nparams()
+    print(f"Q-TensorFormer params: {pq['trainable']:>10,}")
+    print(f"Baseline params:       {pb['trainable']:>10,}")
+    print(f"Compression ratio:     {pb['trainable']/max(pq['trainable'],1):>10.1f}x\n")
+
+    train_dl = make_data(cfg.vocab, cfg.max_seq, 500, 16)
+    val_dl   = make_data(cfg.vocab, cfg.max_seq, 100, 16)
+    E = 8
+
+    print("=" * 50)
+    print(" TRAINING Q-TENSORFORMER")
+    print("=" * 50)
+    oq = torch.optim.AdamW(qt.parameters(), lr=cfg.lr)
+    sq = torch.optim.lr_scheduler.CosineAnnealingLR(oq, E*len(train_dl))
+    for e in range(1, E+1): train_epoch(qt, train_dl, oq, sq, e, "Q-TF")
+
+    print("\n" + "=" * 50)
+    print(" TRAINING BASELINE")
+    print("=" * 50)
+    ob = torch.optim.AdamW(bl.parameters(), lr=cfg.lr)
+    sb = torch.optim.lr_scheduler.CosineAnnealingLR(ob, E*len(train_dl))
+    for e in range(1, E+1): train_epoch(bl, train_dl, ob, sb, e, "BSL")
+
+    ql,qp = evaluate(qt, val_dl)
+    bl_val,bp = evaluate(bl, val_dl)
+
+    torch.save(qt.state_dict(), '/tmp/qt.pt')
+    torch.save(bl.state_dict(), '/tmp/bl.pt')
+    qsz = os.path.getsize('/tmp/qt.pt')/(1024*1024)
+    bsz = os.path.getsize('/tmp/bl.pt')/(1024*1024)
+
+    print("\n" + "=" * 65)
+    print(" RESULTS")
+    print("=" * 65)
+    print(f"{'Metric':<30} {'Q-TensorFormer':>15} {'Baseline':>15}")
+    print("-" * 60)
+    print(f"{'Parameters':<30} {pq['trainable']:>13,}  {pb['trainable']:>13,}")
+    print(f"{'Val Loss':<30} {ql:>15.4f} {bl_val:>15.4f}")
+    print(f"{'Val Perplexity':<30} {qp:>15.2f} {bp:>15.2f}")
+    print(f"{'Model Size (MB)':<30} {qsz:>15.1f} {bsz:>15.1f}")
+
+    ps = (1-pq['trainable']/pb['trainable'])*100
+    ss = (1-qsz/bsz)*100
+    pr = qp/bp
+    print(f"\nParameter reduction: {ps:.1f}%")
+    print(f"Size reduction:      {ss:.1f}%")
+    print(f"PPL ratio (Q-TF/BL): {pr:.2f}x")
+
+    if pr < 1.1:
+        print(f"\n  >> VERDICT: Significant compression with minimal quality loss! <<")
+    elif pr < 1.3:
+        print(f"\n  >> VERDICT: Moderate trade-off — compression worth the cost <<")
+    else:
+        print(f"\n  >> VERDICT: Quality gap too large, needs tuning <<")
+
+    print("\nDone!")
+    return {'params_q':pq['trainable'],'params_b':pb['trainable'],'qloss':ql,'qppl':qp,'bloss':bl_val,'bppl':bp,'qsz':qsz,'bsz':bsz,'comp':ps,'sred':ss,'ppl_ratio':pr}
+
+if __name__ == '__main__':
+    results = main()