Nexa SmallGPT
A GPT-style autoregressive language model built entirely from scratch in PyTorch — no AutoModelForCausalLM, no GPT2LMHeadModel, no pretrained Transformer implementation. Every component (causal self-attention, grouped-query attention, transformer decoder blocks, training loop) is hand-written and independently verified.
Trained on the TinyStories dataset using a single NVIDIA T4 GPU (Google Colab free tier).
Model Details
- Architecture: Decoder-only Transformer (GPT-style) with grouped-query attention
- Parameters: ~13.6M
- Layers: 4
- Embedding dimension: 128
- Attention heads: 4 query heads, 2 KV heads (GQA)
- Context length: 128 tokens
- Feed-forward dimension: 512
- Vocabulary: GPT-2 BPE (50,257 tokens, via
tiktoken) - Tokenizer: GPT-2 BPE (encode/decode only — no pretrained model weights used)
Training
- Dataset: 500,000 stories from TinyStories
- Steps: 30,000
- Batch size: 32
- Optimizer: AdamW (lr=3e-4, weight_decay=0.1)
- Hardware: 1x NVIDIA T4 GPU (Google Colab free tier)
Results
| Metric | Random init | Step 30,000 |
|---|---|---|
| Validation loss | 10.92 | 2.28 |
| Validation perplexity | ~54,900 | 9.60 |
Sample Generation
Prompt: "Once upon a time" (temperature=0.8, top_k=40)
Once upon a time, there was a little girl named Lily. She had a big, soft, fluffy pillow that she loved to sleep safe. One day, Lily's mommy told her to be careful by mistake...
Usage
import torch
from huggingface_hub import hf_hub_download
# Download the checkpoint
ckpt_path = hf_hub_download(repo_id="yadavkapil23/nexa-smallgpt", filename="model.pt")
# Load with the model code from https://github.com/kapilverse/SLM
from config import GPTConfig
from model import SmallGPT
from checkpoint import load_checkpoint
from generate import generate
from tokenizer import Tokenizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SmallGPT(GPTConfig()).to(device)
load_checkpoint(ckpt_path, model, map_location=device)
model.eval()
tok = Tokenizer()
print(generate(model, tok, "Once upon a time", max_new_tokens=100, temperature=0.8, top_k=40, device=device))