CoolFace
Modelpublic

Akshat-Dwivedi/micro-gpt

sourceHugging Faceapache-2.0updated 2mo agoView on Hugging Face
1likes32downloads
Model Card

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 / SettingSpecification
Model TypeDense Decoder-only GPT Transformer
Total Parameters50,039,520 (~50.0M)
Vocabulary Size32,000 tokens (Byte-Level BPE with NFC normalization)
EmbeddingsShared 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 TypePyTorch Causal Scaled Dot-Product Attention
Feed-Forward / MLPSwiGLU activation with a hidden dimension of 1,352
NormalizationRMSNorm (pre-layer normalization)
Positional EmbeddingsLearned absolute position embeddings
Dropout0.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`: 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

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