once upon a time tom and lily saw things lily were sad her house he heartd them ilily and tom said yes she saw a little girl smiled tom was so excited her mom said yes
The MOS 6502 is an 8-bit microprocessor released in 1975, powering the BBC Micro and the Apple II. I am lucky enough to have access to my dad's BBC Model B from the 80s; I wanted to see, using modern machine learning, what the strongest language model we could fit on this machine was.
Unsurprisingly, this poses significant challenges. The model weights and inference code need to be contained within 25KB of user-space memory — my final configuration was 9KB inference code and 13KB model weights. The CPU only operates on an 8-bit integer datatype, and doesn't include multiplication in its instruction set.
CC65 is used for the inference code, enabling compilation of C to the 6502 instruction set. A binary of a model trained on my MacBook can then be written to the BBC Micro using PlayUEF and a custom 3.5mm-to-tape cable I DIY'ed. This convinces the BBC that it's listening to a tape drive, while my laptop plays audio out of its headphone jack.
The sim65 emulator allows a parity check between the C inference binary and the reference Python model implementation. The full inference engine can be tested on the jsbeeb emulator before running on the BBC Micro.
You can run it yourself — the link below boots a BBC Micro in your browser, loads the UEF tape image straight from GitHub, and auto-types the commands to run the model. No emulator install required (note that generation takes a few minutes).
Modelling #
The goal of this project is to build an autoregressive language model — a language model that produces tokens one-by-one, similar to frontier language models. The model is a function $f$ that produces the next token from the existing context:
In large-scale language modeling, a 'token' would be a word or sub-word part. For this post, the vocabulary (list of tokens) used will be 26 letters plus the ' ' character. For models on this small scale, a larger vocabulary (eg. word or subword vocab) would lead to the vocabulary encoder / decoder layers consuming too much of the parameter budget.
An embedding layer maps tokens into the hidden dimension (dim=56 in our case) of the model. 3Blue1Brown's video on neural networks is great for understanding how spatial token embeddings work.
$$ g: {\text{a}, \text{b}, ... , \text{z}, '\text{ }'} \to \mathbb R^{56} $$Once tokens are mapped to our high dimensional space, mixing layers are used to model recurrent dependencies between tokens (see recurrent layers).
BitNet
BitNet was introduced as a method for fast inference on CPU. A matrix multiplication $Y = XW$ is a set of dot products of rows of $X$ with columns of $W$:
BitNet quantizes $W$ such that its values lie in the ternary set ${-1, 0, 1}$. This reduces the dot product to a sequence of add/subtract operations:
$$ \begin{align*} Y_{ij} &= X_{i1} W_{1j} + X_{i2} W_{2j} + \cdots + X_{in} W_{nj} \ &= X_{i1} - X_{i2} + ... - X_{in}\ \end{align*} $$The 6502 processor's instruction set doesn't contain multiply: instead, a multiply is built from repeated bit-shift-and-add operations. A single 8×8 multiply and accumulate would cost 150 clock cycles. By contrast a ternary accumulate would take 30 clock cycles, making inference much faster with higher-quantization weights.
Each BitNet parameter takes only $\log_2(3) = 1.58$ bits of storage, compared with 8 bits for int8, 32 for float32, etc. We can pack 4 or 5 parameters per byte: 5 parameters per byte is more data-efficient, but unpacking into ternary values requires repeated floor-divide-by-3 operations — not a native instruction on the 6502, making unpacking costly. By contrast, packing 4 parameters per byte assigns a chunk of 2 bits to each parameter, and unpacking just requires a right-shift. We opt for 4 parameters per byte for inference speed.
In 13KB at 4 parameters per byte, 52k BitNet parameters can be stored. Experimental validation shows that the higher parameter count at lower quantization provides better bang-for-buck than fewer high-precision parameters.
To train BitNet parameters, a similar method to other quantized training is used. The parameter is stored in full float32 precision, quantized to ternary during the forward pass, but gradients flow at full precision in the backward pass:
def ternary_quantize(w: torch.Tensor) -> torch.Tensor:
"""Round to {-1, 0, +1} with straight-through estimator on the backward pass."""
q = torch.clamp(torch.round(w), -1.0, 1.0) # quantize forwards
return w + (q - w).detach() # full precision gradient for backwards pass
In practice the final LM head is kept in int4 — the output projection needs more resolution to spread probability across the vocabulary cleanly. Every other matrix parameter in the model is ternary.
Recurrent Layer Architecture #
Attention
Traditional language models such as GPT-3 used attention for modeling sequential relationships. Attention assigns vectors to each token, and uses a dot-product between each pair of vectors to allow tokens to exchange information. With a dimension of $d$ per token and a total of $s$ tokens in context, a dot product is required between our token and every token in the context window: $O(ds)$ FLOPs total. The inference forward pass (what we'll actually run on the BBC Micro) changes with context size. This is challenging with the 6502, as tensor sizes increase during inference, growing memory usage as we generate tokens. As transformers generate tokens, the KV cache grows by $O(\text{n layer} \times \text{hidden dim})$ per token generated. With our 32KB RAM budget, this would eat up the space we would prefer to use for storing model weights.
What is the benefit that attention provides? It allows exact recall from each token; this allows transformers to perform well at copying and needle-in-a-haystack problems. But do we need this level of precision? For our dataset we just need enough short-term recall to learn how to spell words and produce simple grammar forms.
Recurrent Models
Other architectures have the property that each model forward pass has identical computation shape. State stored by the model is contained within a fixed-size state vector $h$:
$$f: (t_i, h) \to t_{i+1}$$where $t_i$ is token number $i$, and $h$ is a fixed size model hidden state. Both SSMs (eg. S4, Mamba) and RNN-style models (GRUs, LSTMs) satisfy this property, making them much more suitable for constrained-memory inference on the 6502.
Why Not GRU?
Recurrent models are known to be victim to the vanishing/exploding gradient problem, often leading to unstable training. Each step of the recurrence compounds errors by the magnitude of the largest eigenvalue of the forward matrix. An eigenvalue slightly larger than 1 can be catastrophic: as sequence length $s$ increases, errors compound as $\lambda^s, \lambda = 1 + \varepsilon$. In the BitNet regime, the spectral radius of our forward matrices typically far exceeds 1: in order to have a spectral radius of $\approx 1$, we would require that $98%$ of weights are $0$.
These instabilities from the GRU mean that we see divergence in every training run. The only way to avoid these is to store the primary weight matrix in a higher quantization (eg. int4). Since these are the largest matrices of the model, storing weights in int4 massively reduces the possible dimension of the model.
In contrast, Mamba performs a per-channel scalar update, where the decay value per step is computed at inference-time (in the range $[0, 128) / 128$), and so by construction never exceeds 1, making explosion impossible. Mamba is chosen as the target model for this project.
16-Bit Accumulate and Activation Functions
Activations of our model are stored in 8 bit. Each accumulator term is $\leq 128$ in magnitude, so up to 256 terms can be accumulated into 16-bit without overflow. This puts a cap on the model dimension used. After a layer, activations need to be mapped back to 8 bit ready for the next layer. The hardtanh
activation function does exactly this clipping — in float, it is expressed as clip(x, min=-1, max=1)
.
A pure clip of 16-bit values into 8-bit clip(x, min=-128, max=127)
loses most of the dynamic range of the values. Consider accumulating 64 $\text{Uniform}(-128, 127)$ values into 16 bit. The accumulated value has standard deviation $\sigma \approx 591$ — so 83% of values would saturate the clip at -127 or 128. Instead we choose to introduce a learned scaling. Our activation becomes the following:
def activation(x: int16, shr: int) -> int8:
return clip(x >> shr, min=-128, max=127)
The model is highly sensitive to changes in shr
— changing the shr
scaling parameter by 1 doubles / halves the post-activation magnitude. We find that the model performs best if the scale parameters are allowed to vary in the first half of the training run, and then frozen during the second half.
Inference #
Matmul Loop
See below the C implementation of our ternary-by-char matrix multiplication. This primitive is composed to build the Mamba layer; the full set of inference building-blocks can be seen here.
Token Sampling
Greedy sampling of tokens on the 6502 is simple — a max over all logit values. But greedy sampling of language leads to less interesting generation. Top-k softmax sampling requires an exponential function, which we can't natively implement on the 6502:
$$\text{softmax}\left(y\right) = \frac{e^{y_i}}{\sum_j e^{y_j}}$$To perform this in integers, we store a lookup table of precomputed exponential values based on the temperature $T$ to sample with:
In the case of $T=0.9$, $\text{lookup} := [255, 83, 27, 9, 3, 1, 0, 0, ...]$.
We use the standard stability trick of subtracting the maximum value in softmax. If our top 5 logits were $y = [50, 48, 47, 30, 25]$, then $y - \max(y) = [0, 2, 3, 20, 25]$. Using our lookup table gives
This is then sampled using a pseudo-random 16-bit integer taken modulo the sum of the array.
Note that while this is an 'rng', there is no way to inject a random seed into the 6502 (without user input). The seed is therefore fixed, and the model generates the same (pseudo-randomly sampled) output every time.
Results #
The result of this modeling and inference work is a Mamba-based model that actually runs inference on the BBC Micro. While it isn't particularly intelligent, it does demonstrate running a model similar to contemporary models on hardware from the 1980s. Pretty awesome!
This project exemplified the need for mechanical sympathy when designing ML models: ensure modeling decisions are made with hardware in mind. This also aligns with Sarah Hooker's thesis on the hardware lottery — that state-of-the-art architectures and algorithms have evolved alongside the available hardware. If a different set of hardware constraints existed (similar to those we imposed by using the BBC Micro), the architecture or optimizer of choice might be different.