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