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.

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

Full training log: metrics.jsonl.

Files

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).