pretrain core 4

scripts/pretrain_core_model_4.yaml

@@ -60,7 +60,7 @@ train:
   # Number of samples between optimizer steps across data-parallel ranks (type: int, default: 512)
   # global_batch_size: 512
   # global_batch_size: 256
-  global_batch_size: 64
+  global_batch_size: 32
 
   # Number of samples per data-parallel rank (type: int, default: 4)
   micro_batch_size: 1
@@ -124,31 +124,15 @@ eval:
 #       - 0.9
 #       - 0.999
 
-# optimizer:
-#   class_path: sophia_opt.SophiaG
-#   init_args:
-#     lr: 1e-4
-#     betas:
-#       - 0.9
-#       - 0.95
-#     rho: 0.05
-#     weight_decay: 0.1
-
 optimizer:
-  # class_path: torch.optim.AdamW
-  # class_path: torchao.prototype.low_bit_optim.AdamW8bit
-  # class_path: torchao.prototype.low_bit_optim.AdamW4bit
-  # class_path: bitsandbytes.optim.AdamW8bit
-  class_path: bitsandbytes.optim.PagedAdamW8bit
+  class_path: sophia_opt.SophiaG
   init_args:
-    # (type: float, default: 0.001)
     lr: 1e-4
-    # (type: float, default: 0.01)
-    weight_decay: 0.01
-    # (type: tuple, default: (0.9,0.999))
     betas:
       - 0.9
-      - 0.999
+      - 0.95
+    rho: 0.05
+    weight_decay: 0.1
 
 # How many devices/GPUs to use. Uses all GPUs by default. (type: Union[int, str], default: auto)
 devices: auto