regnant-io/kw5-lite-base

KW5-Lite: 109.5M Swahili Language Model

A decoder-only transformer language model trained from scratch on 2.1B Swahili tokens

Trained from scratch by the Regnant.io team on a single NVIDIA T4 GPU

---

Table of Contents

• Model Overview
• Key Features
• Model Details
• Architecture
• Training Details
• Usage
• Performance
• Limitations
• Intended Use Cases
• Citation
• Acknowledgments

---

Model Overview

KW5-Lite is a 109.5M parameter decoder-only transformer language model specifically designed for the Swahili language (Kiswahili). This model represents one of the few open-source, from-scratch trained Swahili language models, addressing the critical need for NLP resources in East African languages.

Why KW5-Lite?

• Low-Resource Language Focus: Swahili is spoken by 200M+ people but remains underrepresented in NLP
• Open Source: Fully open model, weights, and training methodology
• Research-Friendly: Designed for fine-tuning and downstream task adaptation
• Resource-Efficient: Trainable on consumer hardware (single T4 GPU)
• Educational: Complete training pipeline available for reproducibility

Quick Stats

Metric	Value
Parameters	109.5M (trainable) / 134M (total with embeddings)
Architecture	12-layer decoder-only transformer
Training Tokens	1.88B (~89% of 2.1B target)
Vocabulary Size	32,000 BPE tokens
Context Length	2,048 tokens
Training Time	~60 hours on single T4 GPU
Languages	Swahili (sw) primary

---

Key Features

• Modern Architecture: RMSNorm, RoPE, SwiGLU activation, GQA-ready design
• Comprehensive Training: 2.1B tokens from deduplicated Swahili web corpus (FineWeb-2)
• Efficient Design: Optimized for single-GPU training with gradient checkpointing
• Interruptible Training: Built-in checkpoint system for Colab session management
• Production-Ready Format: Standard PyTorch model, HuggingFace compatible
• Documented Training: Complete training metrics and methodology available

---

Model Details

Model Specifications

Component	Details
Model Type	Causal Language Model (Decoder-Only Transformer)
Language	Swahili (sw / Kiswahili)
Parameters	109.5M trainable parameters
Layers	12 transformer blocks
Hidden Dimension	768
Attention Heads	12 (multi-head attention)
KV Heads	12 (can be reduced for GQA)
FFN Dimension	2,048 (SwiGLU)
Vocabulary	32,000 BPE tokens (SentencePiece)
Max Context	2,048 tokens
Position Encoding	RoPE (Rotary Position Embeddings) with NTK scaling
Normalization	RMSNorm (pre-norm)
Activation	SwiGLU
Precision	FP32 weights (trained with FP16 mixed precision)

Training Status

Metric	Value
Training Step	8,176
Tokens Seen	1,880,396,800 (~1.88B)
Training Progress	89.3% of 2.1B token target
Training Phase	Stable (WSD scheduler)
Final Learning Rate	~3e-4 (peak)

---

Architecture

KW5-Lite follows a modern decoder-only transformer architecture with several optimizations for efficiency and quality:

Core Components

Input → Token Embedding (32K vocab × 768 dim)
  ↓
12× Transformer Blocks:
  ├─ RMSNorm
  ├─ Multi-Head Attention (12 heads)
  │  └─ RoPE Position Encoding
  ├─ RMSNorm
  └─ SwiGLU FFN (768 → 2048 → 768)
  ↓
RMSNorm → LM Head (tied with input embeddings)

Architecture Highlights

Feature	Implementation	Benefit
Normalization	RMSNorm (pre-norm)	Faster, more stable than LayerNorm
Position Encoding	RoPE with NTK scaling	Relative positions, length extrapolation
Activation	SwiGLU	Better performance than ReLU/GELU
Attention	Scaled Dot-Product (SDPA)	Efficient on T4 (no FlashAttention-2 needed)
Embeddings	Tied input/output	Reduced parameters, better generalization
Precision	FP16 mixed precision	2× memory efficiency, faster training

Design Decisions

• No FlashAttention-2: T4 is compute 7.5 (sm_75), FA2 requires 8.0+. Uses PyTorch SDPA instead.
• FP16 not BF16: T4 lacks BF16 Tensor Cores, uses FP16 with gradient scaling.
• Gradient Checkpointing: Enabled on all layers to maximize batch size on 16GB VRAM.
• 8-bit Optimizer: Uses bitsandbytes AdamW for 4× optimizer memory reduction.

---

Training Details

Hardware & Environment

