config.py

from __future__ import annotations

EPOCHS = 20
CONTEXT = 1024
data_path = "./training_data"
batch_size = 8
learning_rate = 1e-4
weight_decay = 0.01

vocab_size = 8000
D_MODEL = 256
numberoflayers = 8
numberofheads = 4
num_kv_heads = 2
d_Latent = 64
ffn_mult = 3.5
swa_window = 512
use_moe = False
moe_num_experts = 8
moe_top_k = 2

GLOBAL_DTYPE = "bfloat16"   # "float32", "float16", "bfloat16"
MATMUL_PRECISION = "high"   # "default", "high", "highest"
USE_FLASH_ATTENTION = False
USE_SPLASH_ATTENTION = True