---
license: mit
language:
  - en
library_name: flame
pipeline_tag: text-generation
tags:
  - flame
  - rust
  - candle
  - tinystories
  - llama-style
  - tiny-llm
  - educational
datasets:
  - roneneldan/TinyStories
---

# FLAME TinyStories-5M

A 5.87M-parameter decoder-only language model trained from scratch with
[FLAME](https://github.com/RTensorLabs/FLAME) (Framework for Lightweight AI Model Engineering),
a small Rust + [Candle](https://github.com/huggingface/candle) training framework, on the
[TinyStories](https://huggingface.co/datasets/roneneldan/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]`):

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

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

```js
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](https://huggingface.co/datasets/roneneldan/TinyStories), CDLA-Sharing-1.0 (which does not restrict results produced by computing on the data, such as trained weights).