olmputer-tiny / example.py
jgalego's picture
Upload folder using huggingface_hub
a080b45 verified
Raw History Blame Contribute Delete
1.61 kB
# /// 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:])