regnant-io/kw5-149M

KW5 149M

A Swahili (Kiswahili) base language model, pretrained from scratch on 1.97B tokens across two Kaggle TPU v5e-8 sessions.

Built by Regnant.

149M transformer parameters, 173.6M total. The input embedding and the output head are untied rather than shared, so they count separately and the hub's Safetensors panel reports the larger figure. Both describe the same model.

Instruction-tuned version: kw5-149M-instruct.

This is a base model, though not a naive one. Next-token prediction is what it was trained for, and it is not instruction-tuned. It does partially recognise the chat format, because the cooldown phase oversampled chat-formatted documents 12x and those role-token embeddings really were trained. See Does it follow instructions? below for the measurement, which is more interesting than a plain yes or no.

Read this before loading

The input embedding and the output head are NOT tied, even though the config this run was launched with said they were. The tie broke when the trainer moved the model to the TPU. nn.Module._apply rebuilds a Parameter per entry when it cannot reuse the source storage, and CPU to XLA cannot, so the run trained two independent tensors.

If you rebuild this model with tied embeddings and call load_state_dict, both keys are written into one storage, the second overwrites the first, and your input embedding becomes the output head. Nothing raises. The shapes match. The model then scores roughly the unigram baseline (7.32 nats instead of 3.16) and generates whitespace. This cost the project a month and ~27 TPU-hours of false conclusions; the shipped config.json sets tie_embeddings: false, which is the truth, and modeling_kw5v2.py honours it.

Quick start (runs in Google Colab as-is)

No GPU required; a free CPU runtime works (a few seconds per line), and a T4 is faster. The repository ships its own modeling code, so there is nothing to clone.

!pip install -q huggingface_hub sentencepiece

import sys, torch, sentencepiece as spm
from huggingface_hub import snapshot_download

path = snapshot_download("regnant-io/kw5-149M")   # ~700 MB, cached
sys.path.insert(0, path)                    # this repo ships modeling_kw5v2.py
from modeling_kw5v2 import KW5V2ForCausalLM

device = "cuda" if torch.cuda.is_available() else "cpu"
model = KW5V2ForCausalLM.from_pretrained(path).to(device)
sp = spm.SentencePieceProcessor(model_file=f"{path}/tokenizer.model")

def swahili(prompt, max_new_tokens=60):
    ids = model.generate(
        sp.encode(prompt),                  # no <s>, see below
        max_new_tokens=max_new_tokens,
        temperature=0.2, top_p=0.9, repetition_penalty=1.3,
    )
    return sp.decode(ids)

print(swahili("Tanzania ni nchi"))

