Audit fix: Update model/configs/train_rtx4090.yaml for ViuMini-Dense-360M

model/configs/train_rtx4090.yaml

@@ -1,80 +1,70 @@
-# Pretraining config — VIU-1.5B MoE 48L (1.856B total / 300.5M active)
-# OPTIMIZED FOR SINGLE NVIDIA GEFORCE RTX 4090 (24GB GDDR6X, Ada Lovelace)
+# Pretraining config — ViuMini-Dense-360M (32L Dense, 366.6M params)
+# OPTIMIZED FOR SINGLE / MULTI NVIDIA GEFORCE RTX 4090 (24GB GDDR6X, Ada Lovelace)
+model_config: ../configs/viu1_dense_config.yaml
 seed: 42
 data_dir: ../../data/raw
 tokenizer_path: ../../tokenizer/outputs/tokenizer.json
+
 # Sequence Length Curriculum:
-# Stage 1: 0 to 50k steps (seq 512, ~2.5x higher throughput to rapidly master vocabulary & syntax)
-# Stage 2: 50k to 200k steps (seq 1024, mid-length reasoning & conversations)
-# Stage 3: 200k+ steps to completion (seq 2048, full long-context mastery)
-# Global effective batch is kept strictly constant at 262,144 tokens per step across all stages!
+# Constant effective batch: 131,072 tokens per step
 curriculum: true
 curriculum_stages:
   - max_step: 50000
     seq_len: 512
-    batch_size: 32
-    grad_accum: 4  # 512 * 32 * 4 = 65,536 tokens/step
+    batch_size: 16
+    grad_accum: 16  # 512 * 16 * 16 = 131,072 tokens/step
   - max_step: 200000
     seq_len: 1024
-    batch_size: 16
-    grad_accum: 4  # 1024 * 16 * 4 = 65,536 tokens/step
+    batch_size: 8
+    grad_accum: 16  # 1024 * 8 * 16 = 131,072 tokens/step
   - max_step: null
     seq_len: 2048
-    batch_size: 8
-    grad_accum: 4  # 2048 * 8 * 4 = 65,536 tokens/step
+    batch_size: 4
+    grad_accum: 16  # 2048 * 4 * 16 = 131,072 tokens/step
 
-seq_len: 512
+seq_len: 1024
 train_split: 0.95
 
-# RTX 4090 24GB (Ada Lovelace): micro-batch 32 with SDPA, PagedAdamW8bit, and Grad Checkpointing (~16 GB VRAM)
-batch_size: 32
-grad_accum: 4
+batch_size: 8
+grad_accum: 16
 target_global_batch: 128
 auto_tune: false
 
-# FULL PRETRAINING: Dynamic Auto-Steps from 125B+ Tokens (2,109 Parquet files on Hub)
-total_tokens: 126500000000  # 126.5B authentic tokens
-max_steps: auto            # Automatically computed: total_tokens // (batch_size * grad_accum * seq_len)
-lr: 3.0e-04
-min_lr: 3.0e-05
-warmup_steps: 15000
+total_tokens: 126500000000
+max_steps: auto
+optimizer: paged_adamw_8bit
+lr: 2.5e-04
+min_lr: 2.5e-05
+warmup_steps: 1000
 weight_decay: 0.1
 betas: [0.9, 0.95]
 grad_clip: 1.0
 
-# Native Ada Lovelace BF16 precision + Muon + Grad Ckpt
 amp: bf16
-optimizer: muon
-muon_lr: 0.02
-muon_momentum: 0.95
-muon_weight_decay: 0.01
 grad_ckpt: true
 
-# 2026 speed stack (opt-in, default OFF for stability-first launch)
 compile: false
 compile_mode: reduce-overhead
-fp8: false  # 4090 (sm_89) pe unsupported — sirf 5090/Blackwell
+fp8: false
 
-eval_every: 500
-save_every: 500
-log_every: 1
+eval_every: 1000
+save_every: 1000
+log_every: 25
 ckpt_dir: ../checkpoints
 log_dir: ../../logs
 
-# Checkpoint auto-push to Hugging Face Hub (background non-blocking thread)
 hf_push: true
 hf_repo: ViuAI/ViuMini-Dense-360M
 hf_push_every: 0
-hf_push_full_every: 500
+hf_push_full_every: 1000
 keep_local_full: 2
 keep_remote_full: 3
 resume: auto
 sample_every: 500
 sample_prompt: "bhai kal party me kya scene hai"
 
-# Streaming & Dataset settings
 hf_dataset: ViuAI/viu-mini-raw-pretrain
-hf_max_blocks: 5000
+hf_max_blocks: 50000
 hf_safety: full
 hf_max_rows: null
 enforce_mix: true
@@ -85,4 +75,3 @@ lang_mix:
 natural_shuffle_buffer: 50000
 hf_round: 0
 hf_round_rows: 200000
-