README.md

---
license: apache-2.0
language:
  - en
library_name: lwt-inference
pipeline_tag: text-generation
tags:
  - linear-attention
  - gated-linear-attention
  - mixture-of-experts
  - rust
  - rocm
  - from-scratch
---

# LWT-80M

An 80M-parameter language model written from scratch in Rust and HIP — no
PyTorch, no JAX, no ML framework of any kind. Every matrix multiply goes through
hipBLAS; everything else is a hand-written GPU kernel. Trained on a single AMD
Strix Halo APU (gfx1151, RDNA 3.5).

The architecture is not a transformer. Attention is replaced by a **gated linear
recurrence with four independent decay banks**, and the feed-forward layer is a
**mixture of Chebyshev-basis experts**.

- **Code + inference engine:** https://github.com/Linesage/LWT-80M
- **Technical report:** [TECH_REPORT.md](https://github.com/Linesage/LWT-80M/blob/main/TECH_REPORT.md)

## Files

| File | What it is |
|---|---|
| `lwt-80m-base.safetensors` | pretrained base, 732k steps over 6B tokens |
| `lwt-80m-sft.safetensors` | + 1000 steps of Alpaca SFT with prompt-loss masking |
| `tokenizer.json` | byte-level BPE, 32768 vocab |

These are **not** transformers-compatible checkpoints. They load with the
project's own Rust runtime, linked above.

## Usage

```bash
git clone https://github.com/Linesage/LWT-80M
cd LWT-80M
make setup      # install Rust, check ROCm
make download   # fetch these weights
make chat       # build and run
```

Requires ROCm 7.x and an AMD GPU.

## Limitations — read this first

This is a 63M-non-embedding-parameter model trained on a single consumer APU.
Set expectations accordingly.

**The base checkpoint continues text; it does not answer questions.** Ask it
"how do I sort a list?" and you get plausible-looking prose, not an answer. Give
it `def quicksort(arr):` and it writes something Python-shaped. That is the
intended behaviour of a base model, not a defect.

**The SFT checkpoint is experimental.** 1000 steps on ~20k Alpaca examples. It
reliably picks up the response format and learns to stop, and it answers short
factual questions:

```
### Instruction:
What is the capital of France?

### Response:
The capital of France is Paris, France.
```

The redundant trailing "France" is representative. The format is right, the
content is shaky.

**Dialogue quality is limited by both model size and SFT data.** Alpaca is
single-turn, English, and synthetic; there is no multi-turn conversation in the
training data at all, so the model has no notion of dialogue history. At this
scale it also confabulates facts confidently and, at higher temperatures, falls
into repetition loops. Use `--temperature 0.3` and keep the default repetition
penalty (1.15).

**What it is good for:** studying a non-transformer architecture end to end,
inspecting how a gated linear recurrence allocates memory across timescales, and
as a working reference for writing GPU kernels without a framework. It is not a
useful assistant.

## Architecture

Per head, the recurrence is

```
S_t = g_t · S_{t-1} + kᵀ_t v_t
y_t = q_t · S_t
```

`g_t` is a learned forget gate. Because the recurrence is linear, the state
summarises the entire prefix in constant space — **there is no KV cache, because
the state is the cache**. Four banks run in parallel with independently learned
gates; measured half-lives after training are ≈3, 6, 11, 28 tokens, and the
12-layer stack composes them into an effective context far longer than any
single bank.

Each block's feed-forward is 8 experts with top-2 routing, where an expert is a
Chebyshev polynomial basis rather than a SwiGLU MLP:

```
z = tanh(RMSNorm(W_premix · x))
T₁ = z,  T₂ = 2z² − 1
y = W_down · (T₁ ⊙ (W_up · T₂))
```

| | |
|---|---|
| Parameters | 80M total, 63M non-embedding |
| Layers | 12 |
| Model dim | 512 |
| Heads | 8 × 64 |
| Decay banks | 4 |
| Experts | 8, top-2, hidden 320 |
| Vocab | 32768 (byte-level BPE) |
| Training context | 4096 |
| Embeddings | tied |

## Training

| | |
|---|---|
| Hardware | 1× AMD Strix Halo (gfx1151), ROCm 7.2 |
| Data | 6B tokens (24 GB): web, synthetic, code, math/reasoning |
| Steps | 732,000 (one epoch), batch 2 × 4096 |
| Optimizer | AdamW, 3e-4 peak, cosine to 3e-5 |
| Throughput | ~17,000 tok/s, ~4 days wall clock |

A known defect of this run: `max_grad_norm` was left at 1.0 while the raw
gradient norm grew to ~8, so effectively every step after ~50k was clipped and by
the end ~87% of each update was discarded. The technical report has the numbers.

## Inference speed

Decoding is O(1) per token — context length does not affect per-token cost:

| Context | Full window | Incremental | Speedup |
|---------|------------|-------------|---------|
| 512 | 42 tok/s | 191 tok/s | 4.5× |
| 1024 | 23 tok/s | 177 tok/s | 7.7× |
| 2048 | 12 tok/s | 177 tok/s | 14.2× |

## License

Apache-2.0.