| """Run SUPERBEAM inside the constructed BeamForCausalLM (beamboot/chip8_lm.py) on the GPU. |
| |
| .venv-gpu/bin/python scripts/run_chip8_lm.py # first 5 frames, verified |
| .venv-gpu/bin/python scripts/run_chip8_lm.py --frames 171 |
| |
| Keys come from roms/superbeam_inputs.json (the keyboard device injects them as KEY tokens at |
| frame boundaries). Every token, every instruction's state (registers, stack, memory, display) and |
| every frame is checked against beamboot/chip8.py, and the frames also against |
| demo/superbeam_frames.json. Writes demo/superbeam_model_frames.json: |
| {rom, ipf, frames_run, inputs, verified, mismatch, matches_superbeam_frames_json, n_tokens, |
| prompt_tokens, decode_seconds, tok_per_s, device, model: {...}, frames: [[32 hex rows], ...]} |
| Prints tok/s and an ETA for all 171 frames (step time measured at longer cache lengths). |
| """ |
| import argparse |
| import json |
| import os |
| import sys |
| import time |
|
|
| ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| sys.path.insert(0, ROOT) |
|
|
| import torch |
|
|
| from beamboot import chip8_lm as C |
|
|
| ROM = os.path.join(ROOT, 'roms', 'superbeam.ch8') |
| INPUTS = os.path.join(ROOT, 'roms', 'superbeam_inputs.json') |
| REF_FRAMES = os.path.join(ROOT, 'demo', 'superbeam_frames.json') |
| OUT = os.path.join(ROOT, 'demo', 'superbeam_model_frames.json') |
| FULL = 171 |
|
|
|
|
| def total_tokens(rom, events, frames, ipf): |
| """Stream length (prompt excluded) for `frames` frames, from the reference.""" |
| keys = C.key_tokens(events) |
| n, nt = 0, 0 |
| for kind, ts, _ in C._ref_records(rom, ipf, 0, lambda fr: keys.get(fr, ())): |
| n += len(ts) |
| nt += kind == 'tick' |
| if nt == frames: |
| return n |
|
|
|
|
| def save(s, model, a, events, ref_frames): |
| res = s.result(tokens=False) |
| match = (res['frame_rows'] == ref_frames[:res['frames']]) if ref_frames else None |
| out = dict(rom='roms/superbeam.ch8', inputs_file=os.path.relpath(a.inputs, ROOT), ipf=a.ipf, frames_run=res['frames'], |
| inputs=[e for e in events if e[0] < res['frames']], verified=res['verified'], |
| mismatch=res['mismatch'], matches_superbeam_frames_json=match, n_tokens=res['n_tokens'], |
| prompt_tokens=res['prompt_tokens'], decode_seconds=round(res['decode_seconds'], 2), |
| tok_per_s=round(res['tok_per_s'], 1), device=torch.cuda.get_device_name(), |
| model=dict(cls=type(model).__name__, layers=model.config.num_hidden_layers, |
| hidden=model.config.hidden_size, vocab=model.config.vocab_size, dtype='float32', |
| attn=model.config._attn_implementation, decoder='GraphDecoder (CUDA graphs)'), |
| frames=res['frame_rows']) |
| with open(a.out + '.tmp', 'w') as f: |
| json.dump(out, f) |
| os.replace(a.out + '.tmp', a.out) |
| print('wrote %s (%d frames)' % (os.path.relpath(a.out, ROOT), res['frames']), flush=True) |
| res['matches'] = match |
| return res |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--frames', type=int, default=5) |
| ap.add_argument('--ipf', type=int, default=300) |
| ap.add_argument('--no-eta', action='store_true') |
| ap.add_argument('--inputs', default=INPUTS, help='key events list, or a schedule file with an "events" key') |
| ap.add_argument('--out', default=OUT) |
| ap.add_argument('--save-every', type=int, default=5, help='rewrite --out every N frames') |
| a = ap.parse_args() |
| rom = open(ROM, 'rb').read() |
| events = json.load(open(a.inputs)) |
| if isinstance(events, dict): |
| events = events['events'] |
| scripted = os.path.abspath(a.inputs) == os.path.abspath(INPUTS) |
| ref_frames = json.load(open(REF_FRAMES))['frames'] if scripted else None |
|
|
| t0 = time.perf_counter() |
| model = C.get_model(a.ipf, 'cuda') |
| print('model built in %.1fs: %s, %d layers, hidden %d, vocab %d, %.1fM params, %s' |
| % (time.perf_counter() - t0, type(model).__name__, model.config.num_hidden_layers, |
| model.config.hidden_size, model.config.vocab_size, |
| sum(p.numel() for p in model.parameters()) / 1e6, torch.cuda.get_device_name())) |
| s = C.Session(rom, 'cuda', a.ipf, verify=True, model=model) |
| print('prompt %d tokens, prefill %.2fs' % (s.prompt_tokens, s.prefill_seconds)) |
| by = {} |
| for g, k, act in events: |
| by.setdefault(g, []).append((k, act == 'down')) |
| for g in range(a.frames): |
| n0, t1 = len(s.tokens), time.perf_counter() |
| s.advance(by.get(g, ())) |
| dt = time.perf_counter() - t1 |
| ok = 'ok' if s.mismatch is None else 'MISMATCH %s' % s.mismatch |
| same = (s.result(tokens=False)['frame_rows'][-1] == ref_frames[g]) if ref_frames else 'n/a' |
| print('frame %3d: %5d tokens %6.1fs %6.1f tok/s verify %s frame==superbeam_frames.json %s' |
| % (g, len(s.tokens) - n0, dt, (len(s.tokens) - n0) / dt, ok, same)) |
| if s.mismatch is not None: |
| break |
| if a.save_every and (g + 1) % a.save_every == 0: |
| save(s, model, a, events, ref_frames) |
| res = save(s, model, a, events, ref_frames) |
| match = res['matches'] |
| print('frames %d, tokens %d, decode %.1fs, %.1f tok/s, verified %s, frames match json %s' |
| % (res['frames'], res['n_tokens'], res['decode_seconds'], res['tok_per_s'], res['verified'], match)) |
|
|
| if a.no_eta: |
| return |
| total = total_tokens(rom, events, FULL, a.ipf) + s.prompt_tokens |
| done = s.prompt_tokens + len(s.tokens) |
| pts = [done + 64] + [x for x in (16384, 32768, 65536, 131072) if done + 64 < x < total] + [total] |
| times = [(L, s.dec.time_at(L)) for L in pts] |
| print('step time vs cache length: ' + ', '.join('%d: %.1f ms (%.0f tok/s)' % (L, 1e3 * t, 1 / t) |
| for L, t in times)) |
| eta = 0.0 |
| for (l0, t0_), (l1, t1_) in zip(times, times[1:]): |
| eta += (l1 - l0) * (t0_ + t1_) / 2 |
| print('171 frames = %d tokens (incl. prompt); remaining ETA %.1f min, whole run ~%.1f min' |
| % (total, eta / 60, (eta + res['decode_seconds']) / 60)) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
