grokadamw.GrokAdamW

scripts/pretrain-core-model.yaml

@@ -63,8 +63,8 @@ train:
   log_interval: 1
 
   # 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: 512
+  # global_batch_size: 256
 
   # Number of samples per data-parallel rank (type: int, default: 4)
   # micro_batch_size: 4
@@ -114,12 +114,24 @@ eval:
 
 # Optimizer-related arguments
 
+# 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
+#   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.99
+
 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: grokadamw.GrokAdamW
   init_args:
     # (type: float, default: 0.001)
     lr: 1e-4
@@ -128,7 +140,7 @@ optimizer:
     # (type: tuple, default: (0.9,0.999))
     betas:
       - 0.9
-      - 0.99
+      - 0.999
 
 # How many devices/GPUs to use. Uses all GPUs by default. (type: Union[int, str], default: auto)
 devices: auto

scripts/requirements.in

@@ -1,29 +1,22 @@
 # pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
 torch>=2.5.0,<2.6.0
 numpy<2.0
-torchao
 
 tqdm
+Pillow
 datasets
 jinja2
 transformers
 wandb
-# litgpt[all]
 litgpt[all] @ git+https://github.com/Lightning-AI/litgpt.git
 mergekit @ git+https://github.com/arcee-ai/mergekit.git
-# litgpt @ git+https://github.com/Lightning-AI/litgpt.git
-# litdata
-# litdata @ git+https://github.com/Lightning-AI/litdata.git
-# lpmm @ git+https://github.com/thu-ml/low-bit-optimizers.git
 # muon @ git+https://github.com/KellerJordan/Muon
 # pytorch-optimizer
-lm_eval[ifeval,math]
+torchao
 bitsandbytes
-# grokadamw
+grokadamw
 # sophia-opt
-# bitsandbytes
 # pyzstd
 # zstd
-unsloth
-
-Pillow
+# unsloth
+lm_eval[ifeval,math]