---
license: mit
language:
- en
tags:
- 1-bit
- bitnet
- tiny
- language-model
- tinystories
datasets:
- roneneldan/TinyStories
pipeline_tag: text-generation
---

# tiny-tiny-stories

Oh, you think tinystories is small? Welcome to tiny-stories on steroids. A 1m param model at 1.5 bit quantization!

A **1-bit (ternary {-1, 0, +1}) transformer language model** trained on [TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories).

## Specs

| | |
|---|---|
| **Parameters** | 998,784 (< 1M) |
| **Weight precision** | 1.58-bit ternary (BitNet b1.58) |
| **Tokenizer** | SentencePiece unigram, 192 vocab |
| **Context length** | 512 tokens |
| **Best val loss** | 1.2087 (perplexity 3.35) |
| **Training** | 100K steps on 2.1M TinyStories |
| **Checkpoint size** | 3.9 MB (FP32 latent), ~350 KB quantized |

## Architecture

- **d_model**: 128
- **Heads**: 4 (head_dim=32)
- **Layers**: 5
- **FFN**: SwiGLU (d_ff=336)
- **Position encoding**: RoPE (no learned positional embeddings)
- **Normalization**: RMSNorm
- **Embeddings**: Tied input/output, full precision
- **All linear layers**: BitLinear with ternary quantization + straight-through estimator

## How it works

All Q/K/V/O attention projections and SwiGLU FFN matrices use **BitLinear**: weights are quantized to {-1, 0, +1} during the forward pass via `round(W / mean(|W|))`, with gradients flowing through a straight-through estimator to full-precision latent weights during training.

## Usage

```python
import torch
import sentencepiece as spm

# Load tokenizer and model
sp = spm.SentencePieceProcessor(model_file='tokenizer.model')

# Load model (see train.py for BitLM class definition)
from train import BitLM, Config
cfg = Config()
cfg.vocab_size = 192
model = BitLM(cfg)
ckpt = torch.load('best.pt', map_location='cpu', weights_only=True)
state = ckpt['model']
if any(k.startswith('_orig_mod.') for k in state):
    state = {k.replace('_orig_mod.', ''): v for k, v in state.items()}
model.load_state_dict(state)
model.eval()

# Generate
ids = [sp.bos_id()] + sp.encode("Once upon a time")
idx = torch.tensor([ids])
out = model.generate(idx, max_new=200, temp=0.8, top_k=40, eos_id=sp.eos_id())
print(sp.decode(out[0].tolist()))
```

## Sample output

> Once upon a time, there was a squirrel. He was very curious and loved to play in the park. One day, he noticed a big tree in the sky. He was already laughing, but he was stronger under his houses. The squirrel was glue of all the trees, exploring the walls...

## Training

Trained on 2x RTX 2080 Ti using mixed-precision (FP16) with AdamW optimizer, cosine LR schedule (1.5e-3 peak, 1000 step warmup), and gradient accumulation (effective batch size 384).

```bash
python train.py --exp-dir ./output --device cuda:0 --compile
```