File size: 4,579 Bytes
324ecd2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | #!/usr/bin/env python3
"""Replay a SUPERBEAM key schedule inside the constructed BeamForCausalLM, fast path
(beamboot/chip8_fast.py: per-frame checkpoints + speculative verification).
.venv-gpu/bin/python scripts/run_chip8_fast.py --inputs results/beam_key_schedule.json \
--frames 175 --out demo/beam_model_frames.json
Writes the same JSON shape as scripts/run_chip8_lm.py (frames = 32 hex rows per frame) plus
per-frame stats, and compares every frame with the reference emulator (superbeam_gym replay).
"""
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 # noqa: E402
from beamboot import chip8_fast as F # noqa: E402
from beamboot.superbeam_gym import SuperBeamEnv, frames_from_events # noqa: E402
def emulator_rows(events, n):
env = SuperBeamEnv(backend='emulator')
env.reset()
out = []
for keys in frames_from_events(events, n):
env.backend.advance(keys)
out.append(env.backend.m.display_rows())
return out
def save(path, s, ref, args, wall):
st = s.stats
k = 0
while k < len(s.frame_rows) and s.frame_rows[k] == ref[k]:
k += 1
secs = [x.seconds for x in st]
out = dict(rom='roms/superbeam.ch8', inputs_file=os.path.relpath(args.inputs, ROOT), ipf=s.ipf,
frames_run=len(s.frame_rows), verified=k == len(s.frame_rows),
frames_equal_emulator=k, mismatch=None if k == len(s.frame_rows) else k,
method='per-frame checkpoint + speculative verification (beamboot/chip8_fast.py)',
frames_all_draft_accepted=sum(x.accepted_all for x in st),
n_tokens=sum(x.frame_tokens for x in st), prompt_tokens_per_frame=st[-1].prompt_tokens,
seconds_per_frame=round(sum(secs) / len(secs), 3), max_seconds_per_frame=round(max(secs), 3),
wall_seconds=round(wall, 1),
tok_per_s=round(sum(x.frame_tokens for x in st) / sum(secs), 1),
device=torch.cuda.get_device_name() if torch.cuda.is_available() else 'cpu',
model=dict(cls=type(s.model).__name__, layers=s.model.config.num_hidden_layers,
hidden=s.model.config.hidden_size, vocab=s.model.config.vocab_size,
dtype='float32', attn='sdpa'),
per_frame=[dict(f=x.frame, tokens=x.frame_tokens, accepted_all=x.accepted_all,
first_reject=x.first_reject, s=round(x.seconds, 3)) for x in st],
frames=s.frame_rows)
with open(path + '.tmp', 'w') as f:
json.dump(out, f)
os.replace(path + '.tmp', path)
return out
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--inputs', default=os.path.join(ROOT, 'results', 'beam_key_schedule.json'))
ap.add_argument('--frames', type=int, default=None)
ap.add_argument('--out', default=os.path.join(ROOT, 'demo', 'beam_model_frames.json'))
ap.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu')
args = ap.parse_args()
sched = json.load(open(args.inputs))
events = sched['events'] if isinstance(sched, dict) else sched
n = args.frames or (sched.get('frames') if isinstance(sched, dict) else None)
if not n: # plain event list: run until the emulator's win
n = SuperBeamEnv(backend='emulator', frame_skip=1).replay_schedule(events)['frames']
rom = open(os.path.join(ROOT, 'roms', 'superbeam.ch8'), 'rb').read()
ref = emulator_rows(events, n)
t0 = time.time()
def on_frame(g, s):
x = s.stats[-1]
ok = s.frame_rows[-1] == ref[g]
print('frame %3d: %4d tokens %.2fs draft %s frame==emulator %s'
% (g, x.frame_tokens, x.seconds, 'accepted' if x.accepted_all else 'REJECTED at %d' % x.first_reject,
ok if ok else 'MISMATCH'), flush=True)
if (g + 1) % 10 == 0:
save(args.out, s, ref, args, time.time() - t0)
s = F.run(rom, events, frames=n, device=args.device, on_frame=on_frame)
out = save(args.out, s, ref, args, time.time() - t0)
print('done: %d frames, equal to emulator %d, all drafts accepted %d, %.2f s/frame, %.0f tok/s, wall %.0fs'
% (out['frames_run'], out['frames_equal_emulator'], out['frames_all_draft_accepted'],
out['seconds_per_frame'], out['tok_per_s'], out['wall_seconds']))
return 0 if out['verified'] else 1
if __name__ == '__main__':
raise SystemExit(main())
|