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

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

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.