Update config.py

config.py

@@ -4,16 +4,15 @@ xla_on_gpu=False
 # ---------------------------------------------------------------------------
 # Core shape -- chosen for GPU efficiency, not just "a number that fits"
 # ---------------------------------------------------------------------------
-CONTEXT = 512
-max_gen_tokens = 512
-yarn_scale = 1.0 # neutralized YaRN extension ramp for clean RoPE sanity checks
+CONTEXT = 1024
+max_gen_tokens = 1024
 gen_headroom = 128
-vocab_size = 8192+4096 
+vocab_size = 16384+8192 
 
-# Missing config fields expected by the shared GPTConfig dataclass and the
-# model factory path. Keeping them explicit here removes the fallback-warning
-# cascade shown in the traceback and stabilizes object construction.
 bias = False
+d_rope = None
+top_k = 64
+moe_top_k = 2
 
 
 
@@ -21,120 +20,59 @@ gradient_checkpointing = False
 num_kv_heads = 4
 pattern = "dense"
 sliding_window_size = 512
-# tensor_parallel_size is the knob the config dataclass reads.
+
 tensor_parallel_size = 1
 
-# YaRN config defaults from the model constructor.
+
 
 mod_capacity = 0.25
-use_mtp    = False # set True to activate
-mtp_depth  = 1      # 1 = predict 2 tokens ahead; 2 = also predict 3 ahead
-mtp_lambda = 0.3 
+use_mtp    = False
+mtp_depth  = 1      
+mtp_lambda = 0.1 
 use_nvfp4 = False
 use_int8 = False
 int8_group_size = 128
-mod_gate_entropy_coeff = 0.01
-# Experimental hardware-specific paths. These are deliberately opt-in and
-# are read live by model.py/prefetch.py so post-import CLI overrides work.
+mod_budget_coeff = 0.05
+mod_gate_entropy_coeff = 0.0
+mod_hard_eval = False
 
-numberoflayers = 16     # 4 complete pyramid_swa groups (4x4) -- no odd
-                           # trailing layer forced into an incomplete group
+numberoflayers = 16     
 numberofheads = 8
-D_MODEL = 384         # multiple of 64 -> head_dim = 192/12 = 16, a clean
-                           # Tensor Core tile-friendly size (fp16 wants
-                           # multiples of 16; 16 divides evenly with no
-                           # padding waste on T4/A100 GEMM tiling)
-      # AdamW-scale LR for a model this size; only
-                           # takes effect if use_static_lr=False below
-ALPHA_COMPILE_MODE="reduce-overhead"
-# ---------------------------------------------------------------------------
-# Attention
-# ---------------------------------------------------------------------------
-num_kv_heads = 4          # 2x GQA compression on global layers (8->4 kv
-                           # heads); pyramid_swa steps kv heads [1,2,4]
-                           # across each group's 3 SWA layers (MQA->GQA),
-                           # verified against the real build_attention_layers
-                           # logic in model.py, not assumed
-
-sliding_window_size = 512  # explicit (was previously left to fall back to
-                           # the default with a warning) -- 1/4 of CONTEXT,
-                           # a reasonable local-attention window for this
-                           # sequence length
+D_MODEL = 384        
+ALPHA_COMPILE_MODE="default"  
+num_kv_heads = 4         
+sliding_window_size = 512  
 use_sdpa=False
 optimizer_8bit = False
 use_wsd = True
-wsd_decay_steps = 700  # or None for auto (10% of MAX_STEPS)
-# ---------------------------------------------------------------------------
-# FFN / MoE
-# ---------------------------------------------------------------------------
-ffn_mult = 8 // 3           # SwiGLU hidden_dim multiplier. model.py rounds
-                           # d_model*ffn_mult UP to the nearest multiple of
-                           # 64 (256 * 8/3 = 682.67 -> 704), avoiding the
-                           # prime-dimension GEMM padding waste a plain
-                           # round() would produce (683 is prime).
-use_moe = False            # plain dense SwiGLU_FFN every layer. AryaSparseMoE
-                           # currently has NO auxiliary load-balancing loss
-                           # (expert_bias is trainable but nothing pushes it
-                           # toward balanced routing) -- leave this False
-                           # unless/until that's added, especially at this
-                           # small a scale where routing collapse compounds
-                           # fastest across depth.
+wsd_decay_steps = 700  
+ffn_mult = 8 // 3           
+use_moe = False            
 n_experts = 8
 n_shared = 1
 # T4 path
 
-# (use_int8 ignored on SM100+ -- nvFP4 takes priority)
-# ---------------------------------------------------------------------------
-# Precision / batching
-# ---------------------------------------------------------------------------
-PRECISION = 'fp16'         # T4 has no bf16 Tensor Cores -- fp16 is correct
-                           # here. On A100+, bf16 is the better choice (same
-                           # speed, wider dynamic range, no GradScaler
-                           # needed) -- change this if/when you move hardware.
+PRECISION = 'fp16'         
 GLOBAL_DTYPE = PRECISION
 USE_ASYNC_LAYER_PREFETCH = True
 intra_op_threads = 2
+num_workers = 2
+GPU_SIDE_BATCHING = False
 DISABLE_MMAP = False
-use_flash_outproj_add_rmsnorm = False
-use_flash_rope = False
-use_flash_swiglu = False
-use_fused_add_rmsnorm = False
+use_flash_ops = False
 use_xsa = False
-batch_size = 64 # starting point, not a measured optimum -- train.py
-                           # now logs `vram X/Y GB` on every step (added
-                           # recently). Watch that number on your first run
-                           # and raise batch_size toward your actual VRAM
-                           # headroom; this value has NOT been validated
-                           # against a real GPU run for this exact config.
-weight_decay = 0.01        # train.py now splits params into decay/no-decay
-                           # groups automatically (biases, RMSNorm gains,
-                           # expert_bias get weight_decay=0.0 regardless of
-                           # this value) -- this number only applies to
-                           # actual weight matrices + the tied embedding.
-
-# ---------------------------------------------------------------------------
-# Schedule
-# ---------------------------------------------------------------------------
-# Defaults to warmup+cosine (NOT static) for this fresh config -- unlike
-# your existing config.py, this file has no prior "it works, don't touch it"
-# history behind a flat LR, so there's no reason to ship it with the known
-# no-warmup risk baked in. WARMUP_STEPS=500 and MAX_STEPS=20000 are
-# hardcoded in train.py's main() (not read from config.py at all) -- change
-# them there directly if this run needs a different horizon.
+batch_size = 32
+weight_decay = 0.01        
 use_static_lr = False
-static_lr = 1e-3        # only used if you flip use_static_lr back to True
+static_lr = 1e-3        
 learning_rate = 1e-3
-# ---------------------------------------------------------------------------
-# Parallelism / misc
-# ---------------------------------------------------------------------------
 
 max_steps = 4000
 data_parallel_only = True
-dropout = 0.0                   # explicit (was previously left to the
-                                  # default with a warning)
+use_fsdp = False
+dropout = 0.0                   
 use_muon = True
-# ---------------------------------------------------------------------------
-# Data
-# ---------------------------------------------------------------------------
 DATA_PATH = './training_data'
 data_path = DATA_PATH
+TOKENIZER_MODEL_PATH = './tokenizer'
+use_streaming_batcher = False