zachwallace/bonsai-sapling-v1-290m-hf

About this model and the process

Hi guys, I'm a high school student with a strong interest in AI and ML, and I wanted to create a model myself to better understand the process. This model is a 290 million parameter decoder-only transformer, and took around 30 days to pretrain on my Macbook Pro M5 Pro with 48GB of RAM. I wanted to make this project the most 'from scratch' as possible to do by myself, so the transformer, BPE tokenizer, data pipeline, and AdamW optimzer are all original implementations.

Parameters	290,497,536
Layers / d_model / heads	20 / 1024 / 16
Context length	1024
Vocabulary	32,768 (custom byte-level BPE)
Training tokens	6.20B (Chinchilla-optimal)
Training time	29.3 days, on a M5 Pro (48 GB)
Held-out perplexity	17.64

This does not load with transformers

There is no config.json and no modeling_*.py, so AutoModel.from_pretrained will not work. The architecture is hand-written and ships in this repo , clone it and you have everything:

git clone https://huggingface.co/zachwallace/bonsai-sapling-v1-290m
cd bonsai-sapling-v1-290m
pip install mlx numpy
python3 generate.py --prompt "The water cycle is"
python3 generate.py --chat --prompt "Why is the sky blue?"

Requires Apple Silicon (MLX). Or load it yourself:

import mlx.core as mx
from mlx.utils import tree_unflatten
from model.config import target
from model.transformer import Transformer
from tokenizer import BPETokenizer

cfg = target()
model = Transformer(cfg)
model.set_dtype(mx.bfloat16)
# mx.load returns a flat dict of dotted keys; update() needs a nested tree.
model.update(tree_unflatten(list(mx.load("weights.safetensors").items())))
model.eval()

tok = BPETokenizer.load("vocab.json")
prompt = mx.array([tok.encode("The water cycle is")])

for t in model.generate(prompt, max_new_tokens=200, temperature=0.75,
                        top_p=0.9, repetition_penalty=1.15, stop_id=tok.eot_id):
    print(tok.vocab[t].decode("utf-8", errors="replace"), end="", flush=True)

vocab.json is required, the token ids are meaningless without the exact vocabulary that produced them.

Files

file	what it is
weights.safetensors	base model - bfloat16, 182 tensors, 581 MB
weights-sft.safetensors	recommended variant — same shape, 581 MB
vocab.json	the byte-level BPE vocabulary (required), 792 KB
generate.py	standalone sampler — run it directly
model/	the transformer: RoPE, attention, RMSNorm, SwiGLU, sampling
tokenizer/	the byte-level BPE implementation
sft/	chat template and loss masking

Two variants

Base is a completion model. It continues text rather than following instructions.

SFT(Recomended) was instruction-tuned for one hour on Dolly-15k plus 2,852 single-turn OpenAssistant pairs. It answers rather than continues, stops on its own, and picks prose or list format. Chat markers use reserved token ids 257 (user) and 258 (assistant), so vocab_size is unchanged between the two:

<|user|> your question <|assistant|>

Training

Data. 28 GB of FineWeb-Edu → 6.20B tokens. Own filtering (Gopher/C4-style heuristics), own dedup (exact hashing plus MinHash/LSH near-duplicate detection at a 0.771 Jaccard threshold), own packing into a flat uint16 stream. 6,342,000 documents in, 95.2% kept.

Tokenizer. Byte-level BPE trained from scratch on a 1 GB sample — 32,768 tokens, 4.553 bytes/token on held-out text. No sentencepiece, no tokenizers.

Optimization. Hand-written AdamW: β=(0.9, 0.95), weight decay 0.1, gradient clipping at 1.0, peak LR 6e-4 with 2% warmup and cosine decay to 10%. Batch 8×1024 with 16-step gradient accumulation = 131,072 tokens per optimizer step, 47,279 steps. bfloat16 parameters with float32 master weights and moments — without the master copy, end-of-schedule updates at 3e-5 fall below bfloat16's resolution near a 0.02 weight and round away entirely.

Limitations

It is fluent and frequently wrong. The model is well versed in regular speech, but sometimes goes off on tangents, or states false information.

On "from scratch"

Written from the mathematics: RoPE rotation math, causal multi-head attention, RMSNorm, SwiGLU, pre-norm blocks, byte-level BPE, AdamW with bias correction and decoupled weight decay, data filtering/dedup/packing, the training loop, KV-cached inference. The model imports exactly five symbols from mlx.nn — Module, Linear, Embedding, silu, gelu — and tests prove each is plain arithmetic. The full forward pass is independently recomputed from raw array operations and reproduces the model's logits bit-identically.

Not claimed: the architecture is a standard decoder-only transformer, essentially the Llama recipe. I implemented it from the papers. MLX's autograd computes the backward pass. The corpora are existing public datasets.

Licensing

The card is tagged cc-by-sa-3.0 because Dolly-15k is CC-BY-SA-3.0 and the SFT variant is trained on it — share-alike is the most restrictive input, so it governs. Upstream: FineWeb-Edu is ODC-By 1.0, OASST1 is Apache-2.0.

The base checkpoint touches no Dolly data and could reasonably carry a more permissive license if you split the repos. Licensing of model weights trained on licensed data is unsettled; pick deliberately rather than inheriting this tag by default.