divagr1925

Model card: lead with the game running inside the transformer

8ef5dd8 verified 12 days ago

5.24 kB

---
license: mit
library_name: transformers
pipeline_tag: text-generation
tags:
- chip-8
- emulator
- mechanistic-interpretability
- compiled-transformer
- custom_code
---

# chipformer

**Tetris, running entirely inside a transformer.**

This model *is* a CHIP-8 computer. Give it a game as the prompt and greedy decoding runs the game:
every token it predicts is one write to the machine's memory (a register, the program counter, a
row of pixels). Paint those writes onto a 64×32 screen and you are playing Tetris. Your keypresses
and the 60 Hz clock go into the context window. Everything else, every frame, every collision,
every point on the scoreboard, comes out of the model.

There is no emulator behind it. It is the same loop every language model runs: read the context,
predict the next token, append it, repeat.

- One set of weights plays Tetris, Space Invaders, Brix, Pong and 100+ other games. Switching games
  means switching the prompt.
- An ordinary decoder-only architecture: 18 layers, d_model 382,
  42 attention heads, 22,276 ReGLU units, vocabulary 8,708, float64.

**Play it in your browser:** [divagr.com/chipformer](https://divagr.com/chipformer) ·
[Hugging Face Space](https://huggingface.co/spaces/divagr1925/chipformer) · Source: [GitHub](https://github.com/divagr18/chipformer)

## How it works

- **The ROM is the prompt.** The prompt is `BOOT`, then one `ADDR a, VAL v` token pair per nonzero
  byte of machine state, then `STEP`.
- **Every output token is a write.** Each instruction becomes a short record of writes such as
  `@V3 =1F @PC =22A STEP`. The screen is drawn from these writes.
- **The host injects the outside world.** A `TICK` token arrives once per 60 Hz frame, and keypresses
  are injected as writes. Every other token is the model's greedy prediction.
- **Attention is memory.** To read register V3, a head attends to the most recent write to V3.
  Keys use a parabola over the address plus a recency term, so the hardmax lands exactly there.
- **ReLU is the ALU.** Addition, comparison, XOR and opcode decoding are all built from ReLUs,
  using identities like step(x) = ReLU(x+1) - ReLU(x).

The architecture is ordinary:
- token embedding plus a position embedding from [t, t²]
- causal softmax attention (β = 1000) with residual connections
- ReGLU feed-forward layers
- a linear unembedding
- no normalization

`modeling_chipformer.py` contains nothing CHIP-8 specific apart from the host loop, which only feeds
tokens in and reads the screen back.

**Where the weights come from.** They are compiled from a description of the CHIP-8 circuit
(written in a small DSL in the GitHub repo) rather than learned, which is what makes every
instruction exact. 139,229 of the 32,328,660 parameters are nonzero.

## Usage

```python
import torch
from huggingface_hub import hf_hub_download
from transformers import AutoModelForCausalLM

repo = "divagr1925/chipformer"
model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True, dtype=torch.float64)

rom = open(hf_hub_download(repo, "roms/tetris.ch8"), "rb").read()
game = model.play(rom, quirks="shift,loadstore,clip,jump", ipf=15)
for frame in range(120):
    game.frame(keys={5} if frame % 20 < 4 else ())  # CHIP-8 key 5 = left
print(game.ascii())
```

`play_tetris.py` animates the game in a terminal: `python play_tetris.py`.

The model is a normal `PreTrainedModel`, so you can also call it directly:
- `model(input_ids, past_key_values=cache)` returns next-token logits over the 8,708-token vocabulary.
- `game.tokens` (with `keep_tokens=True`) gives you the full trace.

### Notes

- **Use float64.** The memory-read heads compare attention scores of about 10¹³, and these are only
  exact in float64. fp16 and bf16 will break the machine.
- **Speed.** Decoding is greedy, one token at a time with a KV cache. On a CPU in PyTorch that is
  about 200 tokens/s, or 2–3 frames per second of Tetris (the game wants 60).
  - `game = model.play(...)` multiplies by sparse copies of the big feed-forward matrices, since only
    0.4% of the weights are nonzero. The outputs are identical; pass `sparse=False` to run the dense
    weights exactly as loaded.
  - The Rust/wasm engine in the GitHub repo uses the same weights with sparse matmuls and exact
    attention shortcuts. It runs the games at full speed in a browser.
- **Context.** Positions must stay below 4,194,304 (2²²). `ChipformerGame` restarts the context
  from a state dump every 250,000 tokens, as the engine does.
- **Coverage.** The weights implement the original CHIP-8 instruction set with Octo's quirk flags.
  SCHIP/XO-CHIP opcodes and the vBlank display-wait quirk are not supported.

## Files

| File | |
|---|---|
| `model.safetensors` | dense float64 weights (what `transformers` loads) |
| `model.chipformer` | the same weights, sparse (nonzeros only), for the Rust/wasm engine |
| `configuration_chipformer.py`, `modeling_chipformer.py` | the architecture plus the host loop |
| `play_tetris.py` | a terminal demo |
| `roms/tetris.ch8` | Tetris by Fran Dachille (1991), freely distributable (Revival Studios CHIP-8 pack) |

	---
	license: mit
	library_name: transformers
	pipeline_tag: text-generation
	tags:
	- chip-8
	- emulator
	- mechanistic-interpretability
	- compiled-transformer
	- custom_code
	---

	# chipformer

	Tetris, running entirely inside a transformer.

	This model is a CHIP-8 computer. Give it a game as the prompt and greedy decoding runs the game:
	every token it predicts is one write to the machine's memory (a register, the program counter, a
	row of pixels). Paint those writes onto a 64×32 screen and you are playing Tetris. Your keypresses
	and the 60 Hz clock go into the context window. Everything else, every frame, every collision,
	every point on the scoreboard, comes out of the model.

	There is no emulator behind it. It is the same loop every language model runs: read the context,
	predict the next token, append it, repeat.

	- One set of weights plays Tetris, Space Invaders, Brix, Pong and 100+ other games. Switching games
	means switching the prompt.
	- An ordinary decoder-only architecture: 18 layers, d_model 382,
	42 attention heads, 22,276 ReGLU units, vocabulary 8,708, float64.

	Play it in your browser: [divagr.com/chipformer](https://divagr.com/chipformer) ·
	[Hugging Face Space](https://huggingface.co/spaces/divagr1925/chipformer) · Source: [GitHub](https://github.com/divagr18/chipformer)

	## How it works

	- The ROM is the prompt. The prompt is `BOOT`, then one `ADDR a, VAL v` token pair per nonzero
	byte of machine state, then `STEP`.
	- Every output token is a write. Each instruction becomes a short record of writes such as
	`@V3 =1F @PC =22A STEP`. The screen is drawn from these writes.
	- The host injects the outside world. A `TICK` token arrives once per 60 Hz frame, and keypresses
	are injected as writes. Every other token is the model's greedy prediction.
	- Attention is memory. To read register V3, a head attends to the most recent write to V3.
	Keys use a parabola over the address plus a recency term, so the hardmax lands exactly there.
	- ReLU is the ALU. Addition, comparison, XOR and opcode decoding are all built from ReLUs,
	using identities like step(x) = ReLU(x+1) - ReLU(x).

	The architecture is ordinary:
	- token embedding plus a position embedding from [t, t²]
	- causal softmax attention (β = 1000) with residual connections
	- ReGLU feed-forward layers
	- a linear unembedding
	- no normalization

	`modeling_chipformer.py` contains nothing CHIP-8 specific apart from the host loop, which only feeds
	tokens in and reads the screen back.

	Where the weights come from. They are compiled from a description of the CHIP-8 circuit
	(written in a small DSL in the GitHub repo) rather than learned, which is what makes every
	instruction exact. 139,229 of the 32,328,660 parameters are nonzero.

	## Usage

	```python
	import torch
	from huggingface_hub import hf_hub_download
	from transformers import AutoModelForCausalLM

	repo = "divagr1925/chipformer"
	model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True, dtype=torch.float64)

	rom = open(hf_hub_download(repo, "roms/tetris.ch8"), "rb").read()
	game = model.play(rom, quirks="shift,loadstore,clip,jump", ipf=15)
	for frame in range(120):
	game.frame(keys={5} if frame % 20 < 4 else ()) # CHIP-8 key 5 = left
	print(game.ascii())
	```

	`play_tetris.py` animates the game in a terminal: `python play_tetris.py`.

	The model is a normal `PreTrainedModel`, so you can also call it directly:
	- `model(input_ids, past_key_values=cache)` returns next-token logits over the 8,708-token vocabulary.
	- `game.tokens` (with `keep_tokens=True`) gives you the full trace.

	### Notes

	- Use float64. The memory-read heads compare attention scores of about 10¹³, and these are only
	exact in float64. fp16 and bf16 will break the machine.
	- Speed. Decoding is greedy, one token at a time with a KV cache. On a CPU in PyTorch that is
	about 200 tokens/s, or 2–3 frames per second of Tetris (the game wants 60).
	- `game = model.play(...)` multiplies by sparse copies of the big feed-forward matrices, since only
	0.4% of the weights are nonzero. The outputs are identical; pass `sparse=False` to run the dense
	weights exactly as loaded.
	- The Rust/wasm engine in the GitHub repo uses the same weights with sparse matmuls and exact
	attention shortcuts. It runs the games at full speed in a browser.
	- Context. Positions must stay below 4,194,304 (2²²). `ChipformerGame` restarts the context
	from a state dump every 250,000 tokens, as the engine does.
	- Coverage. The weights implement the original CHIP-8 instruction set with Octo's quirk flags.
	SCHIP/XO-CHIP opcodes and the vBlank display-wait quirk are not supported.

	## Files

	\| File \| \|
	\|---\|---\|
	\| `model.safetensors` \| dense float64 weights (what `transformers` loads) \|
	\| `model.chipformer` \| the same weights, sparse (nonzeros only), for the Rust/wasm engine \|
	\| `configuration_chipformer.py`, `modeling_chipformer.py` \| the architecture plus the host loop \|
	\| `play_tetris.py` \| a terminal demo \|
	\| `roms/tetris.ch8` \| Tetris by Fran Dachille (1991), freely distributable (Revival Studios CHIP-8 pack) \|