Update config.py

config.py

@@ -7,9 +7,9 @@
 
 # ── Architecture ──────────────────────────────────────────────────────────────
 CONTEXT = 1024
-vocab_size = 8000
+vocab_size = 5000
 D_MODEL = 256+128 
-numberoflayers = 8
+numberoflayers = 10
 numberofheads = 8
 d_Latent = 64
 ffn_mult = 3.5
@@ -24,13 +24,18 @@ 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 = 3e-4      # warmup will auto-scale; works for BF16
+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