Tejascodex001/flame-tinystories-5m

FLAME TinyStories-5M

A 5.87M-parameter decoder-only language model trained from scratch with FLAME (Framework for Lightweight AI Model Engineering), a small Rust + Candle training framework, on the TinyStories dataset of simple children's stories.

See inside it: the interactive visualizer runs this exact model in your browser and shows every intermediate value of the forward pass: https://huggingface.co/spaces/Tejascodex001/flame-visualizer

Architecture

LLaMA-style pre-norm Transformer.

Parameters	5.87M (tied input/output embeddings)
Layers	6
d_model	256
Heads	8 × 32 dims
MLP	SwiGLU, hidden 704
Norm	RMSNorm (ε = 1e-5), pre-norm, final norm
Positions	RoPE (θ = 10000, rotate-half)
Vocabulary	4096 (byte-level BPE trained on TinyStories)
Context	256 tokens

x = tok_emb[ids]
for each block:  x += Wo·Attn(RoPE(Wq·n), RoPE(Wk·n), Wv·n)   n = RMSNorm(x)
                 x += Wdown·(silu(Wgate·m) ⊙ Wup·m)          m = RMSNorm(x)
logits = RMSNorm(x) · tok_embᵀ

Training

Steps	6000
Tokens seen	98.3M
Batch	64 sequences × 256 tokens
Optimiser	AdamW, lr 0.001 → 0.0001 (cosine), 200 warmup steps
Hardware / time	single laptop GPU (RTX 3070 Ti), 1.3 h
Final train loss	1.476
Best validation loss	1.545 (step 6000)

Full training log: metrics.jsonl.

Files

file	contents
model.safetensors	F32 weights, FLAME tensor names (tok_emb.weight, layers.{i}.attn.q_proj.weight, …). SHA-256 eb8ab8d5a83823a8946c7854c4a4a8ddc180ffb3b89ba61a39a03424355f7bf2
config.json	architecture hyper-parameters
flame_tokenizer.json	FLAME's BPE tokenizer ({"kind":"bpe","merges":[...]}); not the HF tokenizers format
metrics.jsonl	training / validation log

Token ids: 0–255 raw bytes, 256 <|endoftext|>, 257 <|pad|>, 258+ BPE merges.

Usage

With FLAME: model.safetensors is the unmodified FLAME checkpoint (runs/tinystories-5m/best/model.safetensors). Inside a FLAME run directory (config.toml + best/model.safetensors, tokenizer path from [data]):

flame generate --run runs/tinystories-5m --prompt "Once upon a time"

In the browser (WebAssembly) with the standalone engine from the visualizer project:

import init, { Engine } from "./flame_engine.js";
await init();
const base = "https://huggingface.co/Tejascodex001/flame-tinystories-5m/resolve/main/";
const [w, c, t] = await Promise.all(["model.safetensors", "config.json", "flame_tokenizer.json"]
  .map((f, i) => fetch(base + f).then(r => i ? r.text() : r.arrayBuffer())));
const engine = new Engine(new Uint8Array(w), c, t);
const logits = engine.logits_last(engine.encode("Once upon a time"));

Limitations

• Trained only on synthetic children's stories: simple English vocabulary, short narratives. It is not a general-purpose or instruction-following model and knows little about the world.
• 256-token context.
• Outputs can be repetitive or incoherent, and may reflect biases of the generated training data.
• Intended for education and research on small language models.

License

Weights and code: MIT. Training data: TinyStories, CDLA-Sharing-1.0 (which does not restrict results produced by computing on the data, such as trained weights).