lilbase

everything before this was small enough to train locally. tinystories is great for learning, but its simple little stories. i wanted a proper base model, one thats actually read a decent chunk of the internet.

that needs real compute, so i rewrote the training code in jax and moved it to a tpu v5e-8 through kaggle, which is 8 of googles ai chips working together.

its llama style, with grouped query attention, rope, rmsnorm and swiglu, and the llama tokenizer with a 32k vocab. it read 6.1 billion tokens of fineweb-edu, which is about the right amount for its size. thats the chinchilla thing, around 20 tokens per param.

params

297m

layers

24

width

1024

query heads

16

kv heads

4

tokens per step

524,288

tokens total

6.1b

steps

11,634

vocab

32k

matmuls

bf16

schedule

1000 warmup, cosine to 10%

held out for eval

first 1500 docs

the chinchilla idea is that a model should read about 20 tokens for every parameter it has. heres where lilbase and sprout land, one bar per model, one tick per token read for each parameter.

each step is 524,288 tokens, data parallel over the 8 chips: 8 chips, times 8 sequences, times 8 accumulation steps, times 1024 tokens. every chip chews through 65,536 of them, then they all agree on one update.

a kaggle session only lasts about 8.4 hours, so the script checkpoints every 500 steps, stops before its time runs out, and the next session picks up from there.

every 250 steps it logs grad norm, learning rate, tokens per second, eta, host ram, peak tpu memory, and a number called q. if q sits near zero, tokenization is the bottleneck and the tpu is waiting on the cpu.

kaggle sometimes gives you a broken tpu with only 1 chip working instead of 8. so now the script just refuses to run unless it sees all 8.

i also tried 16 sequences with 4 accumulation steps instead of 8 and 8. it ran out of memory by like 400mb. so close. now if the first compile overflows anyway, the script halves the micro batch and retries on its own.

every model goes through the same script, zero shot: the first 2000 hellaswag validation examples, 2000 lambada examples and the full arc-easy test set. multiple choice is scored by log likelihood, normalised by the length of each answer, and every sequence starts with an end of text token, the same way training did. the script is in the repo.

it beats gpt-2 small on hellaswag and arc-easy and loses on lambada. to be fair gpt-2 is from 2019, trained on different data, and lilbase is more than twice its size, so its not exactly a fair fight lol.

lambada perplexity came out at 57.37. the learning rate was 4e-4, a bit above the 3e-4 gpt-3 350m used at a similar batch size, so slightly optimistic.

its a base model, so it just continues whatever you give it instead of answering. it gets grammar right and makes up most facts.

ollama run navthings/lilbase "The water cycle begins when"