File size: 6,004 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 | """Find concrete RAM reads in the in-weights CHIP-8 computer (beamboot/chip8_lm.py).
Runs the constructed BeamForCausalLM on one real SUPERBEAM frame (checkpoint prompt + the frame's
tokens, as chip8_fast does) with eager attention. For the RAM-read layers (L4: opcode and memory
reads, L5: Vx/Vy reads) it lists query positions whose attention sits on a single write token while
older writes to the same address exist earlier in the context: the "latest-write RAM" circuit,
read straight out of the model's attention weights.
.venv-gpu/bin/python scripts/attention_probe.py --frames 6 --out demo/attention_probe.json
"""
import argparse
import json
import os
import sys
import torch
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from beamboot import chip8_fast as F # noqa: E402
from beamboot import chip8_lm as C # noqa: E402
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--frames', type=int, default=6, help='advance this many frames first')
ap.add_argument('--draft-tokens', type=int, default=600)
ap.add_argument('--out', default=os.path.join(ROOT, 'demo', 'attention_probe.json'))
args = ap.parse_args()
rom = open(os.path.join(ROOT, 'roms', 'superbeam.ch8'), 'rb').read()
sched = json.load(open(os.path.join(ROOT, 'roms', 'superbeam_inputs.json')))
s = F.run(rom, [e for e in sched if int(e[0]) < args.frames], frames=args.frames, device='cpu')
for g, k, a in sched:
if int(g) == args.frames:
s.c.set_key(int(k) & 15, a == 'down')
p = F.checkpoint_prompt(s.c, s.ipf)
d = s._draft([])[:args.draft_tokens]
seq = p + d
m = s.model
m.config._attn_implementation = 'eager'
with torch.inference_mode():
out = m(input_ids=torch.tensor([seq]), output_attentions=True, logits_to_keep=len(d) + 1)
am = out.logits[0].argmax(-1).tolist()
agree = sum(1 for j, t in enumerate(d) if am[j] == t)
labels = [tuple(C.untok(t)) for t in seq]
def addr(lb):
return lb[0], (lb[2] if lb[0] == 'VW' and len(lb) > 2 else None)
mem = s.c.mem # code bytes do not change during the frame
def pc_at(q): # PC of the record a query belongs to
return next(int(labels[i][1]) for i in range(q - 1, -1, -1) if labels[i][0] == 'PCW')
found = []
for li in (4,): # L5: Vx / Vy reads
A = out.attentions[li][0] # (heads, q, k)
for h in range(A.shape[0]):
for q in range(len(p), len(seq)):
if labels[q][0] != 'CLKL': # the record's plan is computed at its CLKL token
continue
w = A[h, q, :q + 1]
kpos = int(w.argmax())
if float(w[kpos]) < 0.99 or labels[kpos][0] not in ('VW', 'PCW', 'IW'):
continue
if any(addr(labels[i]) == addr(labels[kpos]) for i in range(kpos + 1, q)):
continue # not the latest write before the query
same = [i for i in range(kpos) if addr(labels[i]) == addr(labels[kpos])]
if len(same) < 2:
continue
pc = pc_at(q)
found.append(dict(layer=li + 1, head=h, query=q, pc=pc, opcode='%02X%02X' % (mem[pc], mem[pc + 1]),
older_vals=[str(labels[i][1]) for i in same[-6:]], query_tok=[str(x) for x in labels[q]],
key=kpos, key_tok=[str(x) for x in labels[kpos]], weight=float(w[kpos]),
older_same=len(same), older_pos=same[-20:],
older_weight=float(w[same].sum()) if same else 0.0, ctx=q + 1))
found.sort(key=lambda f: (-f['older_same'], f['layer']))
# Head 0 of L5 is the Vx read. Check it on every instruction of the frame that uses Vx (x = second
# nibble): its argmax must be the newest VW write to register x before the query (or, if x was
# never written in context, no VW of x at all).
A0 = out.attentions[4][0][0]
uses_vx = set('345678ABEF') - {'A', 'B'}
checked = correct = 0
for q in range(len(p), len(seq)):
if labels[q][0] != 'CLKL':
continue
pc = pc_at(q)
op = '%02X%02X' % (mem[pc], mem[pc + 1])
if op[0] not in uses_vx:
continue
x = int(op[1], 16)
writes = [i for i in range(q) if labels[i][0] == 'VW' and int(labels[i][2]) == x]
if not writes:
continue
checked += 1
kpos = int(A0[q, :q + 1].argmax())
correct += int(kpos == writes[-1] and float(A0[q, kpos]) > 0.99)
stats = dict(head='L5 head 0 (Vx read)', instructions_checked=checked, read_newest_write=correct)
fig = None
best = [f for f in found if f['head'] == 0 and f['opcode'][0] in '34' and f['older_same'] >= 5]
if best:
f = best[0]
q, x = f['query'], int(f['key_tok'][2])
w = A0[q, :q + 1]
writes = [i for i in range(q) if labels[i][0] == 'VW' and int(labels[i][2]) == x]
fig = dict(query=q, ctx=q + 1, pc=f['pc'], opcode=f['opcode'], reg=x,
writes=[dict(pos=i, val=int(labels[i][1]), w=float(w[i])) for i in writes],
other_mass=float(w.sum() - w[writes].sum()),
kinds=[labels[i][0] for i in range(q + 1)])
res = dict(frame=args.frames, prompt_tokens=len(p), checked_tokens=len(d), draft_agree=agree,
n_reads=len(found), vx_read=stats, figure=fig, examples=found[:40])
print(json.dumps(stats))
os.makedirs(os.path.dirname(args.out), exist_ok=True)
json.dump(res, open(args.out, 'w'), indent=1)
print(json.dumps({k: v for k, v in res.items() if k != 'examples'}))
for f in found[:10]:
print({k: v for k, v in f.items() if k != 'older_pos'})
if __name__ == '__main__':
main()
|