more cli arguments

EleutherAI · Nov 20, 2023 · ede55f4 · ede55f4
1 parent d87e4a3
commit ede55f4
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/megatron/neox_arguments/arguments.py b/megatron/neox_arguments/arguments.py
@@ -258,6 +258,14 @@ def consume_deepy_args(cls):
             "--warmup_iters",
             type=int,
         )
+        group.add_argument(
+            "--lr",
+            type=float,
+        )
+        group.add_argument(
+            "--min_lr",
+            type=float,
+        )
         group.add_argument(
             "--lr_decay_iters",
             type=int,
@@ -1031,10 +1039,12 @@ def calculate_derived(self):
         self.update_values(
             {
                 "optimizer_type": opt_params.get("type", OPT_DEFAULT),
-                "lr": opt_params["params"].get("lr", OPT_PARAMS_DEFAULTS["lr"]),
             }
         )
 
+        if not self.lr:
+            self.update_values({"lr": opt_params["params"].get("lr", OPT_PARAMS_DEFAULTS["lr"])})
+
         if self.optimizer_type.lower() == "onebitadam":
             # onebitadam needs to instantiated by deepspeed, and so we need to pass deepspeed scheduler args
             # for all other optimizers, the scheduling is handled by megatron