README.md

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