SLM — Small Language Model trained on TinyStories
A GPT-style small language model (~16M parameters) built entirely from scratch using PyTorch and trained on the TinyStories dataset.
This model was built as a learning project to understand transformer architecture, training loops, gradient accumulation, and mixed-precision training from the ground up.
Model details
| Property | Value |
|---|---|
| Architecture | GPT-style decoder-only transformer |
| Parameters | ~16M |
| Layers | 4 transformer blocks |
| Attention heads | 4 |
| Embedding dimension | 256 |
| Context window | 64 tokens |
| Vocabulary | GPT-2 BPE (50,257 tokens) |
| Training data | TinyStories (50k stories subset) |
| Training iterations | 3,000 |
| Optimizer | AdamW (beta1=0.9, beta2=0.95, wd=0.1) |
| LR schedule | Linear warmup -> Cosine decay |
How to use
Install dependencies
pip install transformers torch huggingface_hub
Generate text
import torch
from transformers import GPT2Tokenizer
from modeling_slm import SLMModel
# Load
tokenizer = GPT2Tokenizer.from_pretrained("YOUR_HF_USERNAME/slm-tinystories")
model = SLMModel.from_pretrained("YOUR_HF_USERNAME/slm-tinystories",
trust_remote_code=True)
model.eval()
# Generate
prompt = "Once upon a time there was a little girl"
inputs = tokenizer(prompt, return_tensors="pt")
output = model.generate(inputs["input_ids"], max_new_tokens=200,
temperature=1.0, top_k=50)
print(tokenizer.decode(output.squeeze().tolist(), skip_special_tokens=True))
Training approach
The model was trained using:
- Gradient accumulation (8 steps) to simulate an effective batch size of 128 sequences without needing that much GPU memory at once
- Mixed precision (bfloat16) for faster training and lower memory usage
- Gradient clipping (max norm 0.5) for training stability
- Weight tying between the token embedding matrix and the output projection layer — saves parameters and improves learning
Limitations
- Context window of 64 tokens is short — long prompts will be truncated
- Trained on a small 50k-story subset, so vocabulary and topics are limited
- Output quality is intentionally modest — this is a from-scratch learning project, not a production model
- Not suitable for any real-world application
What I learned building this
- How transformer blocks work internally (attention, MLP, residual connections, layer norm)
- How gradient accumulation simulates large batch training on limited hardware
- How the training loop, optimizer, and LR scheduler fit together
- How to package a custom PyTorch model for HuggingFace Hub
References
- TinyStories paper
- nanoGPT by Karpathy — architecture reference
- Vizuara SLM tutorial — workshop this was built from