Tanzania ni nchi ya Afrika Mashariki na moja kati ya mataifa yenye uchumi mkubwa duniani.
Tanzania inasifika kwa kuwa na maliasili nyingi, lakini pia ina madini mengi ambayo yanaifanya iwe miongoni mwa nchi zenye utajiri wa rasilimali za asili barani Afrika.
Kwa mujibu wa ripoti iliyotolewa hivi karibuni na Shirika la Utafiti wa Madini Duniani (

More prompts. This is a base model, so give it a prefix to continue rather than an instruction:

for p in ["Mji mkuu wa Tanzania ni", "Kiswahili ni lugha", "Elimu ni muhimu kwa sababu"]:
    print(swahili(p, 40))
    print()

Mji mkuu wa Tanzania ni Dodoma.

Mji huo ulianzishwa mwaka 1964 na wakoloni Waingereza kwa jina la "Tanganyika Territory" (sasa: Tanganyika) katika eneo la Ziwa Nyasa, ambalo sasa lina

Kiswahili ni lugha ya Kibantu nchini Angola inayozungumzwa na Wazambia. Mwaka wa 1983 idadi ya wasemaji wa Kiswahili imehesabiwa kuwa watu 20,500. Kufuatana na uainishaji wa lugha za

Elimu ni muhimu kwa sababu ya umuhimu wake katika maisha yetu. Ni lazima tufahamu kuwa elimu ndiyo msingi wa maendeleo, na hivyo basi tunahitaji kufahamu jinsi tunavyotumia maarifa hayo ili tuweze kufikia malengo tuliyojiwekea maishani mwetu.
Katika makala

Sampling is stochastic, so your output will differ. Decoding defaults live in generation_config.json and were chosen by measurement. See Decoding below.

Read those samples carefully: the capital of Tanzania is right, and "Kiswahili ni lugha ya Kibantu nchini Angola inayozungumzwa na Wazambia" is confidently wrong. That is what 149M parameters buys you: fluency, not knowledge.

Do not prepend <s>. Base packing never prepended BOS: every document began with its own first token and ended with </s>. A leading <s> is a format this model has never seen.

A stock AutoModelForCausalLM will not load this correctly: the architecture has Canon layers, which no model in the Llama family has.

Does it follow instructions?

Partly, and it is worth being precise about which part. 12 instructions, each run twice with identical text, once wrapped in the chat format and once as plain continuation, using the same decoder settings:

	stopped cleanly	content-word overlap
chat format	6 / 12	0.26
plain text	0 / 12	0.27

The format is learned. The chat wrapper makes the model terminate, by emitting <|end|> or a role token, about half the time. Plain continuation never terminates, not once in 12 tries. In v1 those ids were reserved and never trained, and prompting with them returned noise.

The register follows too. Same instruction, two formats:

"Eleza kazi ya mwalimu."          (explain a teacher's job)

chat  -> "Mwalimu ni mtu mwenye ujuzi mkubwa ambaye anaweza kusaidia
          wanafunzi katika masomo yao, kutoa ushauri na mwongozo..."
plain -> "- Tambua aina za maneno na misemo katika sentensi.
          - Andika orodha ya majina, alama na vishazi vya kila moja..."

The chat version answers. The plain version continues someone else's grammar worksheet.

But constraints are ignored and facts are unreliable. "Write ONE sentence about the sea" produced several sentences about a lake. "List three things found on a farm" produced teacups and cold water. "Name two colours of the Tanzanian flag" produced red and blue; the flag is green, yellow, blue and black.

So it adopts the answer register most of the time, stops about half the time, and gets the content wrong often. Treat it as a base model with a head start, not as an assistant. The instruction-tuned model is a separate release.

To use the format, build the ids directly. sp.encode("<|user|>") prepends a ▁ piece, which is not the configuration measured above:

USER, ASSISTANT, END = 5, 6, 7

def chat(instruction, max_new_tokens=48):
    nl = chr(10)
    ids = [USER] + sp.encode(nl + instruction + nl) + [ASSISTANT] + sp.encode(nl)
    out = model.generate(ids, max_new_tokens=max_new_tokens,
                         temperature=0.2, top_p=0.9, repetition_penalty=1.3)
    body = [t for t in out[len(ids):] if t not in (USER, ASSISTANT, END, 2)]
    return sp.decode(body)

print(chat("Kwa nini usingizi ni muhimu?"))

Usingizi husaidia mwili kunyonya na kuchakata virutubisho, homoni, nishati,
na kemikali nyingine. Pia husaidia kudhibiti kiwango cha sukari katika damu,
shinikizo la damu, cholesterol ya juu, na hali za kiafya kama vile

Model details

Parameters	173,608,448 stored: 123.9M transformer body, 24.6M input embedding, 24.6M untied output head, 0.5M Canon
Layers / hidden / FFN	20 / 768 / 2048
Attention	Grouped-query: 12 query heads, 3 KV heads, head dim 64
Normalization	RMSNorm, pre-norm, plus RMSNorm on Q and K
Canon layers	ABCD: causal depthwise 1-D convolutions, kernel 4, at four points per block (Allen-Zhu, arXiv 2512.17351)
Activation	SwiGLU
Position encoding	RoPE, theta 204,522 (NTK-rebased from 100,000 during cooldown)
Context length	2048
Vocabulary	32,000 SentencePiece BPE, split_digits, byte_fallback
Embeddings	Untied
Precision	FP32 weights (trained under bf16 autocast on TPU v5e-8)

Tokenizer

32k BPE trained on the v2 mixture. NFC normalization applied before training. Special ids are fixed: <unk>=0, <s>=1, </s>=2, <pad>=3, then 16 user-defined symbols at 4–19 (<|system|>, <|user|>, <|assistant|>, <|end|>, <mawazo>, </mawazo>, <think>, </think>, <final>, </final>, and six reserved), then the 256 byte-fallback pieces at 20–275.

The chat and thinking tokens are present in the pretraining corpus (the cooldown oversampled chat_format 12×), unlike v1 where they were reserved but never trained.

Training

Data	2,618,265,226 packed tokens, 3,737,560 documents, 8 registers: web, news, encyclopedia, translated educational text, parliamentary long-form, religious, mixed long-form, chat format
Tokens seen	1,968,701,440, which is 75.2% of one corpus pass, or ~13.2 unique tokens/parameter
Schedule	Warmup-Stable-Decay, decay completed over the final 10%
Context	1024 during the stable phase, 2048 during cooldown with NTK-aware RoPE rebasing
Cooldown registers	translated educational, encyclopedia, parliamentary, chat format
Optimizer	Muon (matrix parameters) + AdamW (embeddings, norms, head, Canon taps)
Packing	Document-boundary attention masking and label masking
Hardware	Kaggle TPU v5e-8, ~27 TPU-hours
Validation loss	3.1560 on the held-out split

v1, for comparison, saw ~705M unique tokens (6.4 per parameter) and its WSD decay never ran, because of a token-budget constant set from a corpus ceiling that turned out to be 4x too high. This model is the version of that plan that finished.

Evaluation

Validation loss 3.1560 on held-out Swahili. No downstream benchmark results are published here yet. Belebele-sw, AfriXNLI and the rest have not been run on this checkpoint. When they are, they go in this section with their sample sizes and standard errors, not before.

What this model demonstrably does: fluent, well-formed Swahili continuation under greedy decoding, without a repetition penalty. That is a stronger claim than it sounds, because greedy argmax is the decoding mode small models fail first.

Intended use

A starting point for Swahili fine-tuning, whether instruction tuning, domain adaptation or classification heads, where training from scratch is too expensive and large multilingual models are too big to serve.

Not for direct deployment: no instruction following, no safety tuning, no factual reliability.

Limitations

• Not an instruction model. It recognises the chat format and often answers in register, but it ignores constraints, stops only about half the time, and is wrong often. See the measurement above.
• Factual reliability is poor. A 149M model has very little room for world knowledge. It will state confident falsehoods.
• Primarily Tanzanian Swahili, reflecting the corpus.
• No safety tuning at all. It reproduces whatever is in web text.
• Part of the corpus is machine-translated educational text; translation artifacts are present.

Files

file	what it is
model.safetensors	FP32 weights, untied embeddings
modeling_kw5v2.py	Self-contained inference model. Required, because this is not a Llama.
config.json	Architecture, with tie_embeddings: false
tokenizer.model	SentencePiece 32k BPE
generation_config.json	Decoder defaults measured on held-out text, see below
tokenizer_config.json	add_bos_token: false, so do not prepend <s>

Decoding

The shipped defaults were chosen by sweeping temperature and nucleus settings and comparing each setting's distinct-4 n-gram rate against the distinct-4 rate of real held-out Swahili, so the target is the language rather than taste. Too cold and the model loops; too hot and it exceeds the repetition structure of real text.

Measured best: temperature 0.2, top-p 0.90, repetition penalty 1.3.

Citation

@misc{kw5v2base2026,
  title  = {KW5-V2: a 149M-parameter Swahili base model trained on Kaggle TPUs},
  author = {Regnant},
  year   = {2026},
  url    = {https://huggingface.co/regnant-io/kw5-149M}
}

Apache 2.0.