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:

---

⚡ 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.

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

[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: Single source of truth for model and training configurations.
• model.py: PyTorch implementation of the GPT model architecture.
• train.py: Main training loop with checkpointing and mixed-precision support.
• dataset.py: Local dataset cleaning, 32k BPE tokenizer training, and binary token streaming pipeline.

Running the Code

# 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