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