lakhera2023/gemma4-nano-tinystories

Gemma 4 Nano — TinyStories (37M parameters)

A 37 million parameter language model built entirely from scratch using the Gemma 4 architecture, trained on the TinyStories dataset to generate children's stories.

This is not a fine-tuned model. Every layer — attention, feed-forward, normalization, positional encoding — was implemented from the ground up in PyTorch, following the Gemma 4 design principles.

Try the live demo

Gemma 4 Architectural Features

Despite being ~8,000x smaller than the full Gemma 4 (31B), this model implements 9 core architectural innovations:

Feature	Description
Dual Head Dimensions	head_dim=48 for sliding attention, global_head_dim=96 for full attention
Proportional RoPE	theta=10,000 for sliding layers, theta=1,000,000 with partial_rotary_factor=0.25 for global layers
Shared KV Cache	Last 6 layers reuse key/value states from earlier donor layers
K=V Attention	Global attention layers share key and value projections
Value Normalization	RMSNorm (without learned scale) applied to value states
Logit Softcapping	tanh-based capping at 30.0 to prevent extreme logit values
Embedding Weight Tying	Input embedding and output head share the same weight matrix
QK Normalization	RMSNorm on queries and keys; attention scaling = 1.0
Zero-centered RMSNorm	Weights initialized to 0, applied as (1 + weight) * normalized

Model Details

Parameter	Value
Total parameters	37.4M
Embedding dimension	384
Layers	20 (17 sliding + 3 full attention)
Attention heads	8
KV heads (sliding)	2 (GQA group size = 4)
KV heads (global)	1 (GQA group size = 8)
Feed-forward dim	1,152 (GeGLU)
Sliding window	512 tokens
Context length	2,048 tokens
Vocabulary	8,000 (custom SentencePiece BPE)

Layer Pattern

3 groups of (5 sliding + 1 full) + 2 extra sliding = 20 layers:

S S S S S F | S S S S S F | S S S S S F | S S

Usage

Quick Start

import torch
import torch.nn as nn
import torch.nn.functional as F
import sentencepiece as spm
from huggingface_hub import hf_hub_download

# Download files
repo = "lakhera2023/gemma4-nano-tinystories"
ckpt_path = hf_hub_download(repo_id=repo, filename="pytorch_model.bin")
tok_path = hf_hub_download(repo_id=repo, filename="tinystories_tokenizer.model")

# Load tokenizer
sp = spm.SentencePieceProcessor()
sp.load(tok_path)

# Load model (requires the model class definitions — see inference.py in this repo)
model = Gemma4Model(CONFIG)
state = torch.load(ckpt_path, map_location="cpu", weights_only=True)
cleaned = {k.replace("_orig_mod.", ""): v.float() for k, v in state.items()}
model.load_state_dict(cleaned, strict=False)
model.eval()

# Generate
prompt = "Once upon a time there was a little rabbit"
ids = torch.tensor([sp.encode(prompt)], dtype=torch.long)
output = model.generate(ids, max_new_tokens=200, temperature=0.8, top_k=50)
print(sp.decode(output[0].tolist()))

Standalone Inference

Download and run inference.py from this repo:

pip install torch sentencepiece huggingface_hub
python inference.py --prompt "Once upon a time" --max_tokens 200 --temperature 0.8

Training Details

Setting	Value
Hardware	NVIDIA A100 40GB
Dataset	roneneldan/TinyStories
Tokenizer	Custom SentencePiece BPE (8,000 vocab, trained on TinyStories)
Precision	bfloat16 with torch.compile
Optimizer	AdamW (lr=3e-4, betas=(0.9, 0.95), weight_decay=0.1)
Scheduler	Linear warmup (500 steps) + cosine decay to 1e-5
Batch size	32 (effective 64 with gradient accumulation = 2)
Sequence length	1,024 tokens
Total steps	50,000
Effective tokens/step	65,536

Tokenizer

A custom 8,000-vocabulary SentencePiece BPE tokenizer trained specifically on TinyStories. Using a domain-specific tokenizer instead of GPT-2's 50,257-token vocabulary saved ~16M embedding parameters — critical for keeping the model under 40M total.

Files:

• tinystories_tokenizer.model — SentencePiece model file (load with spm.SentencePieceProcessor)
• tinystories_tokenizer.vocab — Human-readable vocabulary

Limitations

• Domain-specific: Trained only on children's stories. Will not perform well on other text domains.
• Small vocabulary: The 8K tokenizer is optimized for TinyStories English text. It may struggle with uncommon words, names, or non-English text.
• No instruction following: This is a base language model, not fine-tuned for instructions or chat.
• Short context: Best results with prompts under ~500 tokens; maximum context is 2,048 tokens.
• Custom architecture: This model uses non-standard tensor names and architecture details. It cannot be converted to GGUF/ONNX or run with llama.cpp/Ollama directly.

Files

File	Size	Description
pytorch_model.bin	74.9 MB	Model weights (PyTorch state_dict, bfloat16)
tinystories_tokenizer.model	371 KB	SentencePiece tokenizer model
tinystories_tokenizer.vocab	110 KB	Human-readable vocabulary
config.json	~2 KB	Full model configuration
inference.py	~8 KB	Standalone inference script

Citation

If you use this model, please cite the original Gemma 4 paper and the TinyStories dataset:

@article{tinystories2023,
  title={TinyStories: How Small Can Language Models Be and Still Speak Coherent English?},
  author={Eldan, Ronen and Li, Yuanzhi},
  journal={arXiv preprint arXiv:2305.07759},
  year={2023}
}

License

Apache 2.0