Mini-GPT TinyStories

A compact GPT-2-style language model (~17.7M parameters) trained from scratch on the TinyStories dataset — a collection of 2.1M synthetically generated short stories designed for small-scale language model training.

Visualize in Trackio

Model Details

Parameter Value
Architecture GPT-2 (decoder-only transformer)
Parameters 17.7M
Layers 6
Hidden dimension 256
Attention heads 8
Vocabulary size 50,257 (GPT-2 BPE)
Context length 256 tokens
Tokenizer GPT-2
Training data TinyStories (2,119,719 stories)
Training framework Hugging Face Transformers + Accelerate

Evaluation Results

Metric Score Interpretation
Validation Loss 1.5902 Cross-entropy on held-out stories
Avg Loss (2K samples) 1.5700 Independent re-evaluation
Perplexity 4.81 Excellent — model confidently predicts next tokens
BLEU 0.1082 Reasonable for open-ended generation
Distinct-1 0.4108 Diverse vocabulary usage (>0.3 = good)
Distinct-2 0.7896 Highly diverse phrase generation (>0.5 = good)

A perplexity of 4.81 is excellent for TinyStories — the model has learned the simple vocabulary and story structure well. For comparison, random prediction would yield perplexity ~50,257 (vocab size).

Usage

Quick Start

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model_id = "bhautikv/mini-gpt-tinystories"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16)
model.to("cuda" if torch.cuda.is_available() else "cpu")

prompt = "Once upon a time, a little girl named Lily"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

output = model.generate(
    **inputs,
    max_new_tokens=150,
    do_sample=True,
    temperature=0.7,
    top_p=0.9,
    repetition_penalty=1.1,
    pad_token_id=tokenizer.eos_token_id,
)

print(tokenizer.decode(output, skip_special_tokens=True))

Batch Generation

prompts = [
    "Once upon a time, a little girl named Lily",
    "The brave knight approached the dragon and",
    "One day, a small bird found a shiny",
]

for prompt in prompts:
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    output = model.generate(
        **inputs,
        max_new_tokens=120,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        repetition_penalty=1.1,
        pad_token_id=tokenizer.eos_token_id,
    )
    print(tokenizer.decode(output, skip_special_tokens=True))
    print("---")

Training Procedure

Training Hyperparameters

Hyperparameter Value
Learning rate 3e-4
Train batch size (per device) 32
Eval batch size (per device) 32
Gradient accumulation steps 2
Total effective batch size 128
Optimizer AdamW (betas=(0.9, 0.999), eps=1e-8)
LR scheduler Cosine
Warmup steps 1,000
Number of epochs 3
Total training steps 43,179
Mixed precision fp16 (Native AMP)
Seed 42
Distributed type Multi-GPU (2× T4)
Total eval batch size 64

Training Results

Training Loss Epoch Step Validation Loss
3.1502 0.0695 1000 3.0154
2.4582 0.1390 2000 2.3299
2.2208 0.2084 3000 2.1053
2.1085 0.2779 4000 1.9927
2.0266 0.3474 5000 1.9238
1.9736 0.4169 6000 1.8746
1.9361 0.4864 7000 1.8391
1.9176 0.5558 8000 1.8104
1.8922 0.6253 9000 1.7867
1.8743 0.6948 10000 1.7685
1.8530 0.7643 11000 1.7514
1.8390 0.8338 12000 1.7365
1.8245 0.9032 13000 1.7248
1.8175 0.9727 14000 1.7126
1.7969 1.0422 15000 1.7037
1.7973 1.1117 16000 1.6945
1.7834 1.1811 17000 1.6860
1.7694 1.2506 18000 1.6797
1.7623 1.3201 19000 1.6723
1.7579 1.3896 20000 1.6650
1.7514 1.4591 21000 1.6569
1.7504 1.5285 22000 1.6508
1.7417 1.5980 23000 1.6460
1.7309 1.6675 24000 1.6404
1.7335 1.7370 25000 1.6354
1.7206 1.8065 26000 1.6304
1.7209 1.8759 27000 1.6256
1.7219 1.9454 28000 1.6209
1.7110 2.0149 29000 1.6172
1.7022 2.0843 30000 1.6134
1.7007 2.1538 31000 1.6100
1.6881 2.2233 32000 1.6062
1.6969 2.2928 33000 1.6029
1.6983 2.3623 34000 1.6006
1.6940 2.4318 35000 1.5985
1.6875 2.5012 36000 1.5960
1.6948 2.5707 37000 1.5946
1.6921 2.6402 38000 1.5927
1.6795 2.7097 39000 1.5918
1.6820 2.7792 40000 1.5912
1.6799 2.8486 41000 1.5906
1.6819 2.9181 42000 1.5902
1.6848 2.9876 43000 1.5902
1.6792 3.0 43179 1.5902

Training Infrastructure

  • Platform: Kaggle Notebook (GPU T4 x2)
  • GPUs: 2× NVIDIA T4 (16GB each)
  • Distributed strategy: Accelerate launch with DistributedDataParallel
  • Experiment tracking: Trackio (Hugging Face)
  • Training time: ~6 hours (including data preprocessing)

Data Preprocessing

  1. Loaded 2,119,719 training stories and 21,990 validation stories from TinyStories
  2. Tokenized using GPT-2 BPE tokenizer (vocab size 50,257)
  3. Concatenated all tokens and grouped into fixed 256-token blocks
  4. Used DataCollatorForLanguageModeling with dynamic padding for training

Intended Uses & Limitations

Intended Uses

  • Educational demonstration of small-scale language model training
  • Generating simple children's stories with coherent structure
  • Baseline for experimenting with GPT-2 architecture modifications
  • Learning resource for distributed training with Hugging Face Accelerate

Limitations

  • Only trained on TinyStories — limited to simple vocabulary and short narratives
  • Small model size (17.7M params) limits complexity of generated stories
  • Context length of 256 tokens restricts long-form generation
  • Not suitable for production use, real-world text generation, or handling complex topics
  • May produce repetitive or nonsensical text outside the TinyStories domain

Framework Versions

  • Transformers 5.14.1
  • Pytorch 2.10.0+cu128
  • Datasets 5.0.0
  • Tokenizers 0.22.2
  • Accelerate (latest)
  • Trackio (latest)
Downloads last month
-
Safetensors
Model size
17.7M params
Tensor type
F32
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Dataset used to train bhautikv/mini-gpt-tinystories

Evaluation results