---
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)