olmputer-tiny

Olmputer logo: a pixel computer

olmputer-tiny is a RISC-V computer (RV32IM with the Zba extension) compiled into the weights of a stock OlmoHybridForCausalLM. Nothing is trained, because a compiler works out every weight. The code ROM, data ROM, RAM and registers live in the recurrent state of the Gated DeltaNet layers. To run a program you decode greedily, and each token the model generates is the next token of the CPU's execution trace.

It is the standalone half of Olmputer, which puts the same layers in front of allenai/Olmo-Hybrid-Instruct-DPO-7B so that OLMo's calc tool calls run on the CPU inside its weights. Both port Qwenputer, which does this with Qwen3.5-4B, to OLMo-Hybrid's head shapes. The compiler is the chipformer package in JGalego/olmputer.

Uses the architecture and head shapes of Ai2's OLMo-Hybrid but none of its weights. Not affiliated with or endorsed by Ai2.

The machine

ISA RV32IM + Zba, plus bignum digit-file instructions
Code ROM 256 KB
Data ROM 128 KB, read-only
RAM 32 KB
Layers 11 Gated DeltaNet + SwiGLU layers, plus the zero attention layer HF's cache needs
Shapes OLMo-Hybrid-7B's: 30 heads of 96 × 192 a layer, MLP 11,008 wide, hidden 2,816
Parameters 3.13B, almost all zero
Vocabulary the trace tokens (values in 15-bit chunks, jumps, stores, putchar, ...) plus PAD

Each row below is one layer and each cell one of its 30 heads, tinted by the part of the machine's memory it holds. Layers 2, 10 and 11 hold no memory, and only their MLPs do any work.

The 11 layers as rows of 30 heads, each head tinted by the memory it holds

Porting from Qwen3.5-4B comes down to head shapes. Qwen's heads are 128 × 128, with two value heads per key head, while OLMo's are 96 × 192 with one of each. Four changes keep every memory read exact:

  • Each state row keeps its ballast in 3 of its 192 columns, where Qwen uses 2 of 128. That way the gated output norm still divides by the same power of 2.
  • The code and data ROMs hold bytes, and reading a byte exactly needs a query scale that is a power of 2. OLMo's 96^-0.5 is not, so the query of the selected row is padded with four rows that are never written to (weights 3 : 3, 2, 1, 1). That makes the scale exactly 1/16.
  • Write gates use OLMo's β = 2 sigmoid(b), so b = 0 writes and b ≤ -256 does not.
  • ROM rows hold 32 words instead of 8, so the 256 KB ROM fits in one layer's 30 heads. The RAM holds one word per row and takes three layers.

Checks

Each program also runs on a reference RISC-V emulator, and its trace is the target. Two things must hold:

  • After the chip layers, the residual stream equals the circuit's IR evaluator on every dimension at every position. This runs in bf16, on traces of up to 60k tokens.
  • Free-running greedy generate, started at the program's entry, reproduces the reference trace token for token.

The runs below use two sets of linear-attention kernels: torch is transformers' own fallback and fla is flash-linear-attention.

Run Kernels Programs Exact Trace tokens
c-qsort torch 1 1 240,175
c-sieve-gmpmul torch 2 2 66,825
short-fla fla 55 55 26,677
short-torch torch 55 55 26,677

Usage

example.py runs fibonacci(25) on olmputer-tiny. Greedy generate runs the CPU, and the script reads the result off the trace the model wrote. It is a uv script, so it brings its own dependencies:

uv run https://huggingface.co/jgalego/olmputer-tiny/resolve/main/example.py
# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "chipformer @ https://github.com/JGalego/olmputer/archive/98639dff4c84a18542e3ca12a286a4df563b3a27.tar.gz#subdirectory=experimental/riscv",  # pylint: disable=line-too-long
#     "huggingface_hub",
#     "torch==2.14.1",
#     "transformers==5.19.0",
# ]
# ///
"""Run fibonacci(25) on olmputer-tiny, the RISC-V computer in its weights.

The program's code and registers are written into the model's Gated DeltaNet states, and greedy
generate then runs the CPU: each token it generates is the next token of the execution trace. The
program ends with mv a0, t0 (the result), li a7, 93 and the exit ecall, so the result is the
second-to-last value record the model wrote.
"""

import torch
from chipformer.q35.ir import OLMO
from chipformer.rv import core, firmware, ref, spec
from chipformer.rv.programs import program
from transformers import OlmoHybridForCausalLM

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
MODEL = OlmoHybridForCausalLM.from_pretrained("jgalego/olmputer-tiny", dtype=torch.bfloat16)
MODEL = MODEL.to(DEVICE).eval()
CIRCUIT = core.build(OLMO)  # the CPU as a circuit: it writes the program into the model's cache

TRACE = firmware.run(MODEL, CIRCUIT, ref.Machine.from_program(program("fib")))
VALUES = [spec.decode_value(TRACE[:i + 1]) for i, t in enumerate(TRACE)
          if spec.family(t)[0] in ("F0", "F1", "F2")]
print("fibonacci(25) =", VALUES[-2])
print(f"{len(TRACE)} trace tokens on {DEVICE}, identical to the reference machine:",
      TRACE == ref.Machine.from_program(program("fib")).run()[1:])

On a GPU it takes under a minute, and it also runs on a CPU, only more slowly. firmware.run feeds the model a one-token prefix, loads the program's ROM, RAM and registers into the recurrent states of the linear-attention layers, starts the CPU with the GO token and decodes greedily until HALT.

Downloads last month
457
Safetensors
Model size
3B params
Tensor type
BF16
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Collection including jgalego/olmputer-tiny