Olmputer

Olmputer is allenai/Olmo-Hybrid-Instruct-DPO-7B with a small RISC-V computer built into its weights. When OLMo calls its calc tool, the computer runs the calculation inside the forward pass and writes the exact result back into the conversation, in the same tokens OLMo would use itself.

Nothing here is trained. A compiler works out the weights of 11 new layers that implement the CPU (RV32IM with the Zba extension), and they sit in front of OLMo's own 32 layers. OLMo's weights stay as they were, apart from zero-padding to the wider hidden size and a rescale of the norm gains to make up for it.

The idea and the CPU come from Qwenputer, which does the same inside Qwen3.5-4B. Olmputer ports it to OLMo-Hybrid. The compiler is the chipformer package in JGalego/olmputer, and olmputer-tiny is the computer on its own, without OLMo.

Built on Ai2's OLMo-Hybrid. Not affiliated with or endorsed by Ai2.

How it works

The computer's memory lives in the recurrent state of the new layers' Gated DeltaNet heads. It has 256 KB of code ROM, 128 KB of read-only data and 32 KB of RAM. The program it runs is a bignum calculator that handles + - * / // % ** and gcd on integers of up to 62 digits, as well as exact fractions and decimals.

The CPU works by writing an execution trace, one token at a time, at about 1.2 tokens per instruction. These trace tokens only pass through the chip layers. OLMo's own layers, caches and positions never see them, so from OLMo's side the computer is invisible. This takes a custom decoding loop. OLMo's full-attention layers have no output gate that could silence the trace tokens, so plain generate() does not run the computer (see Usage).

When OLMo writes a tool call, the text between <function_calls> and </function_calls> is copied into a buffer that the calculator reads. At </function_calls> the circuit takes over. It runs the calculation, writes the whole environment turn with the result, and then hands the conversation back to OLMo for its answer.

The CPU writes numbers the way OLMo's tokenizer does, in groups of up to three digits. This matters more than you might expect. In a test where base OLMo asked calc for a product and then had to give the result in its answer, it got the exact value 33 times out of 40 when the number came in its own digit groups, and 0 times when it came one digit per token. The conversation above took 2,717 trace tokens, none of which OLMo saw.

Inside the chip layers

Every trace token passes once through the 11 chip layers. 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. The last two layers hold no memory, and only their MLPs do any work.

Checks

Text parity. Adding the chip layers should not change what OLMo says. To check this, the merged model's next-token distributions are compared with base OLMo's on held-out prompts. The yardstick is base OLMo against itself with two attention implementations (sdpa and eager), which differ only by numerical noise. The merge's mean KL divergence from base OLMo is 5.87e-04, against a noise floor of 5.66e-04. Greedy continuations match base OLMo on 12 of 16 prompts, while the two base implementations agree on 11.

The computer inside the merged model. A few programs run through the full merged model, with trace tokens going through the chip layers and the echoed text through every layer. Every token the circuit owns has to match the reference emulator's.

Program	Positions	Circuit-owned	Wrong
calc	651	651	0
c-strops	1,074	1,074	0
rv32ui-lb	193	193	0
rv32um-div	142	142	0

Conversations. OLMo answers questions using the calc tool. Each token the computer writes is checked against the reference emulator as it is generated, and the whole conversation is then replayed through the circuit. 24 of 24 conversations came out clean. There is also a Triton executor in chipformer, which does the chip layers' arithmetic in fp32 and has been checked to match the circuit on every value. On it, the same conversations come out token for token identical, and each takes 6.5 s instead of 157 s.

Arithmetic

Three setups answer the same problems. The first is Olmputer. The second is base OLMo with a calc tool that the host answers by running the same calculator on a reference emulator, so its replies are byte for byte the ones Olmputer writes. The third is base OLMo with no tool at all. There are five sets of 40 problems. arith is a seeded synthetic set, and the other four are modules from the extrapolation split of DeepMind's mathematics dataset. Decoding is greedy, and an answer counts only if the final boxed value is exactly the right number.

Usage

example.py asks Olmputer a question, given as its first argument, and prints the conversation. Every token the computer writes is checked against a reference RISC-V emulator as it is generated, and at the end the whole conversation is replayed through the circuit. It is a uv script, so it brings its own dependencies:

uv run https://huggingface.co/jgalego/olmputer/resolve/main/example.py "What is 2**64 - 1?"

# /// 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",
# ]
# ///
"""Ask Olmputer a question and print the conversation.

OLMo writes a calc call. At </function_calls> the chip layers run the calculator firmware and
write the environment turn with the exact value, and OLMo reads it and answers. The CPU's trace
tokens stay inside the chip layers, so this decodes with chipformer's RvChat rather than
generate(), on its Triton executor when there is a GPU. Every token the computer writes is checked
against a reference RISC-V emulator as it is generated, and the whole conversation is replayed
through the circuit at the end. The question is the first argument.
"""

import sys
import urllib.request
from pathlib import Path

import torch
from chipformer.rv import bridge, elf, ref
from chipformer.rv.chat import RvChat, replay
from huggingface_hub import snapshot_download

QUESTION = sys.argv[1] if len(sys.argv) > 1 else "Use calc: what is 123456789 * 987654321?"
FIRMWARE = Path("calc-olmo.elf")
if not FIRMWARE.exists():
    urllib.request.urlretrieve("https://raw.githubusercontent.com/JGalego/olmputer/"
                               "98639dff4c84a18542e3ca12a286a4df563b3a27/experimental/riscv/"
                               "tests/data/rv/calc-olmo.elf", FIRMWARE)

BOT = RvChat(snapshot_download("jgalego/olmputer"), fast=False,
             executor=torch.cuda.is_available())
IDS = BOT.prompt_ids([{"role": "user", "content": QUESTION}])
MACHINE, REFERENCE = (ref.Machine.from_program(elf.program(FIRMWARE)) for _ in range(2))
bridge.park(MACHINE)  # the firmware waits for a call
bridge.park(REFERENCE)
LOCK = bridge.Session(BOT.T, REFERENCE)
LOCK.history(IDS)
RESULT = BOT.generate(IDS, MACHINE, 600, lock=LOCK)

print(BOT.text_of(RESULT["tokens"]))
print(f"\n{RESULT['text_tokens']} text tokens, {RESULT['chip_tokens']} trace tokens, "
      f"{RESULT['seconds']:.1f} s; checked against the reference: {RESULT['lock_error'] is None}")
print("replayed through the circuit:", replay(BOT, IDS, RESULT["tokens"], FIRMWARE))

The checkpoint is about 32 GB in bf16, so you need a GPU with at least 40 GB of memory. Decoding goes through chipformer's RvChat, which keeps the trace tokens inside the chip layers and runs them on the Triton executor when there is a GPU.