Component	Specification
GPU	NVIDIA T4 (16GB VRAM, Turing architecture)
Environment	Google Colab (free tier)
Session Duration	5-hour sessions with automatic checkpoint recovery
Precision	FP16 mixed precision (no BF16 on T4)
Total Training Time	~120 GPU hours across multiple sessions

Training Data

Aspect	Details
Primary Source	FineWeb-2 corpus (swh_Latn subset)
Total Tokens	~2.1B Swahili tokens (deduplicated)
Deduplication	Exact hash + MinHash LSH for near-duplicates
Tokenizer	SentencePiece BPE (32K vocab, trained on corpus)
Preprocessing	Unicode normalization, corpus filtering

Training Configuration

Hyperparameters:

Optimizer:          8-bit AdamW (bitsandbytes)
Learning Rate:      3e-4 (peak)
LR Schedule:        Warmup-Stable-Decay (WSD)
  - Warmup:         5% of steps
  - Stable:         85% of steps  
  - Decay:          10% of steps
Weight Decay:       0.1
Gradient Clipping:  1.0
Dropout:            0.0 (not used during stable training)

Batch Configuration:
  Micro-batch:      4 sequences
  Grad Accumulation: 56 steps
  Context Length:   1024 tokens (stable), 2048 (cooldown)
  Effective Batch:  ~230K tokens per step
  
Training Budget:    
  Target Tokens:    2.1B
  Target Steps:     ~9,130 steps
  Completed:        8,176 steps (89.3%)

Memory Optimizations:

• Gradient checkpointing on all layers
• 8-bit AdamW optimizer (bitsandbytes)
• FP16 mixed precision with gradient scaling
• Efficient attention (PyTorch SDPA with memory-efficient backend)

Checkpoint System

The model includes a robust checkpoint system designed for interrupted Colab sessions:

• Automatic checkpoint every 500 steps or 15-20 minutes
• Atomic writes (no partial/corrupt checkpoints)
• Google Drive sync with rotating keep-last-3 policy
• Stores: model, optimizer, scheduler, RNG states, data position

---

Usage

Quick Start with Transformers

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Load model and tokenizer
model_name = "regnant-io/kw5-lite-base"  
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto"
)

# Generate text
prompt = "Tanzania ni nchi yenye"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

outputs = model.generate(
    **inputs,
    max_new_tokens=50,
    temperature=0.1,
    top_p=0.9,
    do_sample=True,   
    repetition_penalty=1.2,
)

generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(generated_text)

Loading with Custom Code

If you want to use the original training code:

import torch
from transformers import AutoConfig, AutoModelForCausalLM

# Load configuration
config = AutoConfig.from_pretrained("regnant-io/kw5-lite-base")

# Load model
model = AutoModelForCausalLM.from_pretrained(
    "regnant/kw5-lite",
    config=config,
    torch_dtype=torch.float16,
)
model.eval()

# Load tokenizer (SentencePiece)
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("regnant-io/kw5-lite-base")

# Generate
prompt = "Elimu ni muhimu kwa sababu"
inputs = tokenizer(prompt, return_tensors="pt")

with torch.no_grad():
    outputs = model.generate(
        inputs.input_ids,
        max_length=100,
        temperature=0.1,
        top_k=50,
        top_p=0.9,
        do_sample=True,
        repetition_penalty=1.2,
    )

print(tokenizer.decode(outputs[0]))

Fine-Tuning with Unsloth

For efficient fine-tuning (recommended for instruction tuning):

from unsloth import FastLanguageModel
import torch

# Load model for fine-tuning
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "regnant/kw5-lite",
    max_seq_length = 2048,
    dtype = torch.float16,
    load_in_4bit = False,  # Set True for further memory saving
)

# Add LoRA adapters for efficient fine-tuning
model = FastLanguageModel.get_peft_model(
    model,
    r = 16,                    # LoRA rank
    lora_alpha = 32,           # LoRA alpha
    lora_dropout = 0.05,
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                      "gate_proj", "up_proj", "down_proj"],
    use_gradient_checkpointing = True,
)

# Now train with your dataset
# ... (see Unsloth documentation for full training loop)

---

Performance

Generation Examples

Example 1: Geography

Prompt: "Tanzania ni nchi yenye"
Output: "Tanzania ni nchi yenye utamaduni mkubwa na watu wenye furaha. 
         Ina historia ndefu ya amani na utulivu..."

Example 2: Daily Life

Prompt: "Watoto wanacheza"
Output: "Watoto wanacheza vizuri, wanaweza kucheza vizuri sana, lakini 
         wanacheza vizuri. Wakati mwingine wanaweza kucheza vizuri sana..."

Example 3: Cultural Context

