| |
| import os |
| os.environ['CUDA_VISIBLE_DEVICES'] = '' |
| os.environ.setdefault('DNA_CK', '/root/dna/ckpt/base.pt') |
| import time, json, resource, numpy as np, torch |
| torch.set_num_threads(int(os.environ.get('THREADS', '12'))) |
| from infer_dna import load_model, generate, make_ram_reader, load_tok |
|
|
| tok = load_tok() |
| m, cfg = load_model('cpu') |
| read_rows = make_ram_reader(m) |
|
|
| def run_once(prompt, n=64): |
| t0 = time.time() |
| _ = generate(m, tok, prompt, n=n, temp=0.0, device='cpu', read_rows=read_rows) |
| return n / (time.time() - t0) |
|
|
| prompts = ["Hello there", "The weather today", "In the beginning", "Science is", |
| "My favorite food", "The ocean is", "A long time ago", "Computers can", |
| "The best way to", "People often say"] |
| run_once("warmup", 8) |
| res = [] |
| for i, p in enumerate(prompts): |
| ts = run_once(p, 64); res.append(ts) |
| print(f'test {i+1}/10 prompt={p!r} tok_s={ts:.2f}', flush=True) |
| rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 |
| meta = {'tok_s': res, 'mean': float(np.mean(res)), 'std': float(np.std(res)), |
| 'rss_mb': rss, 'threads': torch.get_num_threads(), |
| 'io': 'codon table in RAM (warm, compute-bound)', 'gpu': False} |
| json.dump(meta, open('/root/dna/cpu_bench.json', 'w'), indent=2) |
| print('CPU_BENCH_DONE', json.dumps(meta), flush=True) |
|
|
