config.py

# ─────────────────────────────────────────────────────────────────────────────
# Veylon Alpha 1 — 8M Parameter Config
# ─────────────────────────────────────────────────────────────────────────────
# Target: Fast iteration (3-4 day convergence to reasonable loss on T4 dual)
# Constraint: 15-20 GPU hours/week, fit in T4 16GB VRAM with mixed precision
# ─────────────────────────────────────────────────────────────────────────────

# ── Architecture ──────────────────────────────────────────────────────────────
CONTEXT = 1024
vocab_size = 16000
D_MODEL = 256
numberoflayers = 6
numberofheads = 8
d_Latent = 96
ffn_mult = 3.5
swa_window = 512  # SWA: attend to last 512 tokens (decoder-only, no loss of info at T4 scale)
num_kv_heads = 2   # 4× KV compression via GQA (8 heads / 2 KV heads = 4 query groups)

# ── MoE (optional; leave False if you want pure dense for simplicity) ────────
use_moe = False           # start dense, add MoE later via expert_add.py if needed
moe_num_experts = 8       # if use_moe=True, start with 8 experts (8M → ~9.5M params)
moe_top_k = 2             # route to top 2 experts per token

# ── Training ──────────────────────────────────────────────────────────────────
EPOCHS = 10                # 10 full passes on 160MB → quick baseline, then iterate
batch_size = 16           # 16 per T4; 2 T4s in parallel = effective 32 (fits in 16GB)
learning_rate = 1e-3      # warmup will auto-scale; works for BF16
weight_decay = 0.01       # light L2, prevents drift

# ── Data ──────────────────────────────────────────────────────────────────────
data_path = './training_data'  # where prepare_dataset.py outputs .txt files
PRECISION = 'float16'        # BF16 for T4 (mixed_bfloat16 policy) — no loss scaling needed

# ── Checkpointing & Inference ────────────────────────────────────────────────
USE_REMAT = True          # gradient checkpointing — saves ~30% VRAM
MAX_GEN_TOKENS = 256      # RoPE headroom: context + this = table size