Vak (54.9M Parameters)

Vak is a compact decoder-only transformer trained from scratch on consumer laptop hardware (NVIDIA RTX 4050 Laptop GPU, 6 GB VRAM). It features a modern LLaMA-style architecture with tied embeddings, trained on a 1.0 billion token multi-source corpus using a dual Muon + AdamW optimizer stack and aligned for conversational chat via SFT.

Architecture & Specifications

Parameter	Value
Parameters	54,940,416
Layers	18
d_model	512
Attention Heads	8 (head dim 64)
KV Heads	2 (Grouped-Query Attention 4:1)
FFN Dimension	1,408 (SwiGLU)
Vocab Size	8,192 (Byte-level BPE)
Context Length	1,024
Positional Embedding	RoPE (theta = 10,000)
Normalization	RMSNorm (eps = 1e-5) + QK-Norm
Weight Tying	Input embeddings tied to output head

Repository Structure

• 55m/vak_chat_sft.pt: Final 54.9M conversational SFT chat checkpoint.
• 55m/vak_base_pretrain.pt: Final 54.9M base foundation checkpoint (1.0B tokens).
• 28m/vak_chat_sft.pt: Earlier 28.3M conversational SFT checkpoint.
• 28m/vak_base_pretrain.pt: Earlier 28.3M base foundation checkpoint.
• tokenizer/vak_tokenizer.model: SentencePiece byte-level BPE tokenizer.
• config.yaml: Architecture and training hyperparameters.

Pretraining Dataset

1,000,341,504 tokens across a curated mixture:

• FineWeb-Edu (50%)
• Cosmopedia v2 (25%)
• Python-Edu (12%)
• OpenWebMath (8%)
• TinyStories (5%)

Project Status

Development on Vak concluded upon reaching the practical compute and memory limits of local consumer hardware.

GitHub Repository: Ronak-jain-afk/Vak