---
license: apache-2.0
language:
- en
library_name: pytorch
pipeline_tag: text-generation
datasets:
- wikimedia/wikipedia
tags:
- product-key-memory
- memory-layers
- rocm
---

# sparse-memory-lm B-16M

A small Llama-style language model with a product-key memory: a table of 16.8 million learned vectors (6.4B
parameters), of which the model reads a few hundred per token. Per token it uses 33M parameters. It's the B-16M
model from [github.com/re133/sparse-memory-lm](https://github.com/re133/sparse-memory-lm), where the code, the
training setup and all the measurements are.

The point of uploading it: the table doesn't have to be in GPU memory. On my Radeon RX 9070 the model writes about
140 tokens/s with the table memory-mapped from an NVMe SSD, using 0.42 GiB of VRAM.

## Files

| File | What | Size |
|---|---|---|
| `values_q4.bin` | the table, 4 bit, 16,777,216 rows x 384, two values per byte | 3.2 GB |
| `scales_q4.bin` | one fp16 scale per row | 34 MB |
| `rest.pt` | everything else (PyTorch state dict + config), fp32 | 284 MB |
| `hot_rows.npy` | rows sorted by how often they were read in training, for the RAM cache | 134 MB |
| `meta.json` | shapes, the 4-bit format, sha256 of every file | |

The full fp32 table isn't uploaded (25.8 GB). The 4-bit table loses almost nothing: validation PPL 19.98 against
19.96.

## Running it

```bash
git clone https://github.com/re133/sparse-memory-lm.git && cd sparse-memory-lm
python3 -m venv .venv
.venv/bin/pip install torch --index-url https://download.pytorch.org/whl/rocm7.2   # or the CUDA wheel
.venv/bin/pip install -r requirements.txt huggingface_hub
.venv/bin/python scripts/demo_generate.py --download --table nvme    # or ram / vram, -i for your own prompts
```

`--table nvme` reads the table from the SSD, `ram` keeps it in RAM, `vram` on the GPU (~200 tokens/s on the
RX 9070). All three write the same text (with Triton 3.5 the logits are even bit-identical; with Triton 3.8 they can
differ in the last bits). It needs the Triton kernels from the repo, so it doesn't load with
`transformers`.

## Training

- **Data:** 500M tokens of English Wikipedia (`wikimedia/wikipedia`, 20231101.en), GPT-2 tokenizer, every article
  as `Title\n\nText`.
- **Model:** d=384, 12 layers, 6 heads, SwiGLU, RoPE, 1024-token context. Layers 3, 7 and 11 have a memory layer
  instead of the FFN (4 heads, top 32 of 4096² product keys, one table shared by the three layers).
- **Hardware:** one NVIDIA H200, ~3 hours, about 101 GB of GPU memory.

## Results

| | Val PPL (Wikipedia) |
|---|---|
| Same model without the table (21M) | 25.67 |
| B-16M | 19.96 (fp32 table) / 19.98 (this 4-bit table) |

That's about as good as a dense model with ~114M parameters trained on the same data (between 106M and 123M).

## Limitations

- **Small model:** it writes fluent Wikipedia-style English, but most facts in it are made up. It's a research
  model for measuring what such a table costs and brings, not for actually using it.
- **No instruction tuning, no safety tuning.** It just continues Wikipedia articles.
- **Training data:** Wikipedia text, licensed CC BY-SA 4.0. The weights themselves are Apache 2.0.