pretrain.pth was trained using the
openwebtext 1M dataset, using 100 000 steps and achieving loss of 3.75 (on RTX3060 this took about 9 hours)
haiku.pth was created by finetuning the pretrained model on the
statworx/haiku dataset, using 5 000 steps achieving loss of 2.69 (about 30min)
BATCH_SIZE = 32
BLOCK_SIZE = 256
EMBED_SIZE = 512
NUM_HEADS = 8
NUM_LAYERS = 1
ATTENTION_DROP = 0.0
RESID_DROP = 0.0
LEARNING_RATE = 3e-4
# Finetuning
LEARNING_RATE = 5e-5
