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.

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

python

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

python

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

Loaded 2,119,719 training stories and 21,990 validation stories from TinyStories
Tokenized using GPT-2 BPE tokenizer (vocab size 50,257)
Concatenated all tokens and grouped into fixed 256-token blocks
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)