Create README.md

README.md

@@ -0,0 +1,133 @@
+---
+license: apache-2.0
+datasets:
+- roneneldan/TinyStories
+- facebook/recycling_the_web
+language:
+- en
+base_model:
+- Akshat-Dwivedi/micro-gpt
+pipeline_tag: text-generation
+tags:
+- gpt
+- micro-gpt
+---
+
+# Micro-GPT
+
+**Micro-GPT** is a compact **50M parameter (50,039,520 exact)** dense, decoder-only GPT model pre-trained on **~1.028 Billion (1,027,551,376)** tokens with a token density of **~20.5 tokens per parameter**.
+
+I pre-trained this model from scratch on a single **NVIDIA RTX 3050 6GB GPU**, keeping memory usage strictly under **~4.5 GB VRAM** while maintaining high throughput (~13,500 tokens/sec).
+
+---
+
+## 📐 Model Architecture Specs
+
+Micro-GPT uses a modern decoder-only Transformer architecture with tied embeddings and VRAM-efficient design choices:
+
+| Feature / Setting | Specification |
+| :--- | :--- |
+| **Model Type** | Dense Decoder-only GPT Transformer |
+| **Total Parameters** | **50,039,520 (~50.0M)** |
+| **Vocabulary Size** | 32,000 tokens (Byte-Level BPE with NFC normalization) |
+| **Embeddings** | Shared token and output head weight matrix |
+| **Context Length (`block_size`)** | 512 tokens |
+| **Transformer Layers (`n_layers`)** | 12 blocks |
+| **Embedding Dimension (`d_model`)** | 480 |
+| **Attention Heads (`n_heads`)** | 8 heads (60 dimensions per head) |
+| **Attention Type** | PyTorch Causal Scaled Dot-Product Attention |
+| **Feed-Forward / MLP** | SwiGLU activation with a hidden dimension of 1,352 |
+| **Normalization** | RMSNorm (pre-layer normalization) |
+| **Positional Embeddings** | Learned absolute position embeddings |
+| **Dropout** | 0.0 |
+
+---
+
+## ⚡ Training & GPU Setup (RTX 3050 6GB)
+
+The training process was optimized to squeeze maximum speed and efficiency out of an RTX 3050 6GB:
+
+- **Hardware**: Single NVIDIA RTX 3050 6GB GPU (peak VRAM usage ~4.5 GB).
+- **Mixed Precision**: FP16 autocast paired with TF32 matrix multiplication for fast training on RTX Series cards.
+- **VRAM Optimizations**:
+  - **Gradient Checkpointing**: Reduces memory footprint during backpropagation.
+  - **Chunked Loss Calculation**: Calculates cross-entropy loss in smaller steps of 256 tokens to prevent allocating large prediction arrays in VRAM.
+- **Optimizer**: Fused AdamW with a weight decay of 0.1 on weight matrices (0.0 on biases and norm weights).
+- **Learning Rate Schedule**: Cosine decay with a 2% linear warmup (1,220 warmup steps). Peak learning rate of 0.0003, gradually decaying to a minimum of 0.00003.
+- **Batching & Throughput**:
+  - Micro-batch size: 4
+  - Gradient accumulation steps: 8
+  - Effective batch size: 32 samples (**16,384 tokens per optimizer step**)
+  - Training speed: **~13,000 to 13,600 tokens/second**
+- **Total Training Steps**: 61,000 steps (~1.028 Billion tokens streamed).
+
+---
+
+## 📊 Dataset Summary
+
+The training data was cleaned using a multi-core parallel cleaner and converted into a binary token stream.
+
+```bash
+================================================================================
+                     DATASET SUMMARY & METRICS
+================================================================================
+ Total File Size                  : 4.66 GB
+ Total Output Dataset Rows        : 3,077,324
+ Final Total Tokens               : 1.0276 Billion (1,027,551,376)
+ Max Tokens per Sample            : 1,024
+--------------------------------------------------------------------------------
+ DATASET BREAKDOWN BY SOURCE:
+
+ 1. TinyStories:
+    • Output Dataset Row Share    : 68.87%
+    • Token Count                 : 0.4496 Billion (449,583,331)
+    • Output Dataset Token Share  : 43.75%
+
+ 2. Recycling The Web:
+    • Output Dataset Row Share    : 31.13%
+    • Token Count                 : 0.5780 Billion (577,968,045)
+    • Output Dataset Token Share  : 56.25%
+================================================================================
+```
+
+---
+
+## 📈 Pre-training Progress & Results
+
+- **Final Loss**: `3.2339`
+- **Final Perplexity**: `25.38`
+- **Final Learning Rate**: `0.00003`
+
+```bash
+[step 100/61,000] loss=9.5585 ppl=14164.71 lr=2.484e-05 grad=2.223 speed=13,251 tok/s eta=00d:20h:54m:57s
+[step 1,000/61,000] loss=5.5282 ppl=251.69 lr=2.461e-04 grad=1.342 speed=12,949 tok/s eta=00d:20h:44m:22s
+[step 10,000/61,000] loss=3.8061 ppl=44.97 lr=2.859e-04 grad=0.896 speed=13,574 tok/s eta=00d:17h:44m:11s
+[step 20,000/61,000] loss=3.5895 ppl=36.22 lr=2.394e-04 grad=0.994 speed=13,300 tok/s eta=00d:14h:18m:23s
+[step 30,000/61,000] loss=3.4506 ppl=31.52 lr=1.729e-04 grad=1.210 speed=13,061 tok/s eta=00d:10h:35m:47s
+[step 40,000/61,000] loss=3.3449 ppl=28.36 lr=1.042e-04 grad=1.343 speed=13,531 tok/s eta=00d:07h:16m:23s
+[step 50,000/61,000] loss=3.2720 ppl=26.36 lr=5.194e-05 grad=1.443 speed=13,064 tok/s eta=00d:03h:48m:17s
+[step 60,000/61,000] loss=3.2304 ppl=25.29 lr=3.019e-05 grad=1.563 speed=13,572 tok/s eta=00d:00h:20m:27s
+[step 61,000/61,000] loss=3.2339 ppl=25.38 lr=3.000e-05 grad=1.585 speed=13,556 tok/s eta=00d:00h:00m:00s
+```
+
+---
+
+## 📁 Codebase Structure & Usage
+
+- [`config.py`](file:///data/projects/sashvat-ai/micro-kimi-k3/micro-gpt/config.py): Single source of truth for model and training configurations.
+- [`model.py`](file:///data/projects/sashvat-ai/micro-kimi-k3/micro-gpt/model.py): PyTorch implementation of the `GPT` model architecture.
+- [`train.py`](file:///data/projects/sashvat-ai/micro-kimi-k3/micro-gpt/train.py): Main training loop with checkpointing and mixed-precision support.
+- [`dataset.py`](file:///data/projects/sashvat-ai/micro-kimi-k3/micro-gpt/dataset.py): Local dataset cleaning, 32k BPE tokenizer training, and binary token streaming pipeline.
+
+### Running the Code
+
+```bash
+# 1. Clean dataset, train 32k BPE tokenizer, and generate token stream
+uv run micro-gpt/dataset.py
+
+# 2. Start pre-training run
+uv run micro-gpt/train.py
+
+# 3. Resume training from checkpoint
+uv run micro-gpt/train.py --resume
+```