| --- |
| 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](https://github.com/Ronak-jain-afk/Vak) |
| |
