| --- |
| 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. |
|
|
