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.