Ffohturk/Mila_1L

This is a one-layer base model with the LlaMA 2 architecture trained on 6B tokens of the algebraic-stack part of the Proof-pile 2 dataset.
It's output distribution is thus mostly concerned with code. The tokenizer is the LlaMA 2 one. I used the following hyper parameters:
d_model = 512
d_ff = 2048
n_heads = 4
n_ctx = 1024

For the training I used AdamW with weight decay = 0.05 and cosine annealing with 5000 warmup steps and maximum learning rate 1e-4. We used BF16 precision.

Train loss: 2.6228
Test loss: 2.7490