Text Generation
ternary
bitnet
1.58bit
tinystories
llama
File size: 3,842 Bytes
9d70437
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
770b824
 
9d70437
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
---
license: mit
datasets:
  - roneneldan/TinyStories
tags:
  - ternary
  - bitnet
  - 1.58bit
  - tinystories
  - llama
  - text-generation
---

# ternary15M

A 15.19M-parameter Llama-style language model trained from scratch on
[TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories) where **all 42
attention and FFN linear layers are ternary {−1, 0, +1}** with per-output-channel absmean
scales, in the style of BitNet b1.58. Embeddings (tied input/output) and RMSNorm gains
remain FP32.

Training code: [github.com/brianbell-x/ternary15M](https://github.com/brianbell-x/ternary15M)

## Architecture

Mirrors karpathy's stories15M: dim 288 · 6 layers · 6 heads · 6 KV heads · SwiGLU hidden
768 · vocab 32,000 · context 256 · RMSNorm · RoPE · tied embeddings.
**15,191,712 parameters.**

## Files

- `ternary.pt` — deployable hard-ternary checkpoint (43 MB). No latent weights: the 42
  linear layers store only int8 values in {−1, 0, +1} plus one FP32 scale per output
  channel. Verified by reloading from disk and generating text.
- `ckpt.pt` — full latent training checkpoint (182 MB, FP32 latent weights + optimizer +
  RNG, resumable, step 10,000).
- `tokenizer.model` — the Llama 2 SentencePiece tokenizer from
  [karpathy/llama2.c](https://github.com/karpathy/llama2.c) (vocab 32,000).

## Training

- Data: TinyStories (~470M tokens), packed to 256-token sequences with the Llama 2 tokenizer
- 10,000 steps × 65,536 tokens/step = 655M tokens, bf16 autocast, AdamW lr 6e-3 cosine to
  10%, warmup 2%, weight decay 0.1, single L40S GPU (~50 min)
- Straight-through estimator: FP32 latent weights updated by the optimizer; forward passes
  always use the ternarized weights

## Evaluation

| Metric | Value |
|---|---|
| Final validation loss (training) | 1.5895 |
| Validation loss, latent STE | 1.5970 |
| **Validation loss, hard ternary (deployed form)** | **1.6074** |

Hard-ternary inference costs only +0.01 loss over the latent model.

## Usage

Requires the training repo for the model definition:

```python
import sys, torch
sys.path.insert(0, "ternary15M")  # clone of github.com/brianbell-x/ternary15M
from ternary15m.checkpoint import load_hard_model
from ternary15m.runtime import generate_stories, load_tokenizer

model, meta = load_hard_model("ternary.pt", torch.device("cpu"))
tokenizer = load_tokenizer("tokenizer.model")
print(generate_stories(model, tokenizer, torch.device("cpu"), count=1)[0])
```

## Sample output (hard-ternary, CPU)

> Once upon a time, there was a little boy named Tim. Tim loved to bake with his mom. One
> day, they wanted to make cookies for Mom. Tim was very happy.
>
> Tim's mom said, "Let's mix them together." Tim and his mom worked together to mix the
> cookies. Tim was very good at this job. They made the cookies pretty and fun. Tim's mom
> said, "You are a great helper, Tim!"

## Citations

```bibtex
@article{ma2024bitnet158,
  title   = {The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits},
  author  = {Ma, Shuming and Wang, Hongyu and Ma, Lingxiao and Wang, Lei and Wang, Wenhui and Huang, Shaohan and Dong, Li and Wang, Ruiping and Xue, Jilong and Wei, Furu},
  journal = {arXiv preprint arXiv:2402.17764},
  year    = {2024}
}

@article{eldan2023tinystories,
  title   = {TinyStories: How Small Can Language Models Be and Still Speak Coherent English?},
  author  = {Eldan, Ronen and Li, Yuanzhi},
  journal = {arXiv preprint arXiv:2305.07759},
  year    = {2023}
}

@misc{llama2c,
  title  = {llama2.c: Inference Llama 2 in one file of pure C},
  author = {Karpathy, Andrej},
  howpublished = {\url{https://github.com/karpathy/llama2.c}},
  year   = {2023}
}
```

Trained weights released under MIT. TinyStories dataset (Eldan & Li) is MIT-licensed;
the tokenizer file is redistributed from llama2.c under its MIT license.