Devlog: From bigram to building my own tokenizer

Last time I had a bigram model trained on The Wizard of Oz that mostly made gibberish. Since then ARIA has gone through three versions, and I ended up writing most of the pipeline from scratch.

ARIA-01 / 02: a real transformer

I replaced the bigram with a character-level GPT: 8 layers, 512-dim embeddings, 8 attention heads, 256-character context, trained on web text on my RTX 5060 Ti.

Bumped block_size from 8 to 256 so the model can actually see whole sentences
Added bf16 autocast, which cut VRAM use roughly in half
Found out the hard way that when VRAM fills up, Windows quietly spills into system RAM and training gets about 20x slower. Batch 128 was taking 6 to 9 seconds per step; batch 64 was 0.35 seconds
Fixed a bunch of bugs along the way, like my argparse -batch_size flag secretly setting block_size, and eval running 500 batches every 250 steps

Final pretraining loss was about 1.30 per character, and it started writing real English sentences (that didn’t make sense yet).

Teaching it to chat

The base model only continues text, so I fine-tuned it on Alpaca formatted as User: ... / ARIA: ... conversations, with a lower learning rate so it wouldn’t forget English. Loss went down to about 0.91.

Then I had to fix the chat script: wrapping prompts in the right format, cutting replies at the end of a turn, and adding temperature + top-k sampling so it stopped picking random characters. First real reply to “How do you make a cookie?”:

You make a cookie are both a cookie.

Grammatically broken, but it’s an actual answer, in the right place, that stops on its own. The pipeline works.

Writing my own BPE tokenizer

The biggest limit was being character-level: the model wasted most of its capacity on spelling, and 256 characters isn’t much context. So I wrote a byte-level BPE tokenizer from scratch instead of using a library:

get_stats and merge to find and merge the most common pairs
a training loop with stop conditions (no more merging pairs that only appear once)
encode / decode with a round-trip test, plus save/load for the merges
GPT-2 style regex pre-splitting so tokens can’t cross word boundaries (before this, I was getting tokens like "is to " and "00 MPa ")
a cache in encode that made it about 2.3x faster on small text and much more on big files

Trained on about 1MB of educational text plus chat data, with 2048 tokens. It learned real word pieces like " economic", " School" and "rying", and gets about 3x compression. Since it’s byte-level, there are no unknown characters anymore, so Korean works too.

ARIA-03

I encoded 500MB of fineweb-edu into 169.6M tokens, saved as uint16 binary and loaded with np.memmap. Then I rewrote the model:

Fused attention with F.scaled_dot_product_attention instead of 8 separate heads, about 3x faster per step
Pre-norm blocks and GELU
Weight tying between the embedding and output layer
Warmup + cosine learning-rate schedule
state_dict checkpoints with optimizer state and step count, so a crash doesn’t wipe a run

Test run: loss started at 7.71 (right at ln(2048), which is what random guessing should give) and reached 4.27 after 1000 steps, with train and test matching the whole time.