metadata
language:
- en
license: apache-2.0
tags:
- pytorch
- causal-lm
- llama
- muon
- rotary-position-embeddings
- grouped-query-attention
pipeline_tag: text-generation
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