sarimahsan101/nova-141m-tinystories

🌟 Nova 141M (TinyStories)

Nova 141M is a custom 12-layer decoder-only Causal Language Model engineered for fast, efficient text generation. It incorporates modern state-of-the-art Transformer architecture features including:

• Grouped Query Attention (GQA) for efficient memory usage
• SwiGLU Activation Functions for improved model capacity
• QK-Normalization for training stability
• Rotary Position Embeddings (RoPE) for relative position encoding
• Tied Word Embeddings (embed & lm_head)

Trained on the TinyStories dataset, Nova 141M produces grammatically correct and coherent short stories while maintaining an ultra-lightweight footprint (~141M parameters).

---

📊 Training Dashboard & Performance

Key Metrics

Metric	Value
Total Training Steps	1,005 steps
Final Validation Loss	2.176
Gradient Norm (Final)	0.999
Total Parameter Norm	300.87
Weight Update Ratio	~0.0038

---

📐 Model Architecture Specs

Architecture Feature	Details
Model Type	Decoder-Only Transformer (TransformerLM)
Total Layers	12 Transformer Blocks
Parameters	~141 Million
Attention Mechanism	Grouped Query Attention (GQA) with QK-Norm
Positional Embeddings	Rotary Position Embeddings (RoPE)
Feed-Forward Network	SwiGLU FFN
Normalization	RMSNorm
Vocabulary Size	50,257 (gpt2 tokenizer)
Format	PyTorch (model.safetensors)

---

💻 Code & Inference

The complete codebase, training scripts, and fast $O(N)$ KV-Cache inference generation script are available in our GitHub repository:

🔗 GitHub Script Location: /generation/generation.py

Quick Example Usage

import sys
import torch
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer
from safetensors.torch import load_file

# 1. Download model & codebase from Hugging Face
REPO_ID = "sarimahsan/nova-14m-tinystories"  
repo_path = snapshot_download(repo_id=REPO_ID)
sys.path.append(repo_path)

from utils.config import load_config
from models.transformer import TransformerLM

# 2. Load Model
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained("gpt2")
config = load_config(f"{repo_path}/default.yaml")
config.vocab_size = len(tokenizer)

model = TransformerLM(config)
state_dict = load_file(f"{repo_path}/model.safetensors", device="cpu")
if "embed.weight" in state_dict and "lm_head.weight" not in state_dict:
    state_dict["lm_head.weight"] = state_dict["embed.weight"]

model.load_state_dict(state_dict)
model.to(device).eval()

# 3. Simple Generation (See /generation/generation.py in GitHub for full KV-cache script)
prompt = "Once upon a time, a little girl named Lily found a shiny key in the garden."
input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)

output_ids = model.generate(input_ids, max_new_tokens=150) if hasattr(model, 'generate') else input_ids
print(tokenizer.decode(output_ids[0]))