deepgrove / Bonsai

Bonsai: A Small Ternary-Weight Language Model

Model Details

Model Description

Bonsai is a small 500 million parameter ternary weight language model trained by Hespere. Bonsai adopts the Llama architecture and Mistral tokenizer following Danube 3 , with modified linear layers to support ternary weights. The model has been trained primarily using DCLM-Pro and Fineweb-Edu. Bonsai marks a new paradigm of efficiency, being trained in less than 5 billion tokens.

• Developed by: Hespere AI
• Language(s) (NLP): English
• License: Apache-2
• Repository: https://github.com/Hespere/Bonsai
• Paper: https://github.com/Hespere/Bonsai/tree/main/paper/Bonsai.pdf

Usage

Bonsai can be easily used through the Huggingface Transformers library. However, we note that all operations are currently performed in 16 bit precision; we're currently working towards integrating our model design with custom mixed precision kernels. A quick example follows:

from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("hespere-ai/Bonsai", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("hespere-ai/Bonsai", trust_remote_code=True)
text = "What is the capital of France?"
inputs = tokenizer(text, return_tensors="pt")
outputs = model.generate(**inputs, max_length=100)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

We note that Bonsai is not instruction tuned; we highly recommend finetuning the model before usage in a downstream task.

Evaluation

Bonsai achieves competitive performance among its peers, being one of the first ternary models to do so. Evalution results are below; for more detailed results and comparisons to other ternary models, please see the accompanying paper linked above. We use lm-eval for all benchmarks outside of MMLU and lighteval's cloze formulation for MMLU.