nexa-smallgpt

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

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="yadavkapil7155/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))