Hugging Face's logo

Premchan369
/
Q-TensorFormer

Text Generation
Transformers
Safetensors
PyTorch
English
qtensorformer
q-tensorformer
tensor-train
tensor-networks
adaptive-computation
kv-cache
kv-cache-quantization
attention-sinks
pid-control
edge-ai
edge-inference
on-device-ai
model-compression
grouped-query-attention
rotary-position-embeddings
custom_code
Model card Files
xet
 Community
Q-TensorFormer
File size: 4,156 Bytes
eaeea8f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
"""
Experiment Runner: Baselines Comparison.

Compares:
  - Dense Baseline Transformer
  - Static TT-compressed Transformer (fixed rank 4)
  - Quantized Baseline (INT8 / INT4 simulation)
  - Dynamic Depth (Layer skipping)
  - Adaptive Q-TensorFormer
"""

import sys
import os
import time
import math
import json
import argparse
from pathlib import Path

# Add project root
sys.path.insert(0, str(Path(__file__).parent.parent))

import torch
from src.config import ModelConfig
from src.models import QTensorFormer, DenseBaseline
from src.data import load_synthetic_data, load_wikitext2
from src.hardware_cost_model import HardwareCostModel


def evaluate_ppl(model, data_loader, device="cpu", max_batches=20):
    model.eval()
    model.to(device)
    total_loss, total_tokens = 0.0, 0
    with torch.no_grad():
        for i, (inp, tgt) in enumerate(data_loader):
            if i >= max_batches:
                break
            inp, tgt = inp.to(device), tgt.to(device)
            logits = model(inp)
            loss = torch.nn.functional.cross_entropy(
                logits.reshape(-1, logits.size(-1)),
                tgt.reshape(-1),
                ignore_index=0,
                reduction="sum",
            )
            total_loss += loss.item()
            total_tokens += inp.numel()
    avg = total_loss / max(1, total_tokens)
    return math.exp(min(avg, 20.0))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--d-model", type=int, default=128)
    parser.add_argument("--n-layers", type=int, default=2)
    parser.add_argument("--vocab-size", type=int, default=1000)
    parser.add_argument("--seq-len", type=int, default=32)
    parser.add_argument("--device", type=str, default="cpu")
    parser.add_argument("--output", type=str, default="outputs/baseline_results.json")
    args = parser.parse_args()

    os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
    device = torch.device(args.device)

    print("=" * 65)
    print("EXPERIMENT: Standardized Baseline Comparison")
    print("=" * 65)

    loader = load_synthetic_data(vocab_size=args.vocab_size, seq_len=args.seq_len, n_samples=500, batch_size=16)
    hw = HardwareCostModel()

    cfg = ModelConfig(
        vocab_size=args.vocab_size,
        d_model=args.d_model,
        n_layers=args.n_layers,
        max_seq_len=args.seq_len,
        tt_rank=8,
        use_quantum=True,
    )

    models = {
        "Dense Baseline": DenseBaseline(cfg),
        "Static TT (r=4)": QTensorFormer(cfg, preset="QTF_CLASSICAL_ONLY"),
        "Dynamic Depth": QTensorFormer(cfg, preset="QTF_LATENCY"),
        "Q-TensorFormer (Full)": QTensorFormer(cfg, preset="QTF_FULL"),
        "Q-TensorFormer (Balanced)": QTensorFormer(cfg, preset="QTF_BALANCED"),
    }

    results = []
    for name, model in models.items():
        if name == "Static TT (r=4)":
            for b in model.blocks:
                b.set_rank(4)

        model.to(device)
        total_p = model.total_params
        meas = hw.profile_execution(model, batch_size=1, seq_len=args.seq_len, n_repeats=5, device=device)
        ppl = evaluate_ppl(model, loader, device=device, max_batches=10)

        record = {
            "model": name,
            "params": total_p,
            "latency_p50_ms": meas.latency_ms,
            "ttft_ms": meas.ttft_ms,
            "tpot_ms": meas.tpot_ms,
            "throughput_tokens_sec": meas.throughput_tokens_sec,
            "traffic_bytes_per_token": meas.memory_traffic_bytes_per_token,
            "peak_memory_mb": meas.peak_memory_mb,
            "joules_per_token": meas.joules_per_token,
            "energy_level": meas.energy_level,
            "perplexity": round(ppl, 2),
            "classification": "MEASURED",
            "hardware": meas.device_name,
        }
        results.append(record)

        print(f"{name:<26} | Params: {total_p:>8,} | Lat: {meas.latency_ms:>6.2f}ms | J/tok: {meas.joules_per_token*1e6:>6.2f}uJ | PPL: {ppl:>6.2f}")

    with open(args.output, "w") as f:
        json.dump(results, f, indent=2)

    print(f"\nResults saved to {args.output}")


if __name__ == "__main__":
    main()