# /// script # requires-python = ">=3.10" # dependencies = [ # "torch", # "numpy", # "datasets>=2.19", # "huggingface_hub>=0.24", # "safetensors>=0.4", # ] # /// """CPU inference benchmark: Trident (O(1) state) vs the matched Transformer. The product claim is CPU-only inference with bounded memory. This measures, on actual CPU (no GPU), for both models: * throughput (kB/s) scoring bytes, in fp32 and int8-dynamic-quantized form, * quality retention under int8 (BPB fp32 vs int8 - a speedup only counts if quality holds), * how wall-time scales with context length, * the memory each needs to condition on a long context (Trident's constant state vs the transformer's O(context) KV cache). int8 dynamic quantization (weights int8, activations quantized per-op at runtime) is the honest "fortify the CPU moat" lever: it is exactly the kind of optimization that matters for CPU-only deployment and is applied identically to both models so the comparison stays fair. Configurable via TRIDENT_CKPT / XF_CKPT so the same bench serves the 44M and 150M checkpoints. Pushes results/cpu_bench.json (or RESULT_NAME). """ import json import math import os import sys import time from pathlib import Path import torch import torch.nn as nn def env(k, d=None): v = os.environ.get(k) return v if v not in (None, "") else d def log(m): print(f"[cpu] {m}", flush=True) def median_time(fn, reps=3): ts = [] for _ in range(reps): t0 = time.perf_counter() fn() ts.append(time.perf_counter() - t0) ts.sort() return ts[len(ts) // 2] def quantize(model): """int8 dynamic quantization of all Linear layers (weights int8, dynamic activation quant). Leaves everything else in fp32.""" return torch.ao.quantization.quantize_dynamic(model, {nn.Linear}, dtype=torch.qint8) def build_stream(dataset, name, split, field, total): from datasets import load_dataset ds = load_dataset(dataset, name=name, split=split, streaming=True) buf = bytearray() for row in ds: t = row.get(field) if not t: continue buf.extend(t.encode("utf-8")) buf.extend(b"\n\n") if len(buf) >= total: break return bytes(buf[:total]) def int8_bytes(model): """Approximate weight memory after int8 quantization: quantized Linear weights are 1 byte/elt; everything else stays fp32 (4 bytes/elt).""" lin = 0 other = 0 for m in model.modules(): if isinstance(m, nn.Linear): lin += m.weight.numel() if m.bias is not None: other += m.bias.numel() * 4 elif isinstance(m, nn.Embedding): other += m.weight.numel() * 4 return lin * 1 + other def main(): torch.set_grad_enabled(False) repo = env("CODE_REPO", "farguney/trident") threads = int(env("THREADS", "4")) torch.set_num_threads(threads) # pick an available int8 backend (fbgemm on x86 runners, qnnpack on ARM) for eng in ("fbgemm", "qnnpack"): if eng in torch.backends.quantized.supported_engines: torch.backends.quantized.engine = eng log(f"int8 engine={eng}") break result_name = env("RESULT_NAME", "cpu_bench.json") tri_ckpt = env("TRIDENT_CKPT", "checkpoints_tbptt/final") xf_ckpt = env("XF_CKPT", "checkpoints_baseline/final") from huggingface_hub import HfApi, snapshot_download from safetensors.torch import load_file work = Path("/tmp/code") snapshot_download(repo_id=repo, repo_type="model", allow_patterns=["src/**"], local_dir=str(work)) sys.path.insert(0, str(work / "src")) from trident import Trident, TridentConfig from trident.baseline import ByteTransformer ck = Path("/tmp/ck") snapshot_download(repo_id=repo, repo_type="model", allow_patterns=[f"{tri_ckpt}/*", f"{xf_ckpt}/*"], local_dir=str(ck)) tdir = ck / tri_ckpt tcfg = TridentConfig(**json.loads((tdir / "config.json").read_text())) trident = Trident(tcfg).eval() trident.load_state_dict(load_file(str(tdir / "model.safetensors")), strict=False) bdir = ck / xf_ckpt bcfg = json.loads((bdir / "config.json").read_text()) xf = ByteTransformer(**bcfg).eval() xf.load_state_dict(load_file(str(bdir / "model.safetensors")), strict=False) tri_params = sum(p.numel() for p in trident.parameters()) log(f"threads={threads} torch={torch.__version__}") log(f"Trident {tri_params/1e6:.1f}M ({tri_ckpt}) | Transformer {xf.num_params()/1e6:.1f}M ({xf_ckpt})") # int8-quantized copies trident_q = quantize(trident) xf_q = quantize(xf) log(f"int8 weight memory: Trident {int8_bytes(trident)/1e6:.1f} MB (fp32 {tri_params*4/1e6:.1f} MB) | " f"Transformer {int8_bytes(xf)/1e6:.1f} MB (fp32 {xf.num_params()*4/1e6:.1f} MB)") W = 2048 trident_state_bytes = tcfg.n_blocks * tcfg.n_heads * tcfg.d_v * tcfg.d_k * 4 def trident_stream(model, L): data = torch.randint(0, 256, (1, L), dtype=torch.long) def run(): state = None for w in range(0, L, W): out = model(data[:, w:w + W], state0=state, return_state=True) state = out["state"] return run def xf_forward(model, L): data = torch.randint(0, 256, (1, L), dtype=torch.long) return lambda: model(data) results = {"threads": threads, "trident_state_bytes": trident_state_bytes, "window": W, "trident_params_M": tri_params / 1e6, "transformer_params_M": xf.num_params() / 1e6, "int8_weight_bytes": {"trident": int8_bytes(trident), "transformer": int8_bytes(xf)}, "throughput": {}, "quality_retention": {}, "trident": {}, "transformer": {}} # ---- throughput: fp32 vs int8 (single 2 KB window) ---- log("== throughput (single 2KB window) ==") for name, tri_m, xf_m in [("fp32", trident, xf), ("int8", trident_q, xf_q)]: t_tri = median_time(trident_stream(tri_m, W)) t_xf = median_time(xf_forward(xf_m, W)) results["throughput"][name] = {"trident_kBps": W / t_tri / 1e3, "transformer_kBps": W / t_xf / 1e3, "trident_ms": t_tri * 1e3, "transformer_ms": t_xf * 1e3} log(f" [{name}] Trident {W/t_tri/1e3:6.1f} kB/s ({t_tri*1e3:5.0f} ms) | " f"Transformer {W/t_xf/1e3:6.1f} kB/s ({t_xf*1e3:5.0f} ms)") # ---- quality retention: does int8 hold BPB? (real held-out bytes) ---- try: nqw = int(env("QUAL_WINDOWS", "16")) raw = build_stream(env("VAL_DATASET", "Salesforce/wikitext"), env("VAL_NAME", "wikitext-103-raw-v1"), env("VAL_SPLIT", "validation"), env("VAL_TEXT_FIELD", "text"), nqw * W) qd = torch.tensor(list(raw), dtype=torch.long)[: nqw * W].view(nqw, W) def tri_bpb(model): nll, tot = 0.0, 0 for i in range(nqw): win = qd[i:i + 1] out = model(win, return_logits=True) lg = out["logits"].float() nll += torch.nn.functional.cross_entropy(lg.reshape(-1, lg.shape[-1]), win.reshape(-1), reduction="sum").item() tot += win.numel() return nll / tot / math.log(2) def xf_bpb(model): nll, tot = 0.0, 0 for i in range(nqw): win = qd[i:i + 1] out = model(win, return_logits=True) lg = out["logits"].float() nll += torch.nn.functional.cross_entropy(lg.reshape(-1, lg.shape[-1]), win.reshape(-1), reduction="sum").item() tot += win.numel() return nll / tot / math.log(2) tri_fp32, tri_int8 = tri_bpb(trident), tri_bpb(trident_q) xf_fp32, xf_int8 = xf_bpb(xf), xf_bpb(xf_q) results["quality_retention"] = { "trident_bpb_fp32": tri_fp32, "trident_bpb_int8": tri_int8, "transformer_bpb_fp32": xf_fp32, "transformer_bpb_int8": xf_int8, } log("== int8 quality retention (BPB, lower=better) ==") log(f" Trident fp32 {tri_fp32:.4f} -> int8 {tri_int8:.4f} (Δ{tri_int8-tri_fp32:+.4f})") log(f" Transformer fp32 {xf_fp32:.4f} -> int8 {xf_int8:.4f} (Δ{xf_int8-xf_fp32:+.4f})") except Exception as e: # noqa log(f"quality retention skipped ({e})") # ---- wall-time vs context length (fp32 + int8 for Trident stream) ---- log("== wall-time vs context length ==") for L in [2048, 8192, 32768, 98304]: t = median_time(trident_stream(trident, L), reps=2) tq = median_time(trident_stream(trident_q, L), reps=2) results["trident"][L] = {"sec": t, "kBps": L / t / 1e3, "sec_int8": tq, "kBps_int8": L / tq / 1e3, "mem_bytes": trident_state_bytes} log(f" Trident ctx={L:6d} fp32 {t*1e3:7.0f} ms ({L/t/1e3:5.1f} kB/s) | " f"int8 {tq*1e3:7.0f} ms ({L/tq/1e3:5.1f} kB/s) state={trident_state_bytes/1e6:.2f} MB (const)") for L in [1024, 2048, 4096, 8192]: t = median_time(xf_forward(xf, L), reps=2) tq = median_time(xf_forward(xf_q, L), reps=2) kv = 2 * bcfg["n_layers"] * bcfg["d_model"] * L * 4 # fp32 KV bytes results["transformer"][L] = {"sec": t, "kBps": L / t / 1e3, "sec_int8": tq, "kBps_int8": L / tq / 1e3, "kv_bytes": kv} log(f" Transformer ctx={L:6d} fp32 {t*1e3:7.0f} ms ({L/t/1e3:5.1f} kB/s) | " f"int8 {tq*1e3:7.0f} ms ({L/tq/1e3:5.1f} kB/s) KV={kv/1e6:.2f} MB") # ---- dominance: Trident/transformer speed ratio + transformer memory ceiling ---- mem_budget_gb = float(env("MEM_BUDGET_GB", "8")) log("== DOMINANCE (Trident vs Transformer) ==") dominance = {"speed_ratio": {}, "doc_scale": {}, "mem_budget_gb": mem_budget_gb} for L in [2048, 8192]: if L in results["trident"] and L in results["transformer"]: r = results["trident"][L]["kBps"] / results["transformer"][L]["kBps"] rq = results["trident"][L]["kBps_int8"] / results["transformer"][L]["kBps_int8"] dominance["speed_ratio"][L] = {"fp32": r, "int8": rq} log(f" ctx={L:5d}: Trident {r:.2f}x faster (fp32), {rq:.2f}x (int8)") for L in [65536, 262144, 1048576]: kv_gb = 2 * bcfg["n_layers"] * bcfg["d_model"] * L * 4 / 1e9 # transformer fp32 KV feasible = kv_gb <= mem_budget_gb dominance["doc_scale"][L] = {"transformer_kv_GB": kv_gb, "feasible": feasible, "trident_state_MB": trident_state_bytes / 1e6} log(f" doc L={L:7d}: transformer KV={kv_gb:6.2f} GB " f"({'ok' if feasible else 'INFEASIBLE'} @ {mem_budget_gb:.0f}GB) | " f"Trident state={trident_state_bytes/1e6:.2f} MB (const)") results["dominance"] = dominance out = Path("/tmp/results"); out.mkdir(exist_ok=True) (out / result_name).write_text(json.dumps(results, indent=2)) try: HfApi().upload_folder(folder_path=str(out), repo_id=repo, repo_type="model", path_in_repo="results", commit_message=f"cpu inference benchmark ({result_name})") log(f"pushed results/{result_name}") except Exception as e: # noqa log(f"upload failed: {e}") if __name__ == "__main__": main()