Text Generation
ternary
bitnet
1.58bit
tinystories
llama
ternary15M / README.md
brianbellx's picture
Upload README.md with huggingface_hub
770b824 verified
|
Raw
History Blame Contribute Delete
3.84 kB
---
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.