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:

Model Details

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

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

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