KAi — a language model trained from scratch on one RTX 3080 Ti

KAi is a 62.8M-parameter decoder-only transformer. Every weight is randomly initialised and trained here: no pretrained checkpoint, no fine-tuning, no LoRA or other adapters. The tokenizer is trained from scratch on the same corpus. Total training: 2.033B tokens in ~13 hours on a single 12 GB consumer GPU that was also driving a desktop.

Try it

ollama run Gminecrafter/KAi1.0

126 MB. Or locally, chat_with_kai.bat — which adds streaming, /context resizing with NTK-aware RoPE scaling, checkpoint switching and Esc-to-interrupt.

Set expectations first: this is a 63M model. It is fluent and it is frequently wrong. See What it can and cannot do — that section is the most useful thing in this repo.

Architecture

A modern Llama-style decoder, sized so a compute-optimal number of tokens fits in a single overnight run on 12 GB.

• RMSNorm, pre-norm — taken in fp32 and cast back, which keeps training stable under bf16 autocast.
• RoPE — no learned position table, so context can be extended later.
• Grouped-query attention (10 q / 2 kv) — shrinks the KV cache 5× at inference for negligible quality cost.
• QK-norm — bounds attention logits, which is what makes Muon’s aggressive learning rate safe.
• SwiGLU feed-forward, no biases, embedding tied to the output head.
• F.scaled_dot_product_attention — PyTorch’s fused kernel. Note: do not switch to enable_gqa=True. It is numerically identical to the current repeat_interleave, but on Ampere it falls back to the math backend, materialises the full T×T matrix, and OOMs.

Why this size

Sized from measurement, not guesswork. scripts/bench.py times one config per process. On this card KAi sustains ~50k tokens/s at ~23 TFLOP/s. A 100M variant measured 34k tok/s, which over the same wall-clock would have left it at ~6 tokens/param — badly undertrained. 63M was the right point on the curve.

The optimisation that mattered most

A (batch·1024, 32768) fp32 logit tensor is ~3 GB — on its own larger than everything else in the step combined. kai/model.py computes the head and loss in 4096-row chunks under torch.utils.checkpoint, so only one chunk of logits is ever live and the rest are recomputed in the backward pass.

A 27× difference. The mechanism is Windows WDDM: exceeding VRAM does not raise OOM, it silently spills to system RAM, so the failure mode is a collapse in speed rather than a crash — a progress bar that stops moving and no error anywhere.

That trap was hit again on 21 Aug from the other direction. Disabling gradient checkpointing takes the trainer from 3.9 GB to 7.7 GB, which looks like it fits in 12 GB. It does not: the desktop holds ~3 GB, the total reached 11.9 GB, and zero steps completed in three minutes. The usable budget on this machine is ~9 GB, not 12.

Training

Three phases, because the corpus changed twice. --resume continues a phase (model and optimiser, same LR schedule); --init-from starts a new one (weights only, fresh optimiser and schedule).

† Not comparable to phase 2 — the anneal corpus is half UltraChat dialogue, which is inherently more predictable than web text.

Total: 2.033B tokens, 32.4 tokens per parameter.

• Optimiser: Muon + AdamW. Muon orthogonalises the momentum update of every hidden 2-D matrix with a 5-step Newton–Schulz iteration; embeddings, the tied head and RMSNorm gains stay on AdamW. Weight decay is deliberately off for the Muon group — its update is already norm-controlled, so a decoupled decay would shrink those matrices toward zero over thousands of steps.
• bf16 autocast, no GradScaler (bf16 keeps fp32’s exponent range).
• Gradient accumulation — micro-batch 32 × 4 = 131,072 tokens per step.
• Schedule — cosine in 1.0. --schedule wsd is now implemented and is the better choice: a constant plateau with decay only over the last 10% means a run can be stopped, resumed or extended without restarting the schedule. Cosine must know the total step count in advance, and every --init-from restarting its decay from peak caused a visible loss bump.
• --compile — torch.compile on each block, in place, measured at 1.24×. It needs Triton, which is not installed by default on Windows: pip install triton-windows. Without it you get torch._inductor.exc.TritonMissing.

Data (all public and permissively licensed)

Tokenised into a single interleaved stream, so any slice carries the target mixture and the training loader does no mixing at all.

Phase 2 mixture: FineWeb 54.6%, UltraChat 12.9%, UltraChat-gen 11.8%, TinyStories 10.1%, coding 7.1%, everyday 2.5%, identity 1.2%.

Instruction data is mixed into pretraining rather than applied as a separate fine-tuning stage — KAi is a single from-scratch run.