Prompt: "Mama yangu"
Output: "Mama yangu alikuwa akipika, alikuwa akiniambia hivi, 'Wewe, 
         Mama, nakwambia, usifanye hivyo tena...'"

Perplexity Scores

Evaluated on held-out Swahili text samples:

Text Type	Perplexity	Description
Philosophical	~30	Abstract concepts and reasoning
Formal	~50	News articles and formal writing
Casual	~119	Conversational text and greetings
Narrative	~131	Stories and descriptions

Note: Lower perplexity indicates better language modeling. Scores vary by domain and formality.

Quality Assessment

Strengths:

• Coherent multi-sentence generation
• Proper Swahili grammar and syntax
• Subject-verb agreement
• Knowledge of Tanzanian culture and geography
• Handles various tenses correctly

Limitations (expected for base model):

• Not instruction-following (requires fine-tuning)
• No chat formatting
• Occasional factual errors
• No safety/alignment training

---

Limitations

Base Model Constraints

This is a BASE MODEL trained only on language modeling (next-token prediction). It has NOT been fine-tuned for:

• ❌ Instruction following
• ❌ Chat/dialogue formatting
• ❌ Safety and alignment
• ❌ Specific tasks (QA, summarization, etc.)

Known Issues

1. Not Chat-Ready: The model doesn't understand instruction formats like "Answer the following question:" or follow system prompts.

2. No Safety Filters: May generate:

   • Biased or stereotypical content
   • Potentially harmful or offensive text
   • Factually incorrect information

3. Domain Limitations:

   • Training focused on web text (FineWeb-2 corpus)
   • May perform poorly on specialized domains (medical, legal, etc.)
   • Limited technical/scientific vocabulary

4. Language Mixing:

   • May occasionally mix English words (common in Swahili web text)
   • Limited code-switching control

Recommended Mitigations

• For Production: Fine-tune on instruction/chat data with safety examples
• For Research: Use as-is for language modeling studies
• For Evaluation: Test on domain-specific data before deployment

⚡ Fine-tuning recommended before deployment in any user-facing application.

---

🎯 Intended Use Cases

✅ Recommended Uses

1. Research & Development

   • Low-resource language modeling research
   • Swahili NLP algorithm development
   • Linguistic analysis and studies
   • Educational purposes

2. Fine-Tuning Base

   • Instruction tuning for chat/assistant models
   • Domain adaptation (news, education, etc.)
   • Task-specific fine-tuning (classification, NER, etc.)
   • Multi-task learning experiments

3. Feature Extraction

   • Swahili text embeddings
   • Transfer learning for downstream tasks
   • Cross-lingual studies

4. Prototyping

   • Proof-of-concept Swahili applications
   • Testing NLP pipelines
   • Benchmarking experiments

❌ Not Recommended

• Direct deployment without fine-tuning
• Safety-critical applications
• Medical, legal, or financial advice
• Content moderation (requires safety training)
• Real-time production systems (without optimization)

---

📝 Citation

If you use this model in your research or applications, please cite:

@misc{kw5lite2026,
  title={KW5-Lite: A 109.5M Parameter Swahili Language Model},
  author={Regnant.io},
  year={2026},
  month={August},
  url={https://huggingface.co/regnant/kw5-lite},
  note={Swahili decoder-only transformer trained on 1.88B tokens}
}

Related Work

For more context on low-resource language modeling and Swahili NLP:

• FineWeb-2 Dataset
• SentencePiece Tokenizer
• Scaling Laws for Language Models

---

🙏 Acknowledgments

This project was made possible by:

• Computing: Google Colab free tier (T4 GPU access)
• Data: HuggingFace FineWeb-2 corpus (swh_Latn subset)
• Framework: PyTorch, HuggingFace Transformers, bitsandbytes
• Community: The Swahili NLP and low-resource ML communities

Special thanks to the open-source ML community for tools and resources that make projects like this possible.

---

📄 License

This model is released under the Apache 2.0 License.

You are free to:

• ✅ Use commercially
• ✅ Modify and distribute
• ✅ Use privately
• ✅ Patent use

Requirements:

• Include license and copyright notice
• State changes made to the model
• Include attribution

See the LICENSE file for full details.

---

🔗 Additional Resources

• Training Code: GitHub Repository (Add your repo link)
• Documentation: Technical Report (If available)
• Fine-tuned Versions: Coming soon (instruction-tuned, chat, etc.)
• Discussions: HuggingFace Discussions

---

📧 Contact

For questions, issues, or collaboration opportunities:

• Model Issues: Use HuggingFace Discussions
• Technical Support: Open an issue on GitHub
• General Inquiries: info@regnant.io

---

Empowering East African languages through open-source AI