TinyStories 8-layer Llama3-block model
Trained on ~100M tokens of roneneldan/TinyStories using MaxText on a Kaggle TPU v5e-8 accelerator.
Architecture (Llama2/Llama3 decoder block)
| Key | Value |
|---|---|
| Decoder block | llama2 (RMSNorm + RoPE + SwiGLU, causal masking) |
| Hidden layers | 8 |
| Hidden size | 512 |
| Query heads | 4 (head_dim=128) |
| KV heads | 4 (MHA) |
| Intermediate (MLP) | 1376 (8/3 x 512, SwiGLU) |
| Vocab size | 32,000 (bundled Llama2 SentencePiece tokenizer) |
| Max position embeddings | 1024 |
| Tie word embeddings | false |
| Approx parameters | 58M |
Training
- Optimizer: AdamW (b1=0.9, b2=0.95, eps=1e-8, weight_decay=0.1, decoupled)
- LR schedule: cosine, peak 3e-4, warmup 5%, final 0.1x peak
- Gradient clip: global-norm 1.0
- Z-loss: 1e-5 (auxiliary softmax z-loss for stability)
- Skip-step-on-spike: enabled
- Mixed precision: bf16 forward/backward, fp32 master weights & Adam state
- Batch: 8 sequences x 1024 tokens x 8 TPU chips = 65,536 tokens/step
- Steps: 1,526 (~100M tokens trained)
- Hardware: Kaggle TPU v5e-8 (single host, 8 chips, 16 GB HBM/chip)
- Framework: MaxText (JAX + Flax NNX + Optax + Orbax)
Tokenizer
The bundled Llama2 SentencePiece tokenizer from MaxText
(src/maxtext/assets/tokenizers/tokenizer.llama2, vocab 32,000) is included
in this repo as tokenizer.model. Load with:
from transformers import LlamaTokenizer
tok = LlamaTokenizer.from_pretrained("Cion-lab/tinystories-8L-llama3-block")
Usage
from transformers import AutoModelForCausalLM, LlamaTokenizer
import torch
REPO = "Cion-lab/tinystories-8L-llama3-block"
m = AutoModelForCausalLM.from_pretrained(REPO, torch_dtype=torch.bfloat16)
tok = LlamaTokenizer.from_pretrained(REPO)
prompt = "Once upon a time, "
ids = tok(prompt, return_tensors="pt").input_ids
out = m.generate(ids, max_new_tokens=64, do_sample=True, temperature=0.8, top_p=0.9)
print(tok.decode(out[0], skip_special_tokens=True))
License
The model weights are released under CDLA-Sharing-1.0 to match the upstream TinyStories dataset license. The bundled Llama2 tokenizer follows MaxText's Apache-2.0 licensing.
Citation
@misc{tinystories,
title = {TinyStories},
author = {Eldan, Ronen and Li, Yuanzhi},
year = 2023,
howpublished = {https://huggingface.co/datasets/roneneldan/TinyStories}
}
@misc{maxtext,
title = {MaxText: A Simple, Performant and User-Friendly OpenSource LLM training Codebase},
author = {Google},
howpublished = {https://github.com/AI-Hypercomputer/maxtext}
}