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:

python

from transformers import LlamaTokenizer
tok = LlamaTokenizer.from_pretrained("Cion-lab/tinystories-8L-llama3-block")

Usage

python

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

bibtex

@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}
}