File size: 8,198 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 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | """Serve the in-weights CHIP-8 computer (beamboot/chip8_lm.py) on a Modal GPU so the demo page can
run one game frame on the real model, live, from a machine state the visitor created.
modal deploy scripts/deploy_inweights_modal.py # prints the public URL
POST /frame {"mem": base64 of 4096 bytes, "v": [16 ints], "i": int, "pc": int, "dt": int,
"keys": [16 bools]} (the machine state at a frame boundary, keys already applied)
-> {"rows": 32 hex strings (the display decoded from the tokens the model accepted), "verified": bool,
"first_reject": int|null, "prompt_tokens", "frame_tokens", "instructions", "seconds", "gpu",
"sample": first token records as text, "state": the machine state after the frame, decoded from the
model's tokens (mem, v, i, pc, dt), so a client can keep playing on the model}
Same as beamboot/chip8_fast.FastSession.advance: checkpoint prompt + emulator draft, one forward pass
of the constructed BeamForCausalLM, every drafted token compared with the model's own argmax. If the
model disagrees anywhere it answers verified=false with the position (no greedy fallback here, to keep
requests short). One GPU container at most (L40S, else A10G, L4, T4), one forward at a time; it scales to zero when idle.
"""
from pathlib import Path
import modal
PKG = Path(__file__).resolve().parent.parent / "beamboot"
image = (
modal.Image.debian_slim(python_version="3.12")
.pip_install("torch==2.11.0", "transformers==5.22.1", "fastapi", "numpy")
.add_local_file(str(PKG / "__init__.py"), "/app/beamboot/__init__.py")
.add_local_file(str(PKG / "chip8.py"), "/app/beamboot/chip8.py")
.add_local_file(str(PKG / "chip8_lm.py"), "/app/beamboot/chip8_lm.py")
.add_local_file(str(PKG / "chip8_fast.py"), "/app/beamboot/chip8_fast.py")
)
app = modal.App("superbeam-inweights", image=image)
ORIGINS = ["https://staratnyte1-beam-boot.static.hf.space", "https://huggingface.co",
"http://localhost:8765"]
SAMPLE_RECORDS = 6
def _crop_to(cache, n):
"""Keep the first n positions of a DynamicCache (this transformers version takes a negative count)."""
extra = cache.get_seq_length() - n
if extra > 0:
cache.crop(-extra)
def _fmt(C, t):
k, v, a = C.untok(t)
return {"VW": "V%X=%d" % (a, v), "PCW": "PC=%03X" % v, "IW": "I=%03X" % v}.get(k, "%s %d" % (k, v))
@app.cls(gpu=["L40S", "A10G", "L4", "T4"], min_containers=0, max_containers=1, scaledown_window=300, timeout=300)
@modal.concurrent(max_inputs=12)
class InWeights:
@modal.enter()
def load(self):
import sys
import time
import torch
sys.path.insert(0, "/app")
from beamboot import chip8_lm as C
t0 = time.perf_counter()
self.model = C.get_model(300, "cuda")
self.model.config._attn_implementation = "sdpa"
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
self.gpu = torch.cuda.get_device_name(0)
self.build_seconds = round(time.perf_counter() - t0, 1)
import threading
self.lock = threading.Lock() # one forward at a time on the GPU; drafting overlaps it
# Drafting only needs the record tokens; skip the per-record state snapshot (a full display
# dump after every instruction), which was ~90% of the draft time. Tokens are unchanged.
C.snapshot = lambda c: None
self.pids, self.pcache = None, None # prompt-prefix KV cache: consecutive frames share it
def frame(self, req):
import base64
import copy
import time
import torch
from transformers import DynamicCache
from beamboot import chip8_fast as F
from beamboot import chip8_lm as C
from beamboot.chip8 import Chip8
mem = base64.b64decode(req["mem"])
if len(mem) != 4096 or len(req["v"]) != 16 or len(req["keys"]) != 16:
raise ValueError("bad machine state")
c = Chip8()
c.mem[:] = mem
c.v = [int(x) & 0xFF for x in req["v"]]
c.i, c.pc, c.dt = int(req["i"]) & 0xFFF, int(req["pc"]) & 0xFFF, int(req["dt"]) & 0xFF
c.keys = [bool(k) for k in req["keys"]]
s = F.FastSession.__new__(F.FastSession)
s.rom, s.ipf, s.c = bytes(mem[0x200:]), 300, c
t0 = time.perf_counter()
p = F.checkpoint_prompt(c, 300)
# Draft (as FastSession._draft) on a copy of the state; keep that copy: it is the state after the
# frame implied by the drafted tokens.
after, d, n_rec = copy.deepcopy(c), [], 0
for kind, ts, _ in C._ref_records(s.rom, 300, 0, lambda fr: (), chip=after):
d += ts
if kind in ("instr", "tick"):
n_rec += 1
if kind == "tick":
break
seq = p + d
td = time.perf_counter()
with torch.inference_mode(), self.lock:
t1 = time.perf_counter() # GPU time is measured after the queue
# Reuse the KV cache of the longest prompt prefix shared with the previous request (at most
# len(p) - 1, so the logits for every draft position are recomputed).
L = 0
if self.pids is not None:
n = min(len(self.pids), len(p) - 1)
while L < n and self.pids[L] == p[L]:
L += 1
if L >= 256:
cache = copy.deepcopy(self.pcache)
_crop_to(cache, L)
else:
cache, L = DynamicCache(config=self.model.config), 0
out_ = self.model(input_ids=torch.tensor([seq[L:]], device="cuda"), past_key_values=cache,
use_cache=True, logits_to_keep=len(d) + 1)
am = out_.logits[0, :len(d)].argmax(-1)
# first position where the model's own next token differs from the draft (one GPU compare)
bad = torch.nonzero(am != torch.tensor(d, device="cuda")).flatten()
rej = int(bad[0]) if len(bad) else None
_crop_to(cache, len(p))
self.pids, self.pcache = p, cache
torch.cuda.synchronize()
t2 = time.perf_counter()
sec = t2 - t1
if rej is None:
# The model accepted every drafted token, so its tokens are exactly the draft and the state
# they encode is `after` (token replay == reference: tests/test_chip8_fast.py).
fin = after
else:
# The model disagreed: decode what it did accept, token by token, and report where it stopped.
replay = F._new_replay(c, 300)
for t in d[:rej]:
replay.feed(t)
fin = replay.c
out = {"rows": fin.display_rows(), "verified": rej is None, "first_reject": rej,
"prompt_tokens": len(p), "frame_tokens": len(d), "instructions": max(0, n_rec - 1),
"seconds": round(sec, 3), "gpu": self.gpu, "draft_seconds": round(td - t0, 3), "cached_prefix": L,
"state": None if rej is not None else {"mem": base64.b64encode(bytes(fin.mem)).decode(),
"v": list(fin.v), "i": fin.i, "pc": fin.pc, "dt": fin.dt}}
out["server_seconds"] = round(time.perf_counter() - t0, 3)
return out
@modal.asgi_app()
def web(self):
from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
api = FastAPI()
api.add_middleware(CORSMiddleware, allow_origins=ORIGINS, allow_methods=["POST", "GET"],
allow_headers=["content-type"])
@api.get("/health")
def health():
return {"ok": True, "gpu": self.gpu, "model_build_seconds": self.build_seconds}
@api.post("/frame")
async def frame(request: Request):
from starlette.concurrency import run_in_threadpool
try:
return await run_in_threadpool(self.frame, await request.json())
except Exception as e: # bad state / unsupported program
raise HTTPException(status_code=400, detail=str(e)[:200])
return api
|