The most expensive data lesson: repetition, not proportion. Identity data was 1.2% of the corpus, which sounds harmless. But the underlying file was only 1.32M tokens, so reaching 24M meant cycling it 18× — the model saw those strings ~11 times against FineWeb’s 0.58. Memorised strings are what a model emits when it is unsure, so a handful of canned lines became the answer to anything it could not parse, including coding requests. Cap any source at 2–3×.

Results

Measured with scripts/eval_lm.py (HellaSwag: 4-way ending choice, random = 25%; LAMBADA: predict a passage’s final word, needs long-range context).

RESULTS_TABLE_PLACEHOLDER

What it can and cannot do

The interesting output of this project. At 63M the failures are systematic, not random, and each has an identifiable cause.

It has the form of an answer before the content. “The capital of France is Paris” is stable and correct across resets. “The capital of Spain is Paris, which is located in central France” is the same mechanism failing: it retrieves the most frequent city associated with the pattern, not the capital. Asked for Australia’s capital it says Sydney. The confident right answers and the confident wrong ones are produced identically.

Errors compound. Each generation step conditions on the text already written, so one wrong noun becomes evidence for the next. Asked where Spain is, it produced a numbered list of Spanish regions in which Italy leaked in at item two and was item four. Practical rule: the first clause is usually right and everything after it is invention — short answers from this model are far more trustworthy than long ones, which is the opposite of what the fluent prose suggests.

Instructions lose to topic words. “Write a short essay about AI” produces an essay. “Write a short essay about the sun” produces a story about a cat, because sun lives in TinyStories and the noun outvotes the verb. Likewise code a code to reverse a string fails while write a code to reverse a string in python. Keep it simple returns s[::-1] — the knowledge was there and the prompt could not reach it.

There is one register called “code”, not several languages. Asked for JavaScript it emitted Java, C, Pascal and Python syntax in a single function.

What works: fluent English; holding a topic for several turns; recalling a name you gave it earlier; and conditioning on constraints that are content (“it’s very rainy” changes shoe advice) rather than about its own output (“maximum 900 words” never works).

QUOTES.md collects the best examples with the mechanism behind each.

Layout

kai/model.py               transformer + chunked-CE loss + KV-cache sampling
kai/muon.py                Muon optimiser + parameter grouping
kai/corpus.py              document iterators + chat template
scripts/download_data.py   fetch the raw corpora
scripts/train_tokenizer.py train the 32k BPE
scripts/prepare_data.py    tokenise + interleave -> data/tokens*/
scripts/train.py           the training run
scripts/chat.py            talk to the trained model
scripts/eval_lm.py         HellaSwag + LAMBADA
scripts/prompt_suite.py    fixed-seed A/B between checkpoints
scripts/bench.py           throughput / VRAM probe
scripts/export_hf.py       -> HuggingFace, with an equivalence assert

Companion documents: PROGRESS.md (build log and what testing established), QUOTES.md (failure modes with examples), KAI_1.1.md (plan for the successor), CLAUDE.md (notes for anyone working on the code).

Reproducing

python -m venv .venv
.venv/Scripts/pip install torch --index-url https://download.pytorch.org/whl/cu128
.venv/Scripts/pip install tokenizers tqdm numpy pyarrow huggingface_hub
.venv/Scripts/pip install triton-windows      # optional, for --compile

python scripts/download_data.py --fineweb-shards 2
python scripts/train_tokenizer.py --sample-mb 1200 --vocab-size 32768
python scripts/prepare_data.py --preset 2b --total-tokens 2.0e9 --out data/tokens_2b

echo go > ALLOW_TRAINING                      # training refuses to start without it
python scripts/train.py --epochs 28 --steps-per-epoch 250 \
    --micro-bs 32 --accum 4 --data-dir data/tokens_2b --schedule wsd
python scripts/train.py --resume              # after any interruption

Training exits with code 3 unless an ALLOW_TRAINING file exists, and watchdog.ps1 kills unauthorised training processes — the GPU is shared with a live desktop, so a run must never start on its own.

Export

export_to_ollama.bat        # checkpoint -> HF -> GGUF -> ollama create

GGUF export goes through Qwen3, not Llama: KAi has QK-norm, which plain Llama lacks and Qwen3 has. scripts/export_hf.py asserts the exported model is numerically identical before writing — measured max logit difference 0.000e+00. llama.cpp identifies BPE pre-tokenizers by a hash table and rejects unknown ones, so scripts/patch_llamacpp.py registers KAi’s hash.

License

Copyright © 2026 Gminecrafter Distributed under the MIT License. You can copy and use this however you like, but you must include this copyright credit notice in all copies.