Update config.py

config.py

@@ -7,12 +7,12 @@
 
 # ── Architecture ──────────────────────────────────────────────────────────────
 CONTEXT = 1024
-vocab_size = 6000
-D_MODEL = 256+128 
-numberoflayers = 10
+vocab_size = 24000
+D_MODEL = 256+128
+numberoflayers = 9
 numberofheads = 8
 d_Latent = 64
-ffn_mult = 3.5
+ffn_mult = 2.75
 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)
 
@@ -33,10 +33,10 @@ PRECISION = 'bf16'        # BF16 for T4 (mixed_bfloat16 policy) — no loss scal
 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_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
+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