config step=22500

config.json

@@ -1,16 +1,11 @@
 {
   "checkpoint_dir": "/workspace/adaptive_model/checkpoints",
   "tokenizer_dir": "/workspace/adaptive_model/tokenizer",
-  "run_name": "drlm_v1_diffusion",
+  "run_name": "phase2_deepseek_exclusive",
   "hf_repo_id": "moudook/adaptive-model",
   "hf_push_every_n": 3,
   "dataset_names": [
-    "open-thoughts/OpenThoughts-114k",
-    "mlabonne/FineTome-100k",
-    "agentlans/TeichAI-thinking-reasoning-x",
-    "m-a-p/TerminalTraj",
-    "nvidia/HelpSteer2",
-    "glaiveai/glaive-function-calling-v2"
+    "a-m-team/AM-DeepSeek-R1-Distilled-1.4M"
   ],
   "local_paths": {},
   "vocab_size": 32000,
@@ -21,21 +16,21 @@
   "intermediate_dim": 5504,
   "max_seq_len": 2048,
   "dtype": "bfloat16",
-  "learning_rate": 0.0003,
+  "learning_rate": 0.0001,
   "weight_decay": 0.1,
   "beta1": 0.9,
   "beta2": 0.95,
   "grad_clip": 1.0,
   "warmup_steps": 200,
-  "total_steps": 50000,
-  "batch_size": 32,
-  "grad_accum": 4,
-  "save_every": 250,
+  "total_steps": 100000,
+  "batch_size": 64,
+  "grad_accum": 2,
+  "save_every": 500,
   "log_every": 10,
-  "keep_last_n": 5,
+  "keep_last_n": 3,
   "use_wandb": false,
   "wandb_project": "adaptive-model",
   "device": "cuda",
   "seed": 42,
-  "prefetch_batches": 4
+  "prefetch_batches": 8
 }