sophia_opt.SophiaG

scripts/pretrain-model.yaml

@@ -115,21 +115,32 @@ eval:
   final_validation: true
 
 # Optimizer-related arguments
+# optimizer:
+#   # class_path: torch.optim.AdamW
+#   class_path: grokadamw.GrokAdamW
+#
+#   init_args:
+#     #   (type: float, default: 0.001)
+#     lr: 4e-04
+#
+#     #   (type: float, default: 0.01)
+#     weight_decay: 0.1
+#
+#     #   (type: tuple, default: (0.9,0.999))
+#     betas:
+#       - 0.9
+#       - 0.95
+
 optimizer:
-  # class_path: torch.optim.AdamW
-  class_path: grokadamw.GrokAdamW
+  class_path: sophia_opt.SophiaG
 
   init_args:
-    #   (type: float, default: 0.001)
-    lr: 4e-04
-
-    #   (type: float, default: 0.01)
-    weight_decay: 0.1
-
-    #   (type: tuple, default: (0.9,0.999))
+    lr: 2e-4
     betas:
-      - 0.9
-      - 0.95
+      - 0.965
+      - 0.99
+    rho: 0.01
+    weight_decay: 1e-1
 
 # How many devices/GPUs to use. Uses all GPUs by default. (type: Union[int, str], default: auto)
 devices: auto

scripts/requirements.in

@@ -17,3 +17,5 @@ grokadamw
 # pyzstd
 # zstd
 Pillow
+
+sophia-opt