CoolFace
Modelpublic

fillay/LWT-80M

sourceHugging Faceapache-2.0updated 2mo agoView on Hugging Face
0likes
Model Card

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

FileWhat it is
lwt-80m-base.safetensorspretrained base, 732k steps over 6B tokens
lwt-80m-sft.safetensors+ 1000 steps of Alpaca SFT with prompt-loss masking
tokenizer.jsonbyte-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₂))
Parameters80M total, 63M non-embedding
Layers12
Model dim512
Heads8 × 64
Decay banks4
Experts8, top-2, hidden 320
Vocab32768 (byte-level BPE)
Training context4096
Embeddingstied

Training

Hardware1× AMD Strix Halo (gfx1151), ROCm 7.2
Data6B tokens (24 GB): web, synthetic, code, math/reasoning
Steps732,000 (one epoch), batch 2 × 4096
OptimizerAdamW, 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:

ContextFull windowIncrementalSpeedup
51242 tok/s191 tok/s4.5×
102423 tok/s177 tok/s7.7×
204812 tok/s177 tok/s14.2×

License

Apache-2.0.