| """chipformer: a decoder-only transformer whose weights were set by hand to execute CHIP-8.
|
|
|
| The network is ordinary: token embedding plus a position embedding from [t, t^2], causal
|
| softmax attention heads with residual connections, ReGLU feed-forward layers and a linear
|
| unembedding, all in float64, with no normalization. Nothing about CHIP-8 is in this file
|
| except the host loop at the bottom (ChipformerGame), which only feeds the boot prompt,
|
| injects TICK and key tokens, and reads the screen back from the model's own writes.
|
|
|
| model = AutoModelForCausalLM.from_pretrained("divagr1925/chipformer", trust_remote_code=True,
|
| dtype=torch.float64)
|
| game = model.play(open("tetris.ch8", "rb").read(), quirks="shift,loadstore,clip,jump", ipf=15)
|
| for _ in range(60):
|
| game.frame()
|
| print(game.ascii())
|
| """
|
|
|
| import warnings
|
| from dataclasses import dataclass
|
| from pathlib import Path
|
|
|
| import torch
|
| from torch import nn
|
| from transformers import PreTrainedModel
|
| from transformers.modeling_outputs import CausalLMOutputWithPast
|
|
|
| from .configuration_chipformer import ChipformerConfig
|
|
|
| SPARSE_MIN = 100_000
|
|
|
|
|
| class ChipformerCache:
|
| """Keys and values per layer, shaped (batch, heads, positions, width), plus the length.
|
|
|
| Buffers grow by doubling, so appending one token costs O(1) amortized.
|
| """
|
|
|
| def __init__(self):
|
| self.kv: dict[int, list[torch.Tensor]] = {}
|
| self.length = 0
|
|
|
| def get_seq_length(self, layer_idx: int = 0) -> int:
|
| return self.length
|
|
|
| def append(self, li: int, k: torch.Tensor, v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
| """Store a layer's new keys/values; returns all keys and values so far."""
|
| start, n = self.length, k.shape[-2]
|
| if li not in self.kv:
|
| cap = max(1024, 2 * n)
|
| self.kv[li] = [x.new_zeros(*x.shape[:-2], cap, x.shape[-1]) for x in (k, v)]
|
| buf = self.kv[li]
|
| if start + n > buf[0].shape[-2]:
|
| cap = max(2 * buf[0].shape[-2], start + n)
|
| for j, old in enumerate(buf):
|
| buf[j] = old.new_zeros(*old.shape[:-2], cap, old.shape[-1])
|
| buf[j][..., :start, :] = old[..., :start, :]
|
| buf[0][..., start:start + n, :] = k
|
| buf[1][..., start:start + n, :] = v
|
| return buf[0][..., :start + n, :], buf[1][..., :start + n, :]
|
|
|
|
|
| class ChipformerHead(nn.Module):
|
| def __init__(self, d: int, qk: int, v: int):
|
| super().__init__()
|
| self.wq = nn.Parameter(torch.zeros(qk, d, dtype=torch.float64))
|
| self.wk = nn.Parameter(torch.zeros(qk, d, dtype=torch.float64))
|
| self.wv = nn.Parameter(torch.zeros(v, d, dtype=torch.float64))
|
| self.wo = nn.Parameter(torch.zeros(d, v, dtype=torch.float64))
|
|
|
|
|
| class ChipformerFFN(nn.Module):
|
| """ReGLU: w_down (relu(w_gate x) * (w_up x))."""
|
|
|
| def __init__(self, d: int, units: int):
|
| super().__init__()
|
| self.w_gate = nn.Parameter(torch.zeros(units, d, dtype=torch.float64))
|
| self.w_up = nn.Parameter(torch.zeros(units, d, dtype=torch.float64))
|
| self.w_down = nn.Parameter(torch.zeros(d, units, dtype=torch.float64))
|
|
|
|
|
| class ChipformerLayer(nn.Module):
|
| def __init__(self, d: int, spec: dict):
|
| super().__init__()
|
| self.attn = nn.ModuleList(ChipformerHead(d, qk, v) for qk, v in spec["heads"])
|
| self.ffn = ChipformerFFN(d, spec["ffn_units"])
|
|
|
|
|
| class ChipformerEmbedding(nn.Module):
|
| def __init__(self, vocab: int, d: int):
|
| super().__init__()
|
| self.token = nn.Parameter(torch.zeros(vocab, d, dtype=torch.float64))
|
| self.pos_poly = nn.Parameter(torch.zeros(d, 2, dtype=torch.float64))
|
|
|
|
|
| class ChipformerPreTrainedModel(PreTrainedModel):
|
| config_class = ChipformerConfig
|
| base_model_prefix = "chipformer"
|
| _supports_cache_class = False
|
|
|
| def _init_weights(self, module):
|
| pass
|
|
|
|
|
| class ChipformerForCausalLM(ChipformerPreTrainedModel):
|
| def __init__(self, config: ChipformerConfig):
|
| super().__init__(config)
|
| d = config.d_model
|
| self.embed = ChipformerEmbedding(config.vocab_size, d)
|
| self.layers = nn.ModuleList(ChipformerLayer(d, s) for s in config.layers)
|
| self.unembed = nn.Parameter(torch.zeros(config.vocab_size, d, dtype=torch.float64))
|
| self._csr: dict[str, torch.Tensor] | None = None
|
| self._fused: list[tuple[torch.Tensor, ...]] | None = None
|
| self.post_init()
|
|
|
| def sparse(self, enable: bool = True) -> "ChipformerForCausalLM":
|
| """Multiply by sparse (CSR) copies of the feed-forward and unembedding matrices.
|
|
|
| Only 0.4% of the weights are nonzero, and the dense feed-forward matrices are most of
|
| the 259 MB that every decoded token otherwise reads. The arithmetic is exact either
|
| way (the circuits only produce values float64 represents exactly), so the outputs
|
| are identical.
|
| """
|
| if not enable:
|
| self._csr = None
|
| return self
|
| with warnings.catch_warnings():
|
| warnings.simplefilter("ignore")
|
|
|
| self._csr = {name: p.detach().to_sparse_csr() for name, p in self.named_parameters()
|
| if (".ffn." in name or name == "unembed") and p.numel() >= SPARSE_MIN}
|
| return self
|
|
|
| def _lin(self, x: torch.Tensor, w: torch.Tensor, name: str) -> torch.Tensor:
|
| """x @ w.T, through the sparse copy when enabled."""
|
| m = None if self._csr is None else self._csr.get(name)
|
| if m is None:
|
| return x @ w.T
|
| y = m @ x.reshape(-1, x.shape[-1]).T
|
| return y.T.reshape(*x.shape[:-1], m.shape[0])
|
|
|
| def _heads(self, li: int) -> tuple[torch.Tensor, ...] | None:
|
| """A layer's heads stacked into (heads, width, d) tensors, zero-padded to one width.
|
|
|
| Padding adds exact zeros, so this is the same computation as running each head alone.
|
| Built on first use; call ``refresh()`` after changing the weights in place.
|
| """
|
| if self._fused is None:
|
| self._fused = []
|
| for layer in self.layers:
|
| if not len(layer.attn):
|
| self._fused.append(None)
|
| continue
|
| qk = max(h.wq.shape[0] for h in layer.attn)
|
| vd = max(h.wv.shape[0] for h in layer.attn)
|
|
|
| def stack(ws, width):
|
| return torch.stack([nn.functional.pad(w.detach(), (0, 0, 0, width - w.shape[0]))
|
| for w in ws])
|
| self._fused.append((
|
| stack([h.wq for h in layer.attn], qk), stack([h.wk for h in layer.attn], qk),
|
| stack([h.wv for h in layer.attn], vd),
|
| stack([h.wo.T for h in layer.attn], vd).transpose(-1, -2)))
|
| return self._fused[li]
|
|
|
| def refresh(self) -> None:
|
| """Rebuild the stacked/sparse copies after editing weights in place."""
|
| self._fused = None
|
| if self._csr is not None:
|
| self.sparse()
|
|
|
| def get_input_embeddings(self):
|
| return None
|
|
|
| def _attend(self, q, k, v, start: int) -> torch.Tensor:
|
| """Causal softmax attention; query i sits at position start + i, key j at position j."""
|
| s = q @ k.transpose(-1, -2)
|
| n, total = s.shape[-2], s.shape[-1]
|
| if n > 1:
|
| qpos = torch.arange(start, start + n, device=s.device)[:, None]
|
| s = s.masked_fill(torch.arange(total, device=s.device)[None, :] > qpos, float("-inf"))
|
| s = s - s.amax(-1, keepdim=True)
|
| p = torch.exp(self.config.beta * s)
|
| return (p / p.sum(-1, keepdim=True)) @ v
|
|
|
| def forward(self, input_ids: torch.LongTensor, past_key_values: ChipformerCache | None = None,
|
| use_cache: bool | None = None, labels: torch.LongTensor | None = None,
|
| logits_to_keep: int = 0, return_dict: bool | None = None, **kwargs):
|
| """input_ids: (batch, n). Positions continue from the cache length (absolute index)."""
|
| cache = past_key_values if past_key_values is not None else ChipformerCache()
|
| start = cache.length
|
| n = input_ids.shape[-1]
|
| if start + n > self.config.max_position:
|
| raise ValueError(f"position {start + n} exceeds max_position {self.config.max_position}; "
|
| "restart the context from a state dump (see ChipformerGame.rollover)")
|
| t = torch.arange(start, start + n, dtype=torch.float64, device=input_ids.device)
|
| r = self.embed.token[input_ids] + torch.stack([t, t * t], 1) @ self.embed.pos_poly.T
|
| for li, layer in enumerate(self.layers):
|
| heads = self._heads(li)
|
| if heads is not None:
|
| wq, wk, wv, wo = heads
|
| x = r.unsqueeze(-3)
|
| k, v = cache.append(li, x @ wk.transpose(-1, -2), x @ wv.transpose(-1, -2))
|
| out = self._attend(x @ wq.transpose(-1, -2), k, v, start)
|
| r = r + (out @ wo.transpose(-1, -2)).sum(-3)
|
| f, p = layer.ffn, f"layers.{li}.ffn"
|
| hidden = torch.relu(self._lin(r, f.w_gate, f"{p}.w_gate")) * self._lin(r, f.w_up, f"{p}.w_up")
|
| r = r + self._lin(hidden, f.w_down, f"{p}.w_down")
|
| logits = self._lin(r if not logits_to_keep else r[..., -logits_to_keep:, :], self.unembed, "unembed")
|
| loss = None
|
| if labels is not None:
|
| loss = nn.functional.cross_entropy(logits[..., :-1, :].reshape(-1, logits.shape[-1]),
|
| labels[..., 1:].reshape(-1))
|
| cache.length = start + n
|
| return CausalLMOutputWithPast(loss=loss, logits=logits,
|
| past_key_values=None if use_cache is False else cache)
|
|
|
| def play(self, rom: "bytes | str | Path", quirks="octo", ipf: int = 10,
|
| keep_tokens: bool = False, sparse: bool = True) -> "ChipformerGame":
|
| """Boot a CHIP-8 ROM; call ``.frame()`` on the result once per 60 Hz frame.
|
|
|
| ``sparse`` switches on the exact sparse matmuls (see ``sparse()``); pass False to run
|
| the dense weights as loaded.
|
| """
|
| if sparse and self._csr is None:
|
| self.sparse()
|
| elif not sparse:
|
| self.sparse(False)
|
| return ChipformerGame(self, rom, quirks=quirks, ipf=ipf, keep_tokens=keep_tokens)
|
|
|
|
|
|
|
|
|
| MAX_RECORD = 2 * 40 + 2
|
|
|
|
|
| @dataclass
|
| class _Machine:
|
| """Token layout and boot constants, read from config.machine."""
|
|
|
| VAL_BASE: int
|
| ADDR_BASE: int
|
| STEP: int
|
| TICK: int
|
| BOOT: int
|
| NUM_CELLS: int
|
| NULL: int
|
| CLS: int
|
| PC: int
|
| RNG: int
|
| KEYS: int
|
| DISPLAY: int
|
| DISPLAY_W: int
|
| DISPLAY_H: int
|
| PROGRAM_START: int
|
| FONT_BASE: int
|
| RNG_SEED: int
|
| QUIRKS: int
|
| FONT: list
|
| QUIRK_NAMES: list
|
| PROFILES: dict
|
|
|
|
|
| class ChipformerGame:
|
| """Runs a CHIP-8 program on the model, the way the Rust/wasm engine does.
|
|
|
| The host never computes machine semantics. It (1) feeds the boot prompt
|
| ``BOOT <state dump> STEP``, (2) injects ``TICK`` and key writes ``ADDR k, VAL v, STEP``
|
| right after a STEP, and (3) keeps a mirror of the cells from the model's own writes so it
|
| can show the screen. Every other token is the model's greedy prediction.
|
| """
|
|
|
| def __init__(self, model: ChipformerForCausalLM, rom, quirks="octo", ipf: int = 10,
|
| keep_tokens: bool = False):
|
| if isinstance(rom, (str, Path)):
|
| rom = Path(rom).read_bytes()
|
| self.model = model
|
| self.m = _Machine(**model.config.machine)
|
| self.ipf = ipf
|
| self.keep_tokens = keep_tokens
|
| self.tokens: list[int] = []
|
| self.instructions = 0
|
| self.rollovers = 0
|
| self.fed = 0
|
| m = self.m
|
| if len(rom) > 0x1000 - m.PROGRAM_START:
|
| raise ValueError(f"ROM too large: {len(rom)} bytes")
|
| cells = [0] * m.NUM_CELLS
|
| cells[m.FONT_BASE:m.FONT_BASE + len(m.FONT)] = m.FONT
|
| cells[m.PROGRAM_START:m.PROGRAM_START + len(rom)] = list(rom)
|
| cells[m.PC] = m.PROGRAM_START
|
| cells[m.RNG] = m.RNG_SEED
|
| for a, v in self._quirk_cells(quirks).items():
|
| cells[a] = v
|
| self.cells = cells
|
| self._new_context()
|
|
|
| def _quirk_cells(self, quirks) -> dict[int, int]:
|
| m = self.m
|
| if quirks is None:
|
| on = []
|
| elif isinstance(quirks, str):
|
| on = m.PROFILES[quirks] if quirks in m.PROFILES else [q for q in quirks.split(",") if q]
|
| else:
|
| on = list(quirks)
|
| bad = set(on) - set(m.QUIRK_NAMES)
|
| if bad:
|
| raise ValueError(f"unknown quirks {sorted(bad)}; known: {m.QUIRK_NAMES} or {list(m.PROFILES)}")
|
| return {m.QUIRKS + i: int(q in on) for i, q in enumerate(m.QUIRK_NAMES)}
|
|
|
| def boot_tokens(self) -> list[int]:
|
| """BOOT, the CLS write, every other nonzero cell in address order, STEP."""
|
| m, c = self.m, self.cells
|
| out = [m.BOOT, m.ADDR_BASE + m.CLS, m.VAL_BASE + c[m.CLS]]
|
| for a, v in enumerate(c):
|
| if a not in (m.CLS, m.NULL) and v:
|
| out += [m.ADDR_BASE + a, m.VAL_BASE + v]
|
| return out + [m.STEP]
|
|
|
| def _new_context(self) -> None:
|
| self.cache = ChipformerCache()
|
| self._feed(self.boot_tokens())
|
|
|
| @torch.inference_mode()
|
| def _feed(self, toks: list[int]) -> None:
|
| ids = torch.tensor([toks], dtype=torch.long, device=self.model.device)
|
| out = self.model(ids, past_key_values=self.cache, use_cache=True, logits_to_keep=1)
|
| self.cache = out.past_key_values
|
| self.pred = int(out.logits[0, -1].argmax())
|
| self.fed += len(toks)
|
| if self.keep_tokens:
|
| self.tokens += toks
|
|
|
| @property
|
| def position(self) -> int:
|
| return self.cache.length
|
|
|
| def _apply(self, writes) -> None:
|
| m = self.m
|
| for a, v in writes:
|
| if a == m.CLS:
|
| n = m.DISPLAY_W * m.DISPLAY_H // 8
|
| self.cells[m.DISPLAY:m.DISPLAY + n] = [0] * n
|
| self.cells[a] = v
|
|
|
| def _record(self) -> list[tuple[int, int]]:
|
| """Let the model write until it emits STEP."""
|
| m, writes = self.m, []
|
| for _ in range(MAX_RECORD):
|
| tok = self.pred
|
| if tok == m.STEP:
|
| self._feed([tok])
|
| self._apply(writes)
|
| return writes
|
| if not m.ADDR_BASE <= tok < m.STEP:
|
| raise RuntimeError(f"model emitted {tok} where ADDR/STEP was expected at {self.position}")
|
| self._feed([tok])
|
| val = self.pred
|
| if not m.VAL_BASE <= val < m.ADDR_BASE:
|
| raise RuntimeError(f"model emitted {val} after an ADDR at {self.position}")
|
| self._feed([val])
|
| writes.append((tok - m.ADDR_BASE, val - m.VAL_BASE))
|
| raise RuntimeError("record never ended")
|
|
|
| def step(self) -> list[tuple[int, int]]:
|
| """Execute one instruction (the model writes its effects)."""
|
| self.instructions += 1
|
| return self._record()
|
|
|
| def tick(self) -> list[tuple[int, int]]:
|
| """60 Hz timer tick: inject TICK, the model writes DT and ST."""
|
| self._feed([self.m.TICK])
|
| return self._record()
|
|
|
| def set_key(self, key: int, down: bool) -> None:
|
| m = self.m
|
| a = m.KEYS + key
|
| if bool(self.cells[a]) == down:
|
| return
|
| self._feed([m.ADDR_BASE + a, m.VAL_BASE + int(down), m.STEP])
|
| self.cells[a] = int(down)
|
|
|
| def rollover(self) -> None:
|
| """Restart the context from a state dump of the mirror (keeps positions bounded)."""
|
| self._new_context()
|
| self.rollovers += 1
|
|
|
| def frame(self, keys=()) -> None:
|
| """One 60 Hz frame: key changes, one TICK, then ``ipf`` instructions.
|
|
|
| ``keys``: the CHIP-8 keys (0..15) held this frame, or a 16-bit mask.
|
| """
|
| held = {k for k in range(16) if keys >> k & 1} if isinstance(keys, int) else set(keys)
|
| for k in range(16):
|
| self.set_key(k, k in held)
|
| self.tick()
|
| for _ in range(self.ipf):
|
| self.step()
|
| if self.position > self.model.config.rollover_at:
|
| self.rollover()
|
|
|
| def pixels(self) -> list[list[int]]:
|
| m = self.m
|
| w8 = m.DISPLAY_W // 8
|
| return [[self.cells[m.DISPLAY + y * w8 + x // 8] >> (7 - x % 8) & 1 for x in range(m.DISPLAY_W)]
|
| for y in range(m.DISPLAY_H)]
|
|
|
| def ascii(self) -> str:
|
| """The 64x32 screen as 16 lines of half-block characters."""
|
| p = self.pixels()
|
| chars = {(0, 0): " ", (1, 0): "▀", (0, 1): "▄", (1, 1): "█"}
|
| return "\n".join("".join(chars[p[y][x], p[y + 1][x]] for x in range(len(p[0])))
|
| for y in range(0, len(p), 2))
|
|
|
