71 lines
2.5 KiB
Python
71 lines
2.5 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""REAL TRANSFORMER INFERENCE BENCHMARK"""
|
||
|
|
try:
|
||
|
|
import torch
|
||
|
|
import torch.nn as nn
|
||
|
|
except ImportError:
|
||
|
|
print("\033[91m[ERROR] torch not installed. Run: pip install torch\033[0m")
|
||
|
|
exit(1)
|
||
|
|
|
||
|
|
import time
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
class SmallTransformer(nn.Module):
|
||
|
|
def __init__(self, d_model=256, n_heads=4, n_layers=2, max_seq=1024):
|
||
|
|
super().__init__()
|
||
|
|
self.d_model = d_model
|
||
|
|
self.embedding = nn.Embedding(10000, d_model)
|
||
|
|
self.pos_embed = nn.Parameter(torch.randn(1, max_seq, d_model) * 0.02)
|
||
|
|
|
||
|
|
attn = nn.MultiheadAttention(d_model, n_heads, dropout=0, batch_first=True)
|
||
|
|
self.attention = nn.ModuleList([attn] * n_layers)
|
||
|
|
self.ffn = nn.Sequential(
|
||
|
|
nn.Linear(d_model, d_model * 4),
|
||
|
|
nn.ReLU(),
|
||
|
|
nn.Linear(d_model * 4, d_model)
|
||
|
|
)
|
||
|
|
self.norm = nn.LayerNorm(d_model)
|
||
|
|
self.lm_head = nn.Linear(d_model, 10000)
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
seq_len = x.size(1)
|
||
|
|
h = self.embedding(x) + self.pos_embed[:, :seq_len, :]
|
||
|
|
mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool().to(x.device)
|
||
|
|
|
||
|
|
for attn_layer in self.attention:
|
||
|
|
attn, _ = attn_layer(h, h, h, attn_mask=mask)
|
||
|
|
h = h + attn
|
||
|
|
h = h + self.ffn(h)
|
||
|
|
h = self.norm(h)
|
||
|
|
return self.lm_head(h)
|
||
|
|
|
||
|
|
def benchmark():
|
||
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||
|
|
print(f"\033[36mDevice:\033[0m {device}")
|
||
|
|
if device == 'cuda':
|
||
|
|
print(f"\033[36mGPU:\033[0m {torch.cuda.get_device_name(0)}")
|
||
|
|
|
||
|
|
config = {'d_model': 256, 'n_heads': 4, 'n_layers': 2, 'max_seq': 1024}
|
||
|
|
model = SmallTransformer(**config).to(device).eval()
|
||
|
|
|
||
|
|
input_ids = torch.randint(0, 10000, (1, 256), device=device)
|
||
|
|
|
||
|
|
times = []
|
||
|
|
print("\033[33mRunning 5 inference cycles...\033[0m")
|
||
|
|
for i in range(5):
|
||
|
|
if device == 'cuda': torch.cuda.synchronize()
|
||
|
|
start = time.perf_counter()
|
||
|
|
with torch.inference_mode():
|
||
|
|
_ = model(input_ids)
|
||
|
|
if device == 'cuda': torch.cuda.synchronize()
|
||
|
|
times.append((time.perf_counter() - start) * 1000)
|
||
|
|
print(f" Cycle {i+1}: {times[-1]:.2f} ms")
|
||
|
|
|
||
|
|
print(f"\n\033[1m=== TRANSFORMER BASELINE RESULTS ===\033[0m")
|
||
|
|
print(f"TTFT (256 tokens): {np.mean(times):.2f} ± {np.std(times):.2f} ms")
|
||
|
|
print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")
|
||
|
|
print(f"Model: {config['n_layers']}-layer, {config['d_model']}d")
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
benchmark()
|