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 = 768
vocab_size =  28000
D_MODEL = 256+128
numberoflayers = 12
numberofheads = 8 
d_Latent = 96
ffn_mult = 2.5 
swa_window = 256+128 # 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 = 20                # 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-4      # 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 = 'bf16'        # BF16 for T4 (mixed_bfloat16 policy) — no loss scaling needed
GLOBAL_DTYPE = PRECISION  # global dtype for model weights, activations, and optimizer

# ── Data Shuffling ────────────────────────────────────────────────────────────
shuffle_samples = True           # shuffle individual samples/lines before tokenization (avoid abrupt domain transitions)
shuffle_seed = 42              # seed for reproducible shuffling (set to different values for different shuffles)
shuffle_buffer_size = 10_000   # number of window indices to hold in shuffle buffer (larger = more randomness, higher memory)
# ── Checkpointing & Inference ────────────────────────────────────────────────
USE_REMAT = False          # gradient checkpointing — saves ~30% VRAM
MAX_GEN_TOKENS = 256 # RoPE headroom: context + this = table size
LR = 3e-4  # if None, will be auto-scaled based on batch size and warmup