fix lion optimizer documentation

EleutherAI · Oct 31, 2023 · 8ebf7c6 · 8ebf7c6
1 parent b02d989
commit 8ebf7c6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/megatron/neox_arguments/neox_args.py b/megatron/neox_arguments/neox_args.py
@@ -390,7 +390,7 @@ class NeoXArgsOptimizer(NeoXArgsTemplate):
         "adam", "onebitadam", "cpu_adam", "cpu_torch_adam", "sm3", "madgrad_wd", "sgd", "lion"
     ] = "adam"
     """
-    Type of optimizer to use. Choose from ['adam', 'onebitadam', 'cpu_adam', 'cpu_torch_adam', 'sm3', 'madgrad_wd', 'sgd']
+    Type of optimizer to use. Choose from ['adam', 'onebitadam', 'cpu_adam', 'cpu_torch_adam', 'sm3', 'madgrad_wd', 'sgd', 'lion']
     NOTE: sgd will use MuSGD from Mup. Mup must be enabled for this optimizer.
     """