Buckets:
| """E4 latency: ms per decision at batch 1 (p50/p95 over 200 Core items) and batch-32 throughput, K = 1, 2, 4, 8. | |
| Latency does not depend on the weights, so any thinking checkpoint works. Eager bf16, like the evals.""" | |
| import argparse | |
| import json | |
| import os | |
| import random | |
| import sys | |
| import time | |
| import numpy as np | |
| import torch | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| from common import Renderer, load_core, tokenizer_and_mapping # noqa: E402 | |
| from eval import build_model # noqa: E402 | |
| from model import collate # noqa: E402 | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--model", required=True) | |
| ap.add_argument("--out", default="/workspace/think-decide/results/latency.json") | |
| ap.add_argument("--n", type=int, default=200) | |
| a = ap.parse_args() | |
| dev = torch.device("cuda") | |
| tok, mp, mid, mt = tokenizer_and_mapping() | |
| res = {} | |
| for nt, label in ((0, "S1 (no slots, no head)"), (32, "thinking (32 slots + head)")): | |
| items, _ = load_core(Renderer(tok, mid, mt, nt), 16384) | |
| items = random.Random(0).sample(items, a.n) | |
| model = build_model(a.model if nt else "s1", nt, nt > 0, dev, mp) | |
| for K in ([1] if nt == 0 else [1, 2, 4, 8]): | |
| for x in items[:5]: | |
| model(collate([x], dev, tok.pad_token_id), K=K) | |
| ts = [] | |
| for x in items: | |
| b = collate([x], dev, tok.pad_token_id) | |
| torch.cuda.synchronize() | |
| t = time.perf_counter() | |
| model(b, K=K) | |
| torch.cuda.synchronize() | |
| ts.append(1000 * (time.perf_counter() - t)) | |
| srt = sorted(items, key=lambda x: len(x["ids"])) | |
| torch.cuda.synchronize() | |
| t = time.perf_counter() | |
| for i in range(0, len(srt), 32): | |
| model(collate(srt[i:i + 32], dev, tok.pad_token_id), K=K) | |
| torch.cuda.synchronize() | |
| thr = len(srt) / (time.perf_counter() - t) | |
| res["%s K=%d" % (label, K)] = dict(p50_ms=round(float(np.percentile(ts, 50)), 2), | |
| p95_ms=round(float(np.percentile(ts, 95)), 2), | |
| batch32_items_per_s=round(thr, 1), | |
| median_tokens=int(np.median([len(x["ids"]) for x in items]))) | |
| print(json.dumps({"%s K=%d" % (label, K): res["%s K=%d" % (label, K)]}), flush=True) | |
| del model | |
| torch.cuda.empty_cache() | |
| json.dump(res, open(a.out, "w"), indent=1) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 2.62 kB
- Xet hash:
- b9a1a8a3e11f3eebddd1f040dfa8e0a1cc5615b34f59eb8d79b5cf5e9a295e62